Skip to content

Commit e343c0c

Browse files
timcornwelltimcornwell
committed
ska-pipelines fixes for Feng Wang
1 parent 08e0459 commit e343c0c

5 files changed

Lines changed: 75 additions & 40 deletions

File tree

‎data_models/data_model_helpers.py‎

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -506,16 +506,22 @@ def convert_skymodel_to_hdf(sm, f):
506506
"""
507507
f.attrs['ARL_data_model'] = 'SkyModel'
508508
f.attrs['fixed'] = sm.fixed
509-
f.attrs['number_skycomponents'] = len(sm.components)
510-
for i, sc in enumerate(sm.components):
511-
cf = f.create_group('skycomponent%d' % i)
512-
convert_skycomponent_to_hdf(sm.components[i], cf)
513-
cf = f.create_group('image')
514-
convert_image_to_hdf(sm.image, cf)
515-
cf = f.create_group('mask')
516-
convert_image_to_hdf(sm.mask, cf)
517-
cf = f.create_group('gaintable')
518-
convert_gaintable_to_hdf(sm.gaintable, cf)
509+
if sm.components is not None:
510+
f.attrs['number_skycomponents'] = len(sm.components)
511+
for i, sc in enumerate(sm.components):
512+
cf = f.create_group('skycomponent%d' % i)
513+
convert_skycomponent_to_hdf(sm.components[i], cf)
514+
else:
515+
f.attrs['number_skycomponents'] = len(sm.components)
516+
if sm.image is not None:
517+
cf = f.create_group('image')
518+
convert_image_to_hdf(sm.image, cf)
519+
if sm.mask is not None:
520+
cf = f.create_group('mask')
521+
convert_image_to_hdf(sm.mask, cf)
522+
if sm.gaintable is not None:
523+
cf = f.create_group('gaintable')
524+
convert_gaintable_to_hdf(sm.gaintable, cf)
519525
return f
520526

521527

‎tests/data_models/test_buffer_data_models.py‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,6 @@ def test_readwriteskymodel(self):
9797
gt = create_gaintable_from_blockvisibility(vis, timeslice='auto')
9898
gt = simulate_gaintable(gt, phase_error=1.0, amplitude_error=0.1)
9999
im = create_test_image()
100-
gt = create_gaintable_from_blockvisibility(vis)
101100
sm = SkyModel(components=[self.comp], image=im, gaintable=gt)
102101

103102
config = {"buffer": {"directory": self.dir},
@@ -113,6 +112,28 @@ def test_readwriteskymodel(self):
113112
assert newsm.gaintable.data.shape == gt.data.shape
114113
assert numpy.max(numpy.abs(newsm.image.data - im.data)) < 1e-15
115114

115+
def test_readwriteskymodel_no_image(self):
116+
vis = create_blockvisibility(self.lowcore, self.times, self.frequency,
117+
channel_bandwidth=self.channel_bandwidth,
118+
phasecentre=self.phasecentre,
119+
polarisation_frame=PolarisationFrame("linear"),
120+
weight=1.0)
121+
gt = create_gaintable_from_blockvisibility(vis, timeslice='auto')
122+
gt = simulate_gaintable(gt, phase_error=1.0, amplitude_error=0.1)
123+
sm = SkyModel(components=[self.comp], gaintable=gt)
124+
125+
config = {"buffer": {"directory": self.dir},
126+
"skymodel": {"name": "test_bufferskymodel.hdf", "data_model": "SkyModel"}}
127+
bdm = BufferSkyModel(config["buffer"], config["skymodel"], sm)
128+
bdm.sync()
129+
new_bdm = BufferSkyModel(config["buffer"], config["skymodel"])
130+
new_bdm.sync()
131+
newsm = bdm.memory_data_model
132+
133+
assert newsm.components[0].flux.shape == self.comp.flux.shape
134+
assert newsm.gaintable.data.shape == gt.data.shape
135+
136+
116137
def test_readwriteimage(self):
117138
im = create_test_image()
118139

‎workflows/scripts/pipelines/ska-pipelines/ska_continuum_imaging_arlexecute_pipeline.py‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
from workflows.arlexecute.pipelines.pipeline_arlexecute import continuum_imaging_list_arlexecute_workflow
2020

2121
from wrappers.arlexecute.execution_support.arlexecute import arlexecute
22+
from wrappers.arlexecute.visibility.coalesce import convert_visibility_to_blockvisibility, \
23+
convert_blockvisibility_to_visibility
2224

2325
import logging
2426

@@ -35,8 +37,8 @@ def init_logging():
3537
log = logging.getLogger()
3638
logging.info("Starting continuum imaging pipeline")
3739

38-
arlexecute.set_client(use_dask=True, threads_per_worker=1, memory_limit=4e9,
39-
local_dir=dask_dir)
40+
arlexecute.set_client(use_dask=True, threads_per_worker=1, memory_limit=32 * 1024 * 1024 * 1024, n_workers=8,
41+
local_dir=dask_dir, verbose=True)
4042
print(arlexecute.client)
4143
arlexecute.run(init_logging)
4244

@@ -46,9 +48,11 @@ def init_logging():
4648
centre = nfreqwin // 2
4749

4850
# Load data from previous simulation
49-
vis_list = [arlexecute.execute(import_blockvisibility_from_hdf5)
51+
block_vislist = [arlexecute.execute(import_blockvisibility_from_hdf5)
5052
(arl_path('%s/ska-pipeline_simulation_vislist_%d.hdf' % (results_dir, v)))
5153
for v in range(nfreqwin)]
54+
55+
vis_list = [arlexecute.execute(convert_blockvisibility_to_visibility, nout=1)(bv) for bv in block_vislist]
5256
print('Reading visibilities')
5357
vis_list = arlexecute.persist(vis_list)
5458

@@ -60,16 +64,18 @@ def init_logging():
6064
polarisation_frame=pol_frame)
6165
for v in vis_list]
6266

67+
print('Creating model images')
6368
model_list = arlexecute.persist(model_list)
6469

70+
print('Creating graph')
6571
continuum_imaging_list = continuum_imaging_list_arlexecute_workflow(vis_list, model_imagelist=model_list,
6672
context='wstack', vis_slices=51,
6773
scales=[0, 3, 10], algorithm='mmclean',
6874
nmoment=3, niter=1000,
6975
fractional_threshold=0.1,
7076
threshold=0.1, nmajor=5, gain=0.25,
71-
deconvolve_facets=8,
72-
deconvolve_overlap=32,
77+
deconvolve_facets=1,
78+
deconvolve_overlap=1,
7379
deconvolve_taper='tukey',
7480
timeslice='auto',
7581
psf_support=64)

‎workflows/scripts/pipelines/ska-pipelines/ska_imaging_arlexecute_pipeline.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
from workflows.arlexecute.imaging.imaging_arlexecute import invert_list_arlexecute_workflow
2020

2121
from wrappers.arlexecute.execution_support.arlexecute import arlexecute
22+
from wrappers.arlexecute.visibility.coalesce import convert_visibility_to_blockvisibility, \
23+
convert_blockvisibility_to_visibility
2224

2325
import logging
2426

@@ -36,7 +38,7 @@ def init_logging():
3638
log = logging.getLogger()
3739
logging.info("Starting Imaging pipeline")
3840

39-
arlexecute.set_client(use_dask=True, threads_per_worker=1, memory_limit=8e9,
41+
arlexecute.set_client(use_dask=True, threads_per_worker=1, memory_limit=32 * 1024 * 1024 * 1024, n_workers=8,
4042
local_dir=dask_dir)
4143
print(arlexecute.client)
4244
arlexecute.run(init_logging)
@@ -47,9 +49,11 @@ def init_logging():
4749
centre = nfreqwin // 2
4850

4951
# Load data from previous simulation
50-
vis_list = [arlexecute.execute(import_blockvisibility_from_hdf5)
51-
(arl_path('%s/ska-pipeline_simulation_vislist_%d.hdf' % (results_dir, v)))
52-
for v in range(nfreqwin)]
52+
block_vislist = [arlexecute.execute(import_blockvisibility_from_hdf5)
53+
(arl_path('%s/ska-pipeline_simulation_vislist_%d.hdf' % (results_dir, v)))
54+
for v in range(nfreqwin)]
55+
56+
vis_list = [arlexecute.execute(convert_blockvisibility_to_visibility, nout=1)(bv) for bv in block_vislist]
5357
print('Reading visibilities')
5458
vis_list = arlexecute.persist(vis_list)
5559

@@ -72,6 +76,4 @@ def init_logging():
7276

7377
arlexecute.close()
7478

75-
show_image(dirty, title='Dirty image', cm='Greys', vmax=0.1, vmin=-0.01)
76-
print(qa_image(dirty, context='Dirty image'))
7779
export_image_to_fits(dirty, '%s/ska-imaging_arlexecute_dirty.fits' % (results_dir))

‎workflows/scripts/pipelines/ska-pipelines/ska_simulate_arlexecute_pipeline.py‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
from wrappers.arlexecute.simulation.testing_support import create_low_test_image_from_gleam
2020
from wrappers.arlexecute.imaging.base import advise_wide_field
21+
from wrappers.arlexecute.visibility.coalesce import convert_blockvisibility_to_visibility, \
22+
convert_visibility_to_blockvisibility
2123

2224
from workflows.arlexecute.imaging.imaging_arlexecute import predict_list_arlexecute_workflow
2325
from workflows.arlexecute.simulation.simulation_arlexecute import simulate_list_arlexecute_workflow, \
@@ -37,9 +39,9 @@ def init_logging():
3739

3840
if __name__ == '__main__':
3941
log = logging.getLogger()
40-
print("Starting ska-pipelines pipeline")
42+
print("Starting ska-pipelines simulation pipeline")
4143

42-
arlexecute.set_client(use_dask=True, threads_per_worker=1, memory_limit=4e9,
44+
arlexecute.set_client(use_dask=True, threads_per_worker=1, memory_limit=32 * 1024 * 1024 * 1024, n_workers=8,
4345
local_dir=dask_dir)
4446
print(arlexecute.client)
4547
arlexecute.run(init_logging)
@@ -57,15 +59,17 @@ def init_logging():
5759
times = numpy.linspace(-numpy.pi / 3.0, numpy.pi / 3.0, ntimes)
5860
phasecentre = SkyCoord(ra=+30.0 * u.deg, dec=-60.0 * u.deg, frame='icrs', equinox='J2000')
5961

60-
vis_list = simulate_list_arlexecute_workflow('LOWBD2',
61-
rmax=rmax,
62-
frequency=frequency,
63-
channel_bandwidth=channel_bandwidth,
64-
times=times,
65-
phasecentre=phasecentre,
66-
order='frequency')
67-
print('%d elements in vis_list' % len(vis_list))
62+
blockvis_list = simulate_list_arlexecute_workflow('LOWBD2',
63+
rmax=rmax,
64+
frequency=frequency,
65+
channel_bandwidth=channel_bandwidth,
66+
times=times,
67+
phasecentre=phasecentre,
68+
order='frequency',
69+
format='blockvis')
70+
print('%d elements in vis_list' % len(blockvis_list))
6871
print('About to make visibility')
72+
vis_list = [arlexecute.execute(convert_blockvisibility_to_visibility, nout=1)(bv) for bv in blockvis_list]
6973
vis_list = arlexecute.persist(vis_list)
7074

7175
# The vis data are on the workers so we run the advice function on the workers
@@ -85,8 +89,6 @@ def init_logging():
8589

8690
# Now make a graph to fill with a model drawn from GLEAM
8791

88-
89-
9092
dprepb_model = [arlexecute.execute(create_low_test_image_from_gleam)(npixel=npixel,
9193
frequency=[frequency[f]],
9294
channel_bandwidth=[channel_bandwidth[f]],
@@ -100,14 +102,12 @@ def init_logging():
100102
# Put the model on the cluster
101103
dprepb_model = arlexecute.persist(dprepb_model)
102104

103-
print('About to make GLEAM model')
105+
print('About to make initial skymodel')
104106
zero_model = [arlexecute.execute(create_empty_image_like)(im) for im in dprepb_model]
105107
zero_model = arlexecute.compute(zero_model, sync=True)
106-
zero_skymodel = SkyModel(images=zero_model)
108+
zero_skymodel = SkyModel(image=zero_model)
107109
export_skymodel_to_hdf5(zero_skymodel, arl_path('%s/ska-pipeline_simulation_skymodel.hdf' % results_dir))
108110

109-
110-
111111
# vis_list = arlexecute.scatter(vis_list)
112112

113113
wstack = True
@@ -119,8 +119,8 @@ def init_logging():
119119
print('Using timeslicing with %d slices' % ntimes)
120120
predicted_vislist = predict_list_arlexecute_workflow(vis_list, dprepb_model, context='timeslice',
121121
vis_slices=ntimes)
122-
123-
corrupted_vislist = corrupt_list_arlexecute_workflow(predicted_vislist, phase_error=1.0, seed=180555)
122+
corrupted_vislist = [arlexecute.execute(convert_visibility_to_blockvisibility)(v) for v in predicted_vislist]
123+
corrupted_vislist = corrupt_list_arlexecute_workflow(corrupted_vislist, phase_error=1.0, seed=180555)
124124

125125
export_list = [arlexecute.execute(export_blockvisibility_to_hdf5)
126126
(corrupted_vislist[v], arl_path('%s/ska-pipeline_simulation_vislist_%d.hdf' % (results_dir, v)))

0 commit comments

Comments
 (0)