def sum_invert_bag_results(invert_list, normalize=True):
"""Sum a list of invert results, optionally normalizing at the end
Can be used in bag.map()
:param invert_list: List of results from invert: Image, weight tuples
:param normalize: Normalize by the sum of weights
"""
assert isinstance(invert_list, collections.Iterable), invert_list
result = None
weight = None
for i, a in enumerate(invert_list):
assert isinstance(a[0], Image), "Item is not an image: %s" % str(a[0])
if i == 0:
result = create_image_from_array(a[0].data * a[1], a[0].wcs, a[0].polarisation_frame)
weight = a[1]
else:
result.data += a[0].data * a[1]
weight += a[1]
assert weight is not None and result is not None, "No valid images found"
if normalize:
result = normalize_sumwt(result, weight)
return result, weight
def invert_with_raster_iterator(vis, im, dopsf=False,
normalize=True, invert_function=invert_2d_base, facets=1,
**kwargs) -> (Image, numpy.ndarray):
""" Predict using image partitions, calling specified predict function
:param vis: Visibility to be inverted
:param im: image template (not changed)
:param image_iterator: Iterator to use for partitioning
:param dopsf: Make the psf instead of the dirty image
:param normalize: Normalize by the sum of weights (True)
:return: resulting image[nchan, npol, ny, nx], sum of weights[nchan, npol]
"""
log.info("invert_with_image_iterator: Inverting by image partitions")
i = 0
nchan, npol, _, _ = im.shape
totalwt = numpy.zeros([nchan, npol])
for dpatch in image_raster_iter(im, facets=facets):
result, sumwt = invert_function(vis, dpatch, dopsf, normalize=False, **kwargs)
totalwt = sumwt
# Ensure that we fill in the elements of dpatch instead of creating a new numpy arrray
dpatch.data[...] = result.data[...]
assert numpy.max(numpy.abs(dpatch.data)), "Partition image %d appears to be empty" % i
i += 1
assert numpy.max(numpy.abs(im.data)), "Output image appears to be empty"
if normalize:
im = normalize_sumwt(im, totalwt)
return im, totalwt
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 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]
im = normalize_sumwt(im, sumwt)
return im, sumwt
def invert_with_vis_iterator(vis: Visibility, im: Image, dopsf=False, normalize=True, vis_iter=vis_slice_iter,
invert=invert_2d, **kwargs):
""" Invert using a specified iterator and invert
This knows about the structure of invert in different execution frameworks but not
anything about the actual processing.
:param vis:
:param im:
:param dopsf: Make the psf instead of the dirty image
:param normalize: Normalize by the sum of weights (True)
:param kwargs:
:return:
"""
resultimage = create_empty_image_like(im)
if type(vis) is not Visibility:
svis = coalesce_visibility(vis, **kwargs)
else:
svis = vis
i = 0
for rows in vis_iter(svis, **kwargs):
if numpy.sum(rows) and svis is not None:
visslice = create_visibility_from_rows(svis, rows)
workimage, sumwt = invert(visslice, im, dopsf, normalize=False, **kwargs)
resultimage.data += workimage.data
if i == 0:
totalwt = sumwt
else:
totalwt += sumwt
i += 1
if normalize:
resultimage = normalize_sumwt(resultimage, totalwt)
return resultimage, totalwt
def invert_function(vis,
im: Image,
dopsf=False,
normalize=True,
context='2d',
inner=None,
**kwargs):
""" Invert using algorithm specified by context:
* 2d: Two-dimensional transform
* wstack: wstacking with either vis_slices or wstack (spacing between w planes) set
* wprojection: w projection with wstep (spacing between w places) set, also kernel='wprojection'
* timeslice: snapshot imaging with either vis_slices or timeslice set. timeslice='auto' does every time
* facets: Faceted imaging with facets facets on each axis
* facets_wprojection: facets AND wprojection
* facets_wstack: facets AND wstacking
* wprojection_wstack: wprojection and wstacking
:param vis:
:param im:
:param dopsf: Make the psf instead of the dirty image (False)
:param normalize: Normalize by the sum of weights (True)
:param context: Imaging context e.g. '2d', 'timeslice', etc.
:param inner: Inner loop 'vis'|'image'
:param kwargs:
:return: Image, sum of weights
"""
c = imaging_context(context)
vis_iter = c['vis_iterator']
image_iter = c['image_iterator']
invert = c['invert']
if inner is None:
inner = c['inner']
if not isinstance(vis, Visibility):
svis = coalesce_visibility(vis, **kwargs)
else:
svis = vis
resultimage = create_empty_image_like(im)
if inner == 'image':
totalwt = None
for rows in vis_iter(svis, **kwargs):
if numpy.sum(rows):
visslice = create_visibility_from_rows(svis, rows)
sumwt = 0.0
workimage = create_empty_image_like(im)
for dpatch in image_iter(workimage, **kwargs):
result, sumwt = invert(visslice,
dpatch,
dopsf,
normalize=False,
**kwargs)
# Ensure that we fill in the elements of dpatch instead of creating a new numpy arrray
dpatch.data[...] = result.data[...]
# Assume that sumwt is the same for all patches
if totalwt is None:
totalwt = sumwt
else:
totalwt += sumwt
resultimage.data += workimage.data
else:
# We assume that the weight is the same for all image iterations
totalwt = None
workimage = create_empty_image_like(im)
for dpatch in image_iter(workimage, **kwargs):
totalwt = None
for rows in vis_iter(svis, **kwargs):
if numpy.sum(rows):
visslice = create_visibility_from_rows(svis, rows)
result, sumwt = invert(visslice,
dpatch,
dopsf,
normalize=False,
**kwargs)
# Ensure that we fill in the elements of dpatch instead of creating a new numpy arrray
dpatch.data[...] += result.data[...]
if totalwt is None:
totalwt = sumwt
else:
totalwt += sumwt
resultimage.data += workimage.data
workimage.data[...] = 0.0
assert totalwt is not None, "No valid data found for imaging"
if normalize:
resultimage = normalize_sumwt(resultimage, totalwt)
return resultimage, totalwt
