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))