def imagerooter(image_list) -> list():
new_image_list = []
for im in image_list:
newim = create_empty_image_like(im)
newim.data = numpy.sqrt(numpy.abs(im.data))
new_image_list.append(newim)
return new_image_list
def invert_ignore_none(vis, model):
if vis is not None:
return invert(vis,
model,
context=context,
dopsf=dopsf,
normalize=normalize,
facets=facets,
vis_slices=vis_slices,
**kwargs)
else:
return create_empty_image_like(model), 0.0
def gather_image_iteration_results(results, template_model):
result = create_empty_image_like(template_model)
i = 0
sumwt = numpy.zeros([template_model.nchan, template_model.npol])
for dpatch in image_scatter_facets(result, facets=facets):
assert i < len(
results), "Too few results in gather_image_iteration_results"
if results[i] is not None:
assert len(results[i]) == 2, results[i]
dpatch.data[...] = results[i][0].data[...]
sumwt += results[i][1]
i += 1
return result, sumwt
def smooth_image(model: Image, width=1.0):
""" Smooth an image with a kernel
"""
import astropy.convolution
assert isinstance(model, Image), model
kernel = astropy.convolution.kernels.Gaussian2DKernel(width)
cmodel = create_empty_image_like(model)
nchan, npol, _, _ = model.shape
for pol in range(npol):
for chan in range(nchan):
cmodel.data[chan, pol, :, :] = astropy.convolution.convolve(
model.data[chan, pol, :, :], kernel, normalize_kernel=False)
if isinstance(kernel, astropy.convolution.kernels.Gaussian2DKernel):
cmodel.data *= 2 * numpy.pi * width**2
return cmodel
# Create test image
frequency = numpy.array([1e8])
phasecentre = SkyCoord(ra=+15.0 * u.deg,
dec=-35.0 * u.deg,
frame='icrs',
equinox='J2000')
model = create_test_image(frequency=frequency,
phasecentre=phasecentre,
cellsize=0.001,
polarisation_frame=PolarisationFrame('stokesI'))
# Rank 0 scatters the test image
if rank == 0:
subimages = image_scatter_facets(model, facets=facets)
subimages = numpy.array_split(subimages, size)
else:
subimages = list()
sublist = comm.scatter(subimages, root=0)
root_images = imagerooter(sublist)
roots = comm.gather(root_images, root=0)
if rank == 0:
results = sum(roots, [])
root_model = create_empty_image_like(model)
result = image_gather_facets(results, root_model, facets=facets)
numpy.testing.assert_array_almost_equal_nulp(result.data**2,
numpy.abs(model.data), 7)
