def test_create_w_term_image(self):
m31image = create_test_image(cellsize=0.001)
im = create_w_term_like(m31image, w=20000.0, remove_shift=True)
im.data = im.data.real
for x in [64, 64 + 128]:
for y in [64, 64 + 128]:
self.assertAlmostEqual(im.data[0, 0, y, x],
0.84946344276442431, 7)
export_image_to_fits(im, '%s/test_wterm.fits' % self.dir)
assert im.data.shape == (1, 1, 256, 256)
self.assertAlmostEqual(numpy.max(im.data.real), 1.0, 7)
def predict_wstack_single(vis,
model,
remove=True,
facets=1,
vis_slices=1,
**kwargs) -> Visibility:
""" Predict using a single w slices.
This processes a single w plane, rotating out the w beam for the average w
:param vis: Visibility to be predicted
:param model: model image
:return: resulting visibility (in place works)
"""
if not isinstance(vis, Visibility):
log.debug("predict_wstack_single: Coalescing")
avis = coalesce_visibility(vis, **kwargs)
else:
avis = vis
log.debug("predict_wstack_single: predicting using single w slice")
avis.data['vis'] *= 0.0
# We might want to do wprojection so we remove the average w
w_average = numpy.average(avis.w)
avis.data['uvw'][..., 2] -= w_average
tempvis = copy_visibility(avis)
# Calculate w beam and apply to the model. The imaginary part is not needed
workimage = copy_image(model)
w_beam = create_w_term_like(model, w_average, vis.phasecentre)
# Do the real part
workimage.data = w_beam.data.real * model.data
avis = predict_2d(avis, workimage, facets=1, vis_slices=1, **kwargs)
# and now the imaginary part
workimage.data = w_beam.data.imag * model.data
tempvis = predict_2d(tempvis,
workimage,
facets=facets,
vis_slices=vis_slices,
**kwargs)
avis.data['vis'] -= 1j * tempvis.data['vis']
if not remove:
avis.data['uvw'][..., 2] += w_average
if isinstance(vis, BlockVisibility) and isinstance(avis, Visibility):
log.debug("imaging.predict decoalescing post prediction")
return decoalesce_visibility(avis)
else:
return avis
def test_convert_image_to_kernel(self):
m31image = create_test_image(cellsize=0.001,
frequency=[1e8],
canonical=True)
screen = create_w_term_like(m31image, w=20000.0, remove_shift=True)
screen_fft = fft_image(screen)
converted = convert_image_to_kernel(screen_fft, 8, 8)
assert converted.shape == (1, 1, 8, 8, 8, 8)
with self.assertRaises(AssertionError):
converted = convert_image_to_kernel(m31image, 15, 1)
with self.assertRaises(AssertionError):
converted = convert_image_to_kernel(m31image, 15, 1000)
def test_create_w_term_image(self):
newimage = create_image(
npixel=1024,
cellsize=0.001,
polarisation_frame=PolarisationFrame("stokesIQUV"),
frequency=numpy.linspace(0.8e9, 1.2e9, 5),
channel_bandwidth=1e7 * numpy.ones([5]))
im = create_w_term_like(newimage,
w=2000.0,
remove_shift=True,
dopol=True)
im.data = im.data.real
for x in [256, 768]:
for y in [256, 768]:
self.assertAlmostEqual(im.data[0, 0, y, x],
-0.46042631800538464, 7)
export_image_to_fits(im, '%s/test_wterm.fits' % self.dir)
assert im.data.shape == (5, 4, 1024, 1024), im.data.shape
self.assertAlmostEqual(numpy.max(im.data.real), 1.0, 7)
def predict_wstack_single(vis,
model,
remove=True,
gcfcf=None,
**kwargs) -> Visibility:
""" Predict using a single w slices.
This processes a single w plane, rotating out the w beam for the average w
:param vis: Visibility to be predicted
:param model: model image
:return: resulting visibility (in place works)
"""
assert isinstance(vis, Visibility), vis
vis.data['vis'][...] = 0.0
log.debug("predict_wstack_single: predicting using single w slice")
# We might want to do wprojection so we remove the average w
w_average = numpy.average(vis.w)
if remove:
vis.data['uvw'][..., 2] -= w_average
tempvis = copy_visibility(vis)
# Calculate w beam and apply to the model. The imaginary part is not needed
workimage = copy_image(model)
w_beam = create_w_term_like(model, w_average, vis.phasecentre)
# Do the real part
workimage.data = w_beam.data.real * model.data
vis = predict_2d(vis, workimage, gcfcf=gcfcf, **kwargs)
# and now the imaginary part
workimage.data = w_beam.data.imag * model.data
tempvis = predict_2d(tempvis, workimage, gcfcf=gcfcf, **kwargs)
vis.data['vis'] -= 1j * tempvis.data['vis']
if remove:
vis.data['uvw'][..., 2] += w_average
return vis
def invert_wstack_single(vis: Visibility,
im: Image,
dopsf,
normalize=True,
remove=True,
facets=1,
vis_slices=1,
**kwargs) -> (Image, numpy.ndarray):
"""Process single w slice
:param vis: Visibility to be inverted
:param im: image template (not changed)
:param dopsf: Make the psf instead of the dirty image
:param normalize: Normalize by the sum of weights (True)
"""
log.debug("invert_wstack_single: predicting using single w slice")
kwargs['imaginary'] = True
assert isinstance(vis, Visibility), vis
# We might want to do wprojection so we remove the average w
w_average = numpy.average(vis.w)
vis.data['uvw'][..., 2] -= w_average
reWorkimage, sumwt, imWorkimage = invert_2d(vis,
im,
dopsf,
normalize=normalize,
facets=facets,
vis_slices=vis_slices,
**kwargs)
if not remove:
vis.data['uvw'][..., 2] += w_average
# Calculate w beam and apply to the model. The imaginary part is not needed
w_beam = create_w_term_like(im, w_average, vis.phasecentre)
reWorkimage.data = w_beam.data.real * reWorkimage.data - w_beam.data.imag * imWorkimage.data
return reWorkimage, sumwt
def create_awterm_convolutionfunction(im,
make_pb=None,
nw=1,
wstep=1e15,
oversampling=8,
support=6,
use_aaf=True,
maxsupport=512):
""" Fill AW projection kernel into a GridData.
:param im: Image template
:param make_pb: Function to make the primary beam model image
:param nw: Number of w planes
:param wstep: Step in w (wavelengths)
:param oversampling: Oversampling of the convolution function in uv space
:return: griddata correction Image, griddata kernel as GridData
"""
d2r = numpy.pi / 180.0
# We only need the griddata correction function for the PSWF so we make
# it for the shape of the image
nchan, npol, ony, onx = im.data.shape
assert isinstance(im, Image)
# Calculate the template convolution kernel.
cf = create_convolutionfunction_from_image(im,
oversampling=oversampling,
support=support)
cf_shape = list(cf.data.shape)
cf_shape[2] = nw
cf.data = numpy.zeros(cf_shape).astype('complex')
cf.grid_wcs.wcs.crpix[4] = nw // 2 + 1.0
cf.grid_wcs.wcs.cdelt[4] = wstep
cf.grid_wcs.wcs.ctype[4] = 'WW'
if numpy.abs(wstep) > 0.0:
w_list = cf.grid_wcs.sub([5]).wcs_pix2world(range(nw), 0)[0]
else:
w_list = [0.0]
assert isinstance(oversampling, int)
assert oversampling > 0
nx = max(maxsupport, 2 * oversampling * support)
ny = max(maxsupport, 2 * oversampling * support)
qnx = nx // oversampling
qny = ny // oversampling
cf.data[...] = 0.0
subim = copy_image(im)
ccell = onx * numpy.abs(d2r * subim.wcs.wcs.cdelt[0]) / qnx
subim.data = numpy.zeros([nchan, npol, qny, qnx])
subim.wcs.wcs.cdelt[0] = -ccell / d2r
subim.wcs.wcs.cdelt[1] = +ccell / d2r
subim.wcs.wcs.crpix[0] = qnx // 2 + 1.0
subim.wcs.wcs.crpix[1] = qny // 2 + 1.0
if use_aaf:
this_pswf_gcf, _ = create_pswf_convolutionfunction(subim,
oversampling=1,
support=6)
norm = 1.0 / this_pswf_gcf.data
else:
norm = 1.0
if make_pb is not None:
pb = make_pb(subim)
rpb, footprint = reproject_image(pb, subim.wcs, shape=subim.shape)
rpb.data[footprint.data < 1e-6] = 0.0
norm *= rpb.data
# We might need to work with a larger image
padded_shape = [nchan, npol, ny, nx]
thisplane = copy_image(subim)
thisplane.data = numpy.zeros(thisplane.shape, dtype='complex')
for z, w in enumerate(w_list):
thisplane.data[...] = 0.0 + 0.0j
thisplane = create_w_term_like(thisplane, w, dopol=True)
thisplane.data *= norm
paddedplane = pad_image(thisplane, padded_shape)
paddedplane = fft_image(paddedplane)
ycen, xcen = ny // 2, nx // 2
for y in range(oversampling):
ybeg = y + ycen + (support * oversampling) // 2 - oversampling // 2
yend = y + ycen - (support * oversampling) // 2 - oversampling // 2
vv = range(ybeg, yend, -oversampling)
for x in range(oversampling):
xbeg = x + xcen + (support *
oversampling) // 2 - oversampling // 2
xend = x + xcen - (support *
oversampling) // 2 - oversampling // 2
uu = range(xbeg, xend, -oversampling)
for chan in range(nchan):
for pol in range(npol):
cf.data[chan, pol, z, y, x, :, :] = paddedplane.data[
chan, pol, :, :][vv, :][:, uu]
cf.data /= numpy.sum(
numpy.real(cf.data[0, 0, nw // 2, oversampling // 2,
oversampling // 2, :, :]))
cf.data = numpy.conjugate(cf.data)
if use_aaf:
pswf_gcf, _ = create_pswf_convolutionfunction(im,
oversampling=1,
support=6)
else:
pswf_gcf = create_empty_image_like(im)
pswf_gcf.data[...] = 1.0
return pswf_gcf, cf
def create_awterm_convolutionfunction(im,
make_pb=None,
nw=1,
wstep=1e15,
oversampling=8,
support=6,
use_aaf=True,
maxsupport=512,
**kwargs):
""" Fill AW projection kernel into a GridData.
:param im: Image template
:param make_pb: Function to make the primary beam model image (hint: use a partial)
:param nw: Number of w planes
:param wstep: Step in w (wavelengths)
:param oversampling: Oversampling of the convolution function in uv space
:return: griddata correction Image, griddata kernel as GridData
"""
d2r = numpy.pi / 180.0
# We only need the griddata correction function for the PSWF so we make
# it for the shape of the image
nchan, npol, ony, onx = im.data.shape
assert isinstance(im, Image)
# Calculate the template convolution kernel.
cf = create_convolutionfunction_from_image(im,
oversampling=oversampling,
support=support)
cf_shape = list(cf.data.shape)
cf_shape[2] = nw
cf.data = numpy.zeros(cf_shape).astype('complex')
cf.grid_wcs.wcs.crpix[4] = nw // 2 + 1.0
cf.grid_wcs.wcs.cdelt[4] = wstep
cf.grid_wcs.wcs.ctype[4] = 'WW'
if numpy.abs(wstep) > 0.0:
w_list = cf.grid_wcs.sub([5]).wcs_pix2world(range(nw), 0)[0]
else:
w_list = [0.0]
assert isinstance(oversampling, int)
assert oversampling > 0
nx = max(maxsupport, 2 * oversampling * support)
ny = max(maxsupport, 2 * oversampling * support)
qnx = nx // oversampling
qny = ny // oversampling
cf.data[...] = 0.0
subim = copy_image(im)
ccell = onx * numpy.abs(d2r * subim.wcs.wcs.cdelt[0]) / qnx
subim.data = numpy.zeros([nchan, npol, qny, qnx])
subim.wcs.wcs.cdelt[0] = -ccell / d2r
subim.wcs.wcs.cdelt[1] = +ccell / d2r
subim.wcs.wcs.crpix[0] = qnx // 2 + 1.0
subim.wcs.wcs.crpix[1] = qny // 2 + 1.0
if use_aaf:
this_pswf_gcf, _ = create_pswf_convolutionfunction(subim,
oversampling=1,
support=6)
norm = 1.0 / this_pswf_gcf.data
else:
norm = 1.0
if make_pb is not None:
pb = make_pb(subim)
rpb, footprint = reproject_image(pb, subim.wcs, shape=subim.shape)
rpb.data[footprint.data < 1e-6] = 0.0
norm *= rpb.data
# We might need to work with a larger image
padded_shape = [nchan, npol, ny, nx]
thisplane = copy_image(subim)
thisplane.data = numpy.zeros(thisplane.shape, dtype='complex')
for z, w in enumerate(w_list):
thisplane.data[...] = 0.0 + 0.0j
thisplane = create_w_term_like(thisplane, w, dopol=True)
thisplane.data *= norm
paddedplane = pad_image(thisplane, padded_shape)
paddedplane = fft_image(paddedplane)
ycen, xcen = ny // 2, nx // 2
for y in range(oversampling):
ybeg = y + ycen + (support * oversampling) // 2 - oversampling // 2
yend = y + ycen - (support * oversampling) // 2 - oversampling // 2
# vv = range(ybeg, yend, -oversampling)
for x in range(oversampling):
xbeg = x + xcen + (support *
oversampling) // 2 - oversampling // 2
xend = x + xcen - (support *
oversampling) // 2 - oversampling // 2
# uu = range(xbeg, xend, -oversampling)
cf.data[..., z, y,
x, :, :] = paddedplane.data[...,
ybeg:yend:-oversampling,
xbeg:xend:-oversampling]
# for chan in range(nchan):
# for pol in range(npol):
# cf.data[chan, pol, z, y, x, :, :] = paddedplane.data[chan, pol, :, :][vv, :][:, uu]
cf.data /= numpy.sum(
numpy.real(cf.data[0, 0, nw // 2, oversampling // 2,
oversampling // 2, :, :]))
cf.data = numpy.conjugate(cf.data)
#====================================
#Use ASKAPSoft routine to crop the support size
crop_ASKAPSOft_like = True
if crop_ASKAPSOft_like:
#Hardcode the cellsize: 1 / FOV
#uv_cellsize = 57.3;#N=1200 pixel and pixelsize is 3 arcseconds
#uv_cellsize = 43.97;#N=1600 pixel and pixelsize is 3 arcseconds
#uv_cellsize = 114.6;#N=1800 pixel with 1 arcsecond pixelsize
#uv_cellsize = 57.3;#N=1800 pixel with 2 arcsecond pixelsize
#uv_cellsize = 1145.91509915;#N=1800 pixxel with 0.1 arcsecond pixelsize
#Get from **kwargs
if kwargs is None:
#Safe solution works for baselines up to > 100km and result in small kernels
uv_cellsize = 1145.91509915
#N=1800 pixxel with 0.1 arcsecond pixelsize
if 'UVcellsize' in kwargs.keys():
uv_cellsize = kwargs['UVcellsize']
#print(uv_cellsize);
#Cutoff param in ASKAPSoft hardcoded as well
ASKAPSoft_cutof = 0.1
wTheta_list = numpy.zeros(len(w_list))
for i in range(0, len(w_list)):
if w_list[i] == 0:
wTheta_list[i] = 0.9
#This is due to the future if statements cause if it is small, the kernel will be 3 which is a clear cutoff
else:
wTheta_list[i] = numpy.fabs(
w_list[i]) / (uv_cellsize * uv_cellsize)
kernel_size_list = []
#We rounded the kernels according to conventional rounding rules
for i in range(0, len(wTheta_list)):
#if wTheta_list[i] < 1:
if wTheta_list[i] < 1: #Change to ASKAPSoft
kernel_size_list.append(int(3.))
elif ASKAPSoft_cutof < 0.01:
kernel_size_list.append(int(6 + 1.14 * wTheta_list[i]))
else:
kernel_size_list.append(
int(numpy.sqrt(49 + wTheta_list[i] * wTheta_list[i])))
log.info('W-kernel w-terms:')
log.info(w_list)
log.info('Corresponding w-kernel sizes:')
log.info(kernel_size_list)
print(numpy.unique(kernel_size_list))
#print(kernel_size_list);
crop_list = []
#another rounding according to conventional rounding rules
for i in range(0, len(kernel_size_list)):
if support - kernel_size_list[i] <= 0:
crop_list.append(int(0))
else:
crop_list.append(int((support - kernel_size_list[i]) / 2))
#Crop original suppor
for i in range(0, nw):
if crop_list[i] != 0:
cf.data[0, 0, i, :, :, 0:crop_list[i], :] = 0
cf.data[0, 0, i, :, :, -crop_list[i]:, :] = 0
cf.data[0, 0, i, :, :, :, 0:crop_list[i]] = 0
cf.data[0, 0, i, :, :, :, -crop_list[i]:] = 0
else:
pass
#Plot
#import matplotlib.pyplot as plt
#cf.data[0,0,i,0,0,...][cf.data[0,0,i,0,0,...] != 0.] = 1+0.j;
#plt.imshow(numpy.real(cf.data[0,0,i,0,0,...]))
#plt.show(block=True)
#plt.close();
#====================================
if use_aaf:
pswf_gcf, _ = create_pswf_convolutionfunction(im,
oversampling=1,
support=6)
else:
pswf_gcf = create_empty_image_like(im)
pswf_gcf.data[...] = 1.0
return pswf_gcf, cf
dirtyFacet2 = create_image_from_visibility(vt, npixel=npixel, npol=1, cellsize=cellsize)
future = invert_list_arlexecute_workflow([vt], [dirtyFacet2], facets=2, context='facets')
dirtyFacet2, sumwt = arlexecute.compute(future, sync=True)[0]
if doplot:
show_image(dirtyFacet2,cm='hot')
plt.title("Dirty Facet image")
plt.savefig('%s/%s_dirtyfacet2.pdf' % (results_dir, str(freq)))
print("Max, min in dirty image = %.6f, %.6f, sumwt = %f" % (
dirtyFacet2.data.max(), dirtyFacet2.data.min(), sumwt))
export_image_to_fits(dirtyFacet2, '%s/imaging-wterm_dirtyFacet2.fits' % (results_dir))
if doplot:
wterm = create_w_term_like(model, phasecentre=vt.phasecentre, w=numpy.max(vt.w))
show_image(wterm)
plt.title("Wterm image")
plt.savefig('%s/%s_wterm.pdf' % (results_dir, str(freq)))
dirtywstack = create_image_from_visibility(vt, npixel=npixel, npol=1, cellsize=cellsize)
future = invert_list_arlexecute_workflow([vt], [dirtywstack], vis_slices=101, context='wstack')
dirtywstack, sumwt = arlexecute.compute(future, sync=True)[0]
show_image(dirtywstack)
plt.title("dirtywstack image")
plt.savefig('%s/%s_dirtywstack.pdf' % (results_dir, str(freq)))
# plt.show()
print("Max, min in dirty image = %.6f, %.6f, sumwt = %f" %
(dirtywstack.data.max(), dirtywstack.data.min(), sumwt))
