Skip to content

Commit 59085d9

Browse files
timcornwelltimcornwell
committed
Merge some changes from master
1 parent 885f20a commit 59085d9

2 files changed

Lines changed: 168 additions & 35 deletions

File tree

‎tests/workflows/test_imaging_arlexecute.py‎

Lines changed: 70 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515
from processing_components.griddata.kernels import create_awterm_convolutionfunction
1616
from workflows.arlexecute.imaging.imaging_arlexecute import zero_list_arlexecute_workflow, \
1717
predict_list_arlexecute_workflow, invert_list_arlexecute_workflow, subtract_list_arlexecute_workflow, \
18-
weight_list_arlexecute_workflow, residual_list_arlexecute_workflow, sum_invert_results_arlexecute
18+
weight_list_arlexecute_workflow, residual_list_arlexecute_workflow, sum_invert_results_arlexecute, \
19+
restore_list_arlexecute_workflow
1920
from workflows.shared.imaging.imaging_shared import sum_invert_results, sum_invert_results_local
2021
from wrappers.arlexecute.execution_support.arlexecutebase import ARLExecuteBase
2122
from wrappers.arlexecute.execution_support.dask_init import get_dask_Client
@@ -48,6 +49,8 @@ def setUp(self):
4849
from data_models.parameters import arl_path
4950
self.dir = arl_path('test_results')
5051

52+
self.persist = False
53+
5154
def tearDown(self):
5255
global arlexecute
5356
arlexecute.close()
@@ -126,8 +129,8 @@ def actualSetUp(self, add_errors=False, freqwin=3, block=False, dospectral=True,
126129
self.model = self.model_list[centre]
127130

128131
self.cmodel = smooth_image(self.model)
129-
export_image_to_fits(self.model, '%s/test_imaging_model.fits' % self.dir)
130-
export_image_to_fits(self.cmodel, '%s/test_imaging_cmodel.fits' % self.dir)
132+
if self.persist: export_image_to_fits(self.model, '%s/test_imaging_model.fits' % self.dir)
133+
if self.persist: export_image_to_fits(self.cmodel, '%s/test_imaging_cmodel.fits' % self.dir)
131134

132135
if add_errors and block:
133136
self.vis_list = [arlexecute.execute(insert_unittest_errors)(self.vis_list[i])
@@ -184,7 +187,7 @@ def _predict_base(self, context='2d', extra='', fluxthreshold=1.0, facets=1, vis
184187
dirty = arlexecute.compute(dirty, sync=True)[centre]
185188

186189
assert numpy.max(numpy.abs(dirty[0].data)), "Residual image is empty"
187-
export_image_to_fits(dirty[0], '%s/test_imaging_predict_%s%s_%s_dirty.fits' %
190+
if self.persist: export_image_to_fits(dirty[0], '%s/test_imaging_predict_%s%s_%s_dirty.fits' %
188191
(self.dir, context, extra, arlexecute.type()))
189192

190193
maxabs = numpy.max(numpy.abs(dirty[0].data))
@@ -200,7 +203,7 @@ def _invert_base(self, context, extra='', fluxthreshold=1.0, positionthreshold=1
200203
dirty = arlexecute.compute(dirty, sync=True)[centre]
201204

202205
print(dirty)
203-
export_image_to_fits(dirty[0], '%s/test_imaging_invert_%s%s_%s_dirty.fits' %
206+
if self.persist: export_image_to_fits(dirty[0], '%s/test_imaging_invert_%s%s_%s_dirty.fits' %
204207
(self.dir, context, extra, arlexecute.type()))
205208

206209
assert numpy.max(numpy.abs(dirty[0].data)), "Image is empty"
@@ -387,6 +390,68 @@ def test_residual_list(self):
387390
assert numpy.abs(qa.data['max'] - 0.35139716991480785) < 1.0, str(qa)
388391
assert numpy.abs(qa.data['min'] + 0.7681701460717593) < 1.0, str(qa)
389392

393+
def test_restored_list(self):
394+
self.actualSetUp(zerow=True)
395+
396+
centre = self.freqwin // 2
397+
psf_image_list = invert_list_arlexecute_workflow(self.vis_list, self.model_list, context='2d', dopsf=True)
398+
residual_image_list = residual_list_arlexecute_workflow(self.vis_list, self.model_list, context='2d')
399+
restored_image_list = restore_list_arlexecute_workflow(self.model_list, psf_image_list, residual_image_list,
400+
psfwidth=1.0)
401+
restored_image_list = arlexecute.compute(restored_image_list, sync=True)
402+
if self.persist: export_image_to_fits(restored_image_list[centre], '%s/test_imaging_invert_%s_restored.fits' %
403+
(self.dir, arlexecute.type()))
404+
405+
qa = qa_image(restored_image_list[centre])
406+
assert numpy.abs(qa.data['max'] - 99.43438263927834) < 1e-7, str(qa)
407+
assert numpy.abs(qa.data['min'] + 0.6328915148563365) < 1e-7, str(qa)
408+
409+
def test_restored_list_noresidual(self):
410+
self.actualSetUp(zerow=True)
411+
412+
centre = self.freqwin // 2
413+
psf_image_list = invert_list_arlexecute_workflow(self.vis_list, self.model_list, context='2d', dopsf=True)
414+
restored_image_list = restore_list_arlexecute_workflow(self.model_list, psf_image_list, psfwidth=1.0)
415+
restored_image_list = arlexecute.compute(restored_image_list, sync=True)
416+
if self.persist: export_image_to_fits(restored_image_list[centre],
417+
'%s/test_imaging_invert_%s_restored_noresidual.fits' %
418+
(self.dir, arlexecute.type()))
419+
420+
qa = qa_image(restored_image_list[centre])
421+
assert numpy.abs(qa.data['max'] - 100.0) < 1e-7, str(qa)
422+
assert numpy.abs(qa.data['min']) < 1e-7, str(qa)
423+
424+
def test_restored_list_facet(self):
425+
self.actualSetUp(zerow=True)
426+
427+
centre = self.freqwin // 2
428+
psf_image_list = invert_list_arlexecute_workflow(self.vis_list, self.model_list, context='2d', dopsf=True)
429+
residual_image_list = residual_list_arlexecute_workflow(self.vis_list, self.model_list, context='2d')
430+
restored_4facets_image_list = restore_list_arlexecute_workflow(self.model_list, psf_image_list,
431+
residual_image_list,
432+
restore_facets=4, psfwidth=1.0)
433+
restored_4facets_image_list = arlexecute.compute(restored_4facets_image_list, sync=True)
434+
435+
restored_1facets_image_list = restore_list_arlexecute_workflow(self.model_list, psf_image_list,
436+
residual_image_list,
437+
restore_facets=1, psfwidth=1.0)
438+
restored_1facets_image_list = arlexecute.compute(restored_1facets_image_list, sync=True)
439+
440+
if self.persist: export_image_to_fits(restored_4facets_image_list[0],
441+
'%s/test_imaging_invert_%s_restored_4facets.fits' %
442+
(self.dir, arlexecute.type()))
443+
444+
qa = qa_image(restored_4facets_image_list[centre])
445+
assert numpy.abs(qa.data['max'] - 99.43438263927833) < 1e-7, str(qa)
446+
assert numpy.abs(qa.data['min'] + 0.6328915148563354) < 1e-7, str(qa)
447+
448+
restored_4facets_image_list[centre].data -= restored_1facets_image_list[centre].data
449+
if self.persist: export_image_to_fits(restored_4facets_image_list[centre],
450+
'%s/test_imaging_invert_%s_restored_4facets_error.fits' %
451+
(self.dir, arlexecute.type()))
452+
qa = qa_image(restored_4facets_image_list[centre])
453+
assert numpy.abs(qa.data['max']) < 1e-10, str(qa)
454+
390455
def test_sum_invert_list(self):
391456
self.actualSetUp(zerow=True)
392457

‎workflows/arlexecute/imaging/imaging_arlexecute.py‎

Lines changed: 98 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,6 @@ def predict_ignore_none(vis, model, g):
9595
g = gcfcf[ivis]
9696
else:
9797
g = gcfcf[0]
98-
# Create the graph to divide an image into facets. This is by reference.
9998
# Create the graph to divide the visibility into slices. This is by copy.
10099
sub_vis_lists = arlexecute.execute(visibility_scatter, nout=vis_slices)(subvis,
101100
vis_iter, vis_slices)
@@ -271,29 +270,62 @@ def residual_list_arlexecute_workflow(vis, model_imagelist, context='2d', gcfcf=
271270
return arlexecute.optimize(result)
272271

273272

274-
def restore_list_arlexecute_workflow(model_imagelist, psf_imagelist, residual_imagelist=None, **kwargs):
273+
def restore_list_arlexecute_workflow(model_imagelist, psf_imagelist, residual_imagelist=None, restore_facets=1,
274+
restore_overlap=0, restore_taper='tukey', **kwargs):
275275
""" Create a graph to calculate the restored image
276276
277277
:param model_imagelist: Model list
278278
:param psf_imagelist: PSF list
279279
:param residual_imagelist: Residual list
280280
:param kwargs: Parameters for functions in components
281+
:param restore_facets: Number of facets used per axis (used to distribute)
282+
:param restore_overlap: Overlap in pixels (0 is best)
283+
:param restore_taper: Type of taper.
281284
:return:
282285
"""
283-
if residual_imagelist is None:
284-
residual_imagelist = []
286+
assert len(model_imagelist) == len(psf_imagelist)
287+
if residual_imagelist is not None:
288+
assert len(model_imagelist) == len(residual_imagelist)
289+
290+
if restore_facets % 2 == 0 or restore_facets == 1:
291+
actual_number_facets = restore_facets
292+
else:
293+
actual_number_facets = max(1, (restore_facets - 1))
285294

286295
psf_list = arlexecute.execute(remove_sumwt, nout=len(psf_imagelist))(psf_imagelist)
287-
if len(residual_imagelist) > 0:
296+
297+
# Scatter each list element into a list. We will then run restore_cube on each
298+
facet_model_list = [arlexecute.execute(image_scatter_facets, nout=actual_number_facets * actual_number_facets)
299+
(model, facets=restore_facets, overlap=restore_overlap, taper=restore_taper)
300+
for model in model_imagelist]
301+
facet_psf_list = [arlexecute.execute(image_scatter_facets, nout=actual_number_facets * actual_number_facets)
302+
(psf, facets=restore_facets, overlap=restore_overlap, taper=restore_taper)
303+
for psf in psf_list]
304+
305+
if residual_imagelist is not None:
288306
residual_list = arlexecute.execute(remove_sumwt, nout=len(residual_imagelist))(residual_imagelist)
289-
result = [arlexecute.execute(restore_cube)(model_imagelist[i], psf_list[i],
290-
residual_list[i], **kwargs)
291-
for i, _ in enumerate(model_imagelist)]
307+
facet_residual_list = [
308+
arlexecute.execute(image_scatter_facets, nout=actual_number_facets * actual_number_facets)
309+
(residual, facets=restore_facets, overlap=restore_overlap, taper=restore_taper)
310+
for residual in residual_list]
311+
facet_restored_list = [[arlexecute.execute(restore_cube, nout=actual_number_facets * actual_number_facets)
312+
(model=facet_model_list[i][im], psf=facet_psf_list[i][im],
313+
residual=facet_residual_list[i][im],
314+
**kwargs)
315+
for im, _ in enumerate(facet_model_list[i])] for i, _ in enumerate(model_imagelist)]
292316
else:
293-
result = [arlexecute.execute(restore_cube)(model_imagelist[i], psf_list[i], **kwargs)
317+
facet_restored_list = [[arlexecute.execute(restore_cube, nout=actual_number_facets * actual_number_facets)
318+
(model=facet_model_list[i][im], psf=facet_psf_list[i][im],
319+
**kwargs)
320+
for im, _ in enumerate(facet_model_list[i])] for i, _ in enumerate(model_imagelist)]
321+
322+
# Now we run restore_cube on each and gather the results across all facets
323+
restored_imagelist = [arlexecute.execute(image_gather_facets)
324+
(facet_restored_list[i], model_imagelist[i], facets=restore_facets,
325+
overlap=restore_overlap, taper=restore_taper)
294326
for i, _ in enumerate(model_imagelist)]
295327

296-
return arlexecute.optimize(result)
328+
return arlexecute.optimize(restored_imagelist)
297329

298330

299331
def deconvolve_list_arlexecute_workflow(dirty_list, psf_list, model_imagelist, prefix='', mask=None, **kwargs):
@@ -347,30 +379,51 @@ def deconvolve(dirty, psf, model, facet, gthreshold, msk=None):
347379
else:
348380
deconvolve_number_facets = deconvolve_facets ** 2
349381

350-
deconvolve_model_imagelist = arlexecute.execute(image_gather_channels, nout=1)(model_imagelist)
382+
scattered_channels_facets_model_list = \
383+
[arlexecute.execute(image_scatter_facets, nout=deconvolve_number_facets)(m, facets=deconvolve_facets,
384+
overlap=deconvolve_overlap,
385+
taper=deconvolve_taper)
386+
for m in model_imagelist]
387+
scattered_facets_model_list = [
388+
arlexecute.execute(image_gather_channels, nout=1)([scattered_channels_facets_model_list[chan][facet]
389+
for chan in range(nchan)])
390+
for facet in range(deconvolve_number_facets)]
351391

352392
# Scatter the separate channel images into deconvolve facets and then gather channels for each facet.
353393
# This avoids constructing the entire spectral cube.
394+
# i.e. SCATTER BY FACET then SCATTER BY CHANNEL
354395
dirty_list_trimmed = arlexecute.execute(remove_sumwt, nout=nchan)(dirty_list)
355396
scattered_channels_facets_dirty_list = \
356397
[arlexecute.execute(image_scatter_facets, nout=deconvolve_number_facets)(d, facets=deconvolve_facets,
357398
overlap=deconvolve_overlap,
358399
taper=deconvolve_taper)
359400
for d in dirty_list_trimmed]
360-
361-
# Now we do a transpose and gather
362-
scattered_facets_list = [
401+
scattered_facets_dirty_list = [
363402
arlexecute.execute(image_gather_channels, nout=1)([scattered_channels_facets_dirty_list[chan][facet]
364403
for chan in range(nchan)])
365404
for facet in range(deconvolve_number_facets)]
366405

367406
psf_list_trimmed = arlexecute.execute(remove_sumwt, nout=nchan)(psf_list)
368-
psf_list_trimmed = arlexecute.execute(image_gather_channels, nout=1)(psf_list_trimmed)
369407

370-
scattered_model_imagelist = \
371-
arlexecute.execute(image_scatter_facets, nout=deconvolve_number_facets)(deconvolve_model_imagelist,
372-
facets=deconvolve_facets,
373-
overlap=deconvolve_overlap)
408+
def extract_psf(psf, facets):
409+
spsf = create_empty_image_like(psf)
410+
cx = spsf.shape[3] // 2
411+
cy = spsf.shape[2] // 2
412+
wx = spsf.shape[3] // facets
413+
wy = spsf.shape[2] // facets
414+
xbeg = cx - wx // 2
415+
xend = cx + wx // 2
416+
ybeg = cy - wy // 2
417+
yend = cy + wy // 2
418+
spsf.data = psf.data[..., ybeg:yend, xbeg:xend]
419+
spsf.wcs.wcs.crpix[0] -= xbeg
420+
spsf.wcs.wcs.crpix[1] -= ybeg
421+
return spsf
422+
423+
psf_list_trimmed = [arlexecute.execute(extract_psf)(p, deconvolve_facets) for p in psf_list_trimmed]
424+
psf_centre = arlexecute.execute(image_gather_channels, nout=1)([psf_list_trimmed[chan]
425+
for chan in range(nchan)])
426+
374427
# Work out the threshold. Need to find global peak over all dirty_list images
375428
threshold = get_parameter(kwargs, "threshold", 0.0)
376429
fractional_threshold = get_parameter(kwargs, "fractional_threshold", 0.1)
@@ -379,32 +432,47 @@ def deconvolve(dirty, psf, model, facet, gthreshold, msk=None):
379432

380433
# Find the global threshold. This uses the peak in the average on the frequency axis since we
381434
# want to use it in a stopping criterion in a moment clean
382-
global_threshold = arlexecute.execute(threshold_list, nout=1)(scattered_facets_list, threshold,
435+
global_threshold = arlexecute.execute(threshold_list, nout=1)(scattered_facets_dirty_list, threshold,
383436
fractional_threshold,
384437
use_moment0=use_moment0, prefix=prefix)
385438

386439
facet_list = numpy.arange(deconvolve_number_facets).astype('int')
387440
if mask is None:
388441
scattered_results_list = [
389-
arlexecute.execute(deconvolve, nout=1)(d, psf_list_trimmed, m, facet, global_threshold)
390-
for d, m, facet in zip(scattered_facets_list, scattered_model_imagelist, facet_list)]
442+
arlexecute.execute(deconvolve, nout=1)(d, psf_centre, m, facet,
443+
global_threshold)
444+
for d, m, facet in zip(scattered_facets_dirty_list, scattered_facets_model_list, facet_list)]
391445
else:
392446
mask_list = \
393447
arlexecute.execute(image_scatter_facets, nout=deconvolve_number_facets)(mask,
394448
facets=deconvolve_facets,
395449
overlap=deconvolve_overlap)
396450
scattered_results_list = [
397-
arlexecute.execute(deconvolve, nout=1)(d, psf_list_trimmed, m, facet, global_threshold, msk)
398-
for d, m, facet, msk in zip(scattered_facets_list, scattered_model_imagelist, facet_list, mask_list)]
451+
arlexecute.execute(deconvolve, nout=1)(d, psf_centre, m, facet,
452+
global_threshold, msk)
453+
for d, m, facet, msk in
454+
zip(scattered_facets_dirty_list, scattered_facets_model_list, facet_list, mask_list)]
399455

456+
# We want to avoid constructing the entire cube so we do the inverse of how we got here:
457+
# i.e. SCATTER BY CHANNEL then GATHER BY FACET
400458
# Gather the results back into one image, correcting for overlaps as necessary. The taper function is is used to
401459
# feather the facets together
402-
gathered_results_list = arlexecute.execute(image_gather_facets, nout=1)(scattered_results_list,
403-
deconvolve_model_imagelist,
404-
facets=deconvolve_facets,
405-
overlap=deconvolve_overlap,
406-
taper=deconvolve_taper)
407-
result_list = arlexecute.execute(image_scatter_channels, nout=nchan)(gathered_results_list, subimages=nchan)
460+
# gathered_results_list = arlexecute.execute(image_gather_facets, nout=1)(scattered_results_list,
461+
# deconvolve_model_imagelist,
462+
# facets=deconvolve_facets,
463+
# overlap=deconvolve_overlap,
464+
# taper=deconvolve_taper)
465+
# result_list = arlexecute.execute(image_scatter_channels, nout=nchan)(gathered_results_list, subimages=nchan)
466+
467+
scattered_channel_results_list = [arlexecute.execute(image_scatter_channels, nout=nchan)(scat, subimages=nchan)
468+
for scat in scattered_results_list]
469+
470+
# The structure is now [[channels] for facets]. We do the reverse transpose to the one above.
471+
result_list = [arlexecute.execute(image_gather_facets, nout=1)([scattered_channel_results_list[facet][chan]
472+
for facet in range(deconvolve_number_facets)],
473+
model_imagelist[chan], facets=deconvolve_facets,
474+
overlap=deconvolve_overlap)
475+
for chan in range(nchan)]
408476

409477
return arlexecute.optimize(result_list)
410478

0 commit comments

Comments
 (0)