def accumulate_results(results, **kwargs):
newim = copy_image(im)
i = 0
for dpatch in iterator(newim, **kwargs):
dpatch.data[...] = results[i].data[...]
i += 1
return newim
def create_generic_image_iterator_graph(imagefunction, im: Image, iterator,
**kwargs) -> delayed:
""" Definition of interface for create_generic_image_graph
This generates a graph for imagefunction. Note that im cannot be a graph itself.
:func imagefunction: Function to be applied to all pixels
:param im: Image to be processed
:param iterator: iterator e.g. raster_iter
:param kwargs: Parameters for functions in graphs
:return: graph
"""
def accumulate_results(results, **kwargs):
newim = copy_image(im)
i = 0
for dpatch in iterator(newim, **kwargs):
dpatch.data[...] = results[i].data[...]
i += 1
return newim
results = list()
for dpatch in iterator(im, **kwargs):
results.append(delayed(imagefunction(copy_image(dpatch), **kwargs)))
return delayed(accumulate_results, pure=True)(results, **kwargs)
def sum_invert_results(image_list):
""" Sum a set of invert results with appropriate weighting
:param image_list: List of [image, sum weights] pairs
:return: image, sum of weights
"""
first = True
sumwt = 0.0
im = None
for i, arg in enumerate(image_list):
if arg is not None:
if isinstance(arg[1], numpy.ndarray):
scale = arg[1][..., numpy.newaxis, numpy.newaxis]
else:
scale = arg[1]
if first:
im = copy_image(arg[0])
im.data *= scale
sumwt = arg[1]
first = False
else:
im.data += scale * arg[0].data
sumwt += arg[1]
assert not first, "No invert results"
im = normalize_sumwt(im, sumwt)
return im, sumwt
def ingest_visibility(self, freq=None, chan_width=None, times=None, reffrequency=None, add_errors=False,
block=True):
if freq is None:
freq = [1e8]
if times is None:
ntimes = 5
times = numpy.linspace(-numpy.pi / 3.0, numpy.pi / 3.0, ntimes)
if chan_width is None:
chan_width = [1e6]
if reffrequency is None:
reffrequency = [1e8]
lowcore = create_named_configuration('LOWBD2-CORE')
frequency = numpy.array([freq])
channel_bandwidth = numpy.array([chan_width])
phasecentre = SkyCoord(ra=+180.0 * u.deg, dec=-60.0 * u.deg, frame='icrs', equinox='J2000')
if block:
vt = create_blockvisibility(lowcore, times, frequency, channel_bandwidth=channel_bandwidth,
weight=1.0, phasecentre=phasecentre,
polarisation_frame=PolarisationFrame("stokesI"))
else:
vt = create_visibility(lowcore, times, frequency, channel_bandwidth=channel_bandwidth,
weight=1.0, phasecentre=phasecentre,
polarisation_frame=PolarisationFrame("stokesI"))
cellsize = 0.001
model = create_image_from_visibility(vt, npixel=self.npixel, cellsize=cellsize, npol=1,
frequency=reffrequency, phasecentre=phasecentre,
polarisation_frame=PolarisationFrame("stokesI"))
flux = numpy.array([[100.0]])
facets = 4
rpix = model.wcs.wcs.crpix - 1.0
spacing_pixels = self.npixel // facets
centers = [-1.5, -0.5, 0.5, 1.5]
comps = list()
for iy in centers:
for ix in centers:
p = int(round(rpix[0] + ix * spacing_pixels * numpy.sign(model.wcs.wcs.cdelt[0]))), \
int(round(rpix[1] + iy * spacing_pixels * numpy.sign(model.wcs.wcs.cdelt[1])))
sc = pixel_to_skycoord(p[0], p[1], model.wcs, origin=1)
comp = create_skycomponent(flux=flux, frequency=frequency, direction=sc,
polarisation_frame=PolarisationFrame("stokesI"))
comps.append(comp)
if block:
predict_skycomponent_blockvisibility(vt, comps)
else:
predict_skycomponent_visibility(vt, comps)
insert_skycomponent(model, comps)
self.model = copy_image(model)
self.empty_model = create_empty_image_like(model)
export_image_to_fits(model, '%s/test_pipeline_bags_model.fits' % self.dir)
if add_errors:
# These will be the same for all calls
numpy.random.seed(180555)
gt = create_gaintable_from_blockvisibility(vt)
gt = simulate_gaintable(gt, phase_error=1.0, amplitude_error=0.0)
vt = apply_gaintable(vt, gt)
return vt
def ingest_visibility(self,
freq=1e8,
chan_width=1e6,
times=None,
reffrequency=None,
add_errors=False):
if times is None:
times = (numpy.pi / 12.0) * numpy.linspace(-3.0, 3.0, 5)
if reffrequency is None:
reffrequency = [1e8]
lowcore = create_named_configuration('LOWBD2-CORE')
frequency = numpy.array([freq])
channel_bandwidth = numpy.array([chan_width])
phasecentre = SkyCoord(ra=+180.0 * u.deg,
dec=-60.0 * u.deg,
frame='icrs',
equinox='J2000')
vt = create_visibility(lowcore,
times,
frequency,
channel_bandwidth=channel_bandwidth,
weight=1.0,
phasecentre=phasecentre,
polarisation_frame=PolarisationFrame("stokesI"))
cellsize = 0.001
model = create_image_from_visibility(
vt,
npixel=self.npixel,
cellsize=cellsize,
npol=1,
frequency=reffrequency,
phasecentre=phasecentre,
polarisation_frame=PolarisationFrame("stokesI"))
flux = numpy.array([[100.0]])
facets = 4
rpix = model.wcs.wcs.crpix
spacing_pixels = self.npixel // facets
centers = [-1.5, -0.5, 0.5, 1.5]
comps = list()
for iy in centers:
for ix in centers:
p = int(round(rpix[0] + ix * spacing_pixels * numpy.sign(model.wcs.wcs.cdelt[0]))), \
int(round(rpix[1] + iy * spacing_pixels * numpy.sign(model.wcs.wcs.cdelt[1])))
sc = pixel_to_skycoord(p[0], p[1], model.wcs, origin=0)
comp = create_skycomponent(
flux=flux,
frequency=frequency,
direction=sc,
polarisation_frame=PolarisationFrame("stokesI"))
comps.append(comp)
predict_skycomponent_visibility(vt, comps)
insert_skycomponent(model, comps)
self.model = copy_image(model)
export_image_to_fits(model,
'%s/test_bags_model.fits' % (self.results_dir))
return vt
def copy_skymodel(sm):
""" Copy a sky model
"""
newsm = SkyModel()
if sm.components is not None:
newsm.components = [copy_skycomponent(comp) for comp in sm.components]
if sm.images is not None:
newsm.images = [copy_image(im) for im in sm.images]
return newsm
def sum_invert_results(image_list):
for i, arg in enumerate(image_list):
if i == 0:
im = copy_image(arg[0])
im.data *= arg[1]
sumwt = arg[1]
else:
im.data += arg[1] * arg[0].data
sumwt += arg[1]
im = normalize_sumwt(im, sumwt)
return im, sumwt
def test_calculate_image_frequency_moments(self):
frequency = numpy.linspace(0.9e8, 1.1e8, 9)
cube = create_low_test_image_from_gleam(npixel=512, cellsize=0.0001, frequency=frequency)
log.debug(export_image_to_fits(cube, fitsfile='%s/test_moments_cube.fits' % (self.dir)))
original_cube = copy_image(cube)
moment_cube = calculate_image_frequency_moments(cube, nmoments=3)
log.debug(export_image_to_fits(moment_cube, fitsfile='%s/test_moments_moment_cube.fits' % (self.dir)))
reconstructed_cube = calculate_image_from_frequency_moments(cube, moment_cube)
log.debug(export_image_to_fits(reconstructed_cube, fitsfile='%s/test_moments_reconstructed_cube.fits' % (
self.dir)))
error = numpy.std(reconstructed_cube.data - original_cube.data)
assert error < 0.2
def __init__(self, images=None, components=None):
""" Holds a model of the sky
"""
if images is not None:
self.images = [copy_image(im) for im in images]
else:
self.images = None
if components is not None:
self.components = [copy_skycomponent(sc) for sc in components]
else:
self.components = None
def predict_timeslice_single(vis: Visibility,
model: Image,
predict=predict_2d_base,
remove=True,
**kwargs) -> Visibility:
""" Predict using a single time slices.
This fits a single plane and corrects the image geometry.
:param vis: Visibility to be predicted
:param model: model image
:param predict:
:param remove: Remove fitted w (so that wprojection will do the right thing)
:return: resulting visibility (in place works)
"""
log.debug("predict_timeslice: predicting using time slices")
inchan, inpol, ny, nx = model.shape
vis.data['vis'] *= 0.0
if not isinstance(vis, Visibility):
avis = coalesce_visibility(vis, **kwargs)
else:
avis = vis
# Fit and remove best fitting plane for this slice
avis, p, q = fit_uvwplane(avis, remove=remove)
# Calculate nominal and distorted coordinate systems. We will convert the model
# from nominal to distorted before predicting.
workimage = copy_image(model)
# Use griddata to do the conversion. This could be improved. Only cubic is possible in griddata.
# The interpolation is ok for invert since the image is smooth but for clean images the
# interpolation is particularly poor, leading to speckle in the residual image.
lnominal, mnominal, ldistorted, mdistorted = lm_distortion(model, -p, -q)
for chan in range(inchan):
for pol in range(inpol):
workimage.data[chan, pol, ...] = \
griddata((mnominal.flatten(), lnominal.flatten()),
values=workimage.data[chan, pol, ...].flatten(),
xi=(mdistorted.flatten(), ldistorted.flatten()),
method='cubic',
fill_value=0.0,
rescale=True).reshape(workimage.data[chan, pol, ...].shape)
avis = predict(avis, workimage, **kwargs)
return avis
def sum_invert_results(image_list):
first = True
for i, arg in enumerate(image_list):
if arg is not None:
if first:
im = copy_image(arg[0])
im.data *= arg[1]
sumwt = arg[1]
first = False
else:
im.data += arg[1] * arg[0].data
sumwt += arg[1]
assert numpy.sum(sumwt) > 0.0
im = normalize_sumwt(im, sumwt)
return im, sumwt
def restore_cube(model: Image, psf: Image, residual=None, **kwargs) -> Image:
""" Restore the model image to the residuals
:params psf: Input PSF
:return: restored image
"""
assert isinstance(model, Image), "Type is %s" % (type(model))
assert isinstance(psf, Image), "Type is %s" % (type(psf))
assert residual is None or isinstance(
residual, Image), "Type is %s" % (type(residual))
restored = copy_image(model)
npixel = psf.data.shape[3]
sl = slice(npixel // 2 - 7, npixel // 2 + 8)
size = get_parameter(kwargs, "psfwidth", None)
if size is None:
# isotropic at the moment!
try:
fit = fit_2dgaussian(psf.data[0, 0, sl, sl])
if fit.x_stddev <= 0.0 or fit.y_stddev <= 0.0:
log.debug(
'restore_cube: error in fitting to psf, using 1 pixel stddev'
)
size = 1.0
else:
size = max(fit.x_stddev, fit.y_stddev)
log.debug('restore_cube: psfwidth = %s' % (size))
except:
log.debug(
'restore_cube: warning in fit to psf, using 1 pixel stddev')
size = 1.0
else:
log.debug('restore_cube: Using specified psfwidth = %s' % (size))
# By convention, we normalise the peak not the integral so this is the volume of the Gaussian
norm = 2.0 * numpy.pi * size**2
gk = Gaussian2DKernel(size)
for chan in range(model.shape[0]):
for pol in range(model.shape[1]):
restored.data[chan, pol, :, :] = norm * convolve(
model.data[chan, pol, :, :], gk, normalize_kernel=False)
if residual is not None:
restored.data += residual.data
return restored
def predict_wstack_single(vis, model, remove=True, **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_base(avis, workimage, **kwargs)
# and now the imaginary part
workimage.data = w_beam.data.imag * model.data
tempvis = predict_2d_base(tempvis, workimage, **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 ingest_visibility(self, freq=1e8, chan_width=1e6, times=None, reffrequency=None, add_errors=False):
if times is None:
times = [0.0]
if reffrequency is None:
reffrequency = [1e8]
lowcore = create_named_configuration('LOWBD2-CORE')
frequency = numpy.array([freq])
channel_bandwidth = numpy.array([chan_width])
# phasecentre = SkyCoord(ra=+180.0 * u.deg, dec=-60.0 * u.deg, frame='icrs', equinox='J2000')
# Observe at zenith to ensure that timeslicing works well. We test that elsewhere.
phasecentre = SkyCoord(ra=+180.0 * u.deg, dec=-60.0 * u.deg, frame='icrs', equinox='J2000')
vt = create_blockvisibility(lowcore, times, frequency, channel_bandwidth=channel_bandwidth,
weight=1.0, phasecentre=phasecentre,
polarisation_frame=PolarisationFrame("stokesI"))
cellsize = 0.001
model = create_image_from_visibility(vt, npixel=self.npixel, cellsize=cellsize, npol=1,
frequency=reffrequency,
polarisation_frame=PolarisationFrame("stokesI"))
flux = numpy.array([[100.0]])
facets = 4
rpix = model.wcs.wcs.crpix - 1.0
spacing_pixels = self.npixel // facets
centers = [-1.5, -0.5, 0.5, 1.5]
comps = list()
for iy in centers:
for ix in centers:
p = int(round(rpix[0] + ix * spacing_pixels * numpy.sign(model.wcs.wcs.cdelt[0]))), \
int(round(rpix[1] + iy * spacing_pixels * numpy.sign(model.wcs.wcs.cdelt[1])))
sc = pixel_to_skycoord(p[0], p[1], model.wcs, origin=1)
comps.append(create_skycomponent(flux=flux, frequency=vt.frequency, direction=sc,
polarisation_frame=PolarisationFrame("stokesI")))
predict_skycomponent_blockvisibility(vt, comps)
insert_skycomponent(model, comps)
self.actualmodel = copy_image(model)
export_image_to_fits(model, '%s/test_imaging_model.fits' % (self.results_dir))
if add_errors:
# These will be the same for all calls
numpy.random.seed(180555)
gt = create_gaintable_from_blockvisibility(vt)
gt = simulate_gaintable(gt, phase_error=1.0, amplitude_error=0.0)
vt = apply_gaintable(vt, gt)
return vt
def sum_inver_image_reduce_kernel(im_sumwt1, im_sumwt2):
first = True
total_sumwt = 0.0
ret_im = None
facet_id = im_sumwt1[0].facet_id
for im, sumwt in [im_sumwt1, im_sumwt2]:
if isinstance(sumwt, numpy.ndarray):
scale = sumwt[..., numpy.newaxis, numpy.newaxis]
else:
scale = sumwt
if first:
ret_im = copy_image(im)
ret_im.facet_id = facet_id
ret_im.data *= scale
total_sumwt = sumwt
first = False
else:
ret_im.data += scale * im.data
total_sumwt += sumwt
return (ret_im, total_sumwt)
def test_invert_bag(self):
peaks = {
'2d': 65.440798589,
'timeslice': 99.7403479215,
'wstack': 100.654001673
}
vis_slices = {'2d': None, 'timeslice': 'auto', 'wstack': 101}
model = copy_image(self.model)
for context in ['wstack', '2d', 'timeslice']:
dirty_bag = invert_bag(self.vis_bag,
model,
dopsf=False,
context=context,
normalize=True,
vis_slices=vis_slices[context])
dirty, sumwt = list(dirty_bag)[0]
export_image_to_fits(
dirty,
'%s/test_bag_%s_dirty.fits' % (self.results_dir, context))
qa = qa_image(dirty, context=context)
assert numpy.abs(qa.data['max'] - peaks[context]) < 1.0e-7, str(qa)
def predict_wstack_single(vis, model, predict_inner=predict_2d_base, **kwargs):
""" 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
:returns: resulting visibility (in place works)
"""
if type(vis) is not Visibility:
avis = coalesce_visibility(vis, **kwargs)
else:
avis = vis
log.info("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)
# Do the real part
workimage.data = w_beam.data.real * model.data
avis = predict_inner(avis, workimage, **kwargs)
# and now the imaginary part
workimage.data = w_beam.data.imag * model.data
tempvis = predict_inner(tempvis, workimage, **kwargs)
avis.data['vis'] -= 1j * tempvis.data['vis']
avis.data['uvw'][...,2] += w_average
return avis
def w_kernel_list(vis: Visibility,
im: Image,
oversampling=1,
wstep=50.0,
kernelwidth=16,
**kwargs):
""" Calculate w convolution kernels
Uses create_w_term_like to calculate the w screen. This is exactly as wstacking does.
Returns (indices to the w kernel for each row, kernels)
Each kernel has axes [centre_v, centre_u, offset_v, offset_u]. We currently use the same
convolution function for all channels and polarisations. Changing that behaviour would
require modest changes here and to the gridding/degridding routines.
:param vis: visibility
:param image: Template image (padding, if any, occurs before this)
:param oversampling: Oversampling factor
:param wstep: Step in w between cached functions
:return: (indices to the w kernel for each row, kernels)
"""
nchan, npol, ny, nx = im.shape
gcf, _ = anti_aliasing_calculate((ny, nx))
assert oversampling % 2 == 0 or oversampling == 1, "oversampling must be unity or even"
assert kernelwidth % 2 == 0, "kernelwidth must be even"
wmaxabs = numpy.max(numpy.abs(vis.w))
log.debug(
"w_kernel_list: Maximum absolute w = %.1f, step is %.1f wavelengths" %
(wmaxabs, wstep))
def digitise(w, wstep):
return numpy.ceil((w + wmaxabs) / wstep).astype('int')
# Find all the unique indices for which we need a kernel
nwsteps = digitise(wmaxabs, wstep) + 1
w_list = numpy.linspace(-wmaxabs, +wmaxabs, nwsteps)
print('====', nwsteps, wstep, len(w_list))
wtemplate = copy_image(im)
wtemplate.data = numpy.zeros(wtemplate.shape, dtype=im.data.dtype)
padded_shape = list(wtemplate.shape)
padded_shape[3] *= oversampling
padded_shape[2] *= oversampling
# For all the unique indices, calculate the corresponding w kernel
kernels = list()
for w in w_list:
# Make a w screen
wscreen = create_w_term_like(wtemplate, w, vis.phasecentre, **kwargs)
wscreen.data /= gcf
assert numpy.max(numpy.abs(wscreen.data)) > 0.0, 'w screen is empty'
wscreen_padded = pad_image(wscreen, padded_shape)
wconv = fft_image(wscreen_padded)
wconv.data *= float(oversampling)**2
# For the moment, ignore the polarisation and channel axes
kernels.append(
convert_image_to_kernel(wconv, oversampling,
kernelwidth).data[0, 0, ...])
# Now make a lookup table from row number of vis to the kernel
kernel_indices = digitise(vis.w, wstep)
assert numpy.max(kernel_indices) < len(kernels), "wabsmax %f wstep %f" % (
wmaxabs, wstep)
assert numpy.min(kernel_indices) >= 0, "wabsmax %f wstep %f" % (wmaxabs,
wstep)
return kernel_indices, kernels
