# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Illustration of A-projection
#
# How to deal with direction- and baseline-dependent delays when imaging using $A$-kernels.

# +
# %matplotlib inline

import sys
sys.path.append('../..')

from matplotlib import pylab

pylab.rcParams['figure.figsize'] = 10, 10

import functools
import numpy
import scipy
import scipy.special
import astropy
import astropy.units as u

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from ipywidgets import interact

from crocodile.synthesis import *
from crocodile.simulate import *
from util.visualize import *
from arl.test_support import create_named_configuration
# -

# Generate baseline coordinates for a short observation with the VLA where the target is near the zenith. This means minimal w-values - the easy case.

# +
theta = 0.04
lam = 18000
grid_size = int(theta * lam)

vlas = create_named_configuration('VLAA_north')
ha_range = numpy.arange(numpy.radians(-30),
                        numpy.radians(30),
                        numpy.radians(90 / 360))
dec = vlas.location.lat
vobs = xyz_to_baselines(vlas.data['xyz'], ha_range, dec)

# Create antenna mapping for visibilities
antennas = vlas.data['xyz']
nant = len(antennas)
ant1,ant2 = baseline_ids(nant, len(ha_range))
ant1xy = vlas.data['xyz'][ant1,:2]
ant2xy = vlas.data['xyz'][ant2,:2]

# Wavelength: 5 metres 
wvl=5
uvw = vobs / wvl

ax = plt.figure().add_subplot(111, projection='3d')
ax.scatter(uvw[:,0], uvw[:,1] , uvw[:,2])
max_uvw = numpy.amax(uvw)
ax.set_xlabel('U [$\lambda$]'); ax.set_xlim((-max_uvw, max_uvw))
ax.set_ylabel('V [$\lambda$]'); ax.set_ylim((-max_uvw, max_uvw))
ax.set_zlabel('W [$\lambda$]'); ax.set_zlim((-max_uvw, max_uvw))
ax.view_init(20, 20)
pylab.show()
# -

# ## Ionosphere
#
# However, let us assume that the ionosphere introduces random delays (refraction) into our data based on the location of the antenna and direction. Normally this delay screen would depend on both time and frequency too, but let us ignore that for the moment to keep things simple:

# +
ion_res = 2000 # m
ion_height = 300000 # m
ion_fov = int(theta * ion_height)
print("Ionospheric field of view:", ion_fov//1000, "km")
ion_size = 74000 + ion_fov # m
print("Delay screen size:", ion_size//1000, "km")
ion_max_delay = 2e-8 # s
numpy.random.seed(0)

ion_delay = ion_max_delay * numpy.random.random((ion_size // ion_res, ion_size // ion_res))

# Visualise, including antenna (ground) positions (for ha=0) to give a sense of scale
ax = plt.subplot()
img = ax.imshow(ion_delay,interpolation='bilinear',
                extent=(-ion_size/2,ion_size/2,-ion_size/2,ion_size/2));
ax.scatter(vlas.data['xyz'][:,0], vlas.data['xyz'][:,1], c='red')
ax.set_title("Ionospheric delay"); plt.colorbar(img)
ax.set_xlabel('X [m]'); ax.set_ylabel('Y [m]');


# -

# For example, here is the phase screen applying to our field of view at the centre of the telescope:

def ion_sample(ant, l, m):
    # Sample image at appropriate position over the antenna
    d = sample_image(ion_delay, (ant[0] + l * ion_height) / ion_res,
                                (ant[1] + m * ion_height) / ion_res)
    # Convert to phase difference for our wavelength
    return(numpy.exp(2j * numpy.pi * d * astropy.constants.c.value / wvl))
ls, ms = theta * coordinates2(5*ion_fov // ion_res)
pylab.rcParams['figure.figsize'] = 16, 10
show_image(ion_sample((0,0), ls, ms), "phase screen", theta);


# Now let's simulate visibilities. The delay will depend on both antennas involved in the baseline *and* the target position. This introduces some considerable "noise" into the phases:

# +
def add_point(l, m):
    phasor = ion_sample(numpy.transpose(antennas[ant1,:2]), l, m) / \
             ion_sample(numpy.transpose(antennas[ant2,:2]), l, m)
    return phasor, phasor * simulate_point(uvw, l,m)

# Grid of points in the middle
vis = numpy.zeros(len(uvw), dtype=complex)
import itertools
for il, im in itertools.product(range(-3, 4), range(-3, 4)):
    vis += add_point(theta/10*il, theta/10*im)[1]
# Extra dot to mark upper-right corner
vis += add_point(theta*0.28, theta*0.28)[1]
# Extra dot to mark upper-left corner
vis += add_point(theta*-0.32, theta*0.28)[1]
# -

# Because we chose low $w$-values, we can use simple imaging here. However, the noise we added to the phases messes quite a bit with our ability to locate sources:

d,p,_=do_imaging(theta, lam, uvw, None, vis, simple_imaging)
show_image(d, "image", theta)


def zoom(l=0, m=0): show_image(d, "image", theta, xlim=(l-theta/10,l+theta/10), ylim=(m-theta/10,m+theta/10))
interact(zoom, l=(-theta/2,theta/2,theta/10), m=(-theta/2,theta/2,theta/10));

# The tricky aspect of this noise is that it is direction-dependent. This means that they have to be removed within imaging where we introduce direction again. As we will be working in the grid, we therefore make $A$-kernels that compensate for the introduced phase error.
#
# Note that normally $A$-kernels are unknowns, so we would at best have approximations of those available:

ion_oversample = 10
ls, ms = theta * coordinates2(ion_oversample * ion_fov // ion_res)
print("A pattern size: %dx%d" % ls.shape)
apattern = []
for ant in range(nant):
    apattern.append(ion_sample(vlas.data['xyz'][ant], ls, ms))
show_image(apattern[0], "apattern", theta)
show_grid(fft(apattern[0]), "akern", theta)

# Our actual kernels will however we for antenna combinations (baselines). Therefore we make combinations. These are our actual kernels, so now we can do the FFT. We reduce the support a bit as well to make imaging faster.

Nkern = min(25, ls.shape[0])
akern_combs = numpy.empty((nant, nant, 1, 1, Nkern, Nkern), dtype=complex)
for a1 in range(nant):
    for a2 in range(a1+1,nant):
        akern_combs[a1, a2, 0, 0] = extract_mid(fft(apattern[a2] / apattern[a1]), Nkern)

# Some examples for the kernels we just generated. Short baselines will see almost exactly the same ionosphere, so the kernel is going to be relatively trivial - dominated by a single dot at $(0,0)$:

show_grid(akern_combs[0,1,0,0], "aakern", theta)

# On the other hand, for long baselines there is much more turbulence to compensate for, so the kernels start looking increasingly chaotic:

longest = numpy.argmax(uvw[:,0]**2+uvw[:,1]**2)
show_grid(akern_combs[ant1[longest], ant2[longest],0,0], "aakern", theta)

# As random as these kernels look, they are exactly what we need to restore imaging performance:

d_w,p_w,_=do_imaging(theta, lam, uvw, numpy.transpose([ant1, ant2]), vis, conv_imaging, kv=akern_combs)
show_image(d_w, "image", theta)


def zoom(l=0, m=0): show_image(d_w, "image", theta, xlim=(l-theta/10,l+theta/10), ylim=(m-theta/10,m+theta/10))
interact(zoom, l=(-theta/2,theta/2,theta/10), m=(-theta/2,theta/2,theta/10));

# Not required, but we can also easily add an anti-aliasing function into the mix and oversample the kernel. Gridding complexity doesn't change, and we get more accuracy:

Qpx = 8; c = 5
aa = anti_aliasing_function(ls.shape, 0, c)
akern_combs2 = numpy.empty((nant, nant, Qpx, Qpx, Nkern, Nkern), dtype=complex)
for a1 in range(nant):
    for a2 in range(a1+1,nant):
        akern_combs2[a1, a2] = kernel_oversample(aa * apattern[a2] / apattern[a1], Qpx, Nkern)
show_grid(akern_combs2[0,1,0,0], "aakern", theta)
show_grid(akern_combs2[ant1[longest], ant2[longest],0,0], "aakern", theta)
d_w2,p_w2,_=do_imaging(theta, lam, uvw, numpy.transpose([ant1, ant2]), vis, conv_imaging, kv=akern_combs2)
d_w2 /= anti_aliasing_function(d_w2.shape, 0, c)
show_image(d_w2, "image", theta)


def zoom(l=0, m=0): show_image(d_w2, "image", theta, xlim=(l-theta/10,l+theta/10), ylim=(m-theta/10,m+theta/10))
interact(zoom, l=(-theta/2,theta/2,theta/10), m=(-theta/2,theta/2,theta/10));

StarcoderdataJupyterScriptsDedupFiltered