def test_blockwise_clone_with_literals(literal):
arr = da.ones(10, chunks=1)
def noop(arr, lit):
return arr
blk = da.blockwise(noop, "x", arr, "x", literal, None)
cln = clone(blk)
assert_no_common_keys(blk, cln, layers=True)
def test_clone(layers):
dsk1 = {("a", h1): 1, ("a", h2): 2}
dsk2 = {"b": (add, ("a", h1), ("a", h2))}
dsk3 = {"c": 1, "d": 1} # Multiple names
if layers:
dsk1 = HighLevelGraph.from_collections("a", dsk1)
dsk2 = HighLevelGraph({
"a": dsk1,
"b": dsk2
},
dependencies={
"a": set(),
"b": {"a"}
})
dsk3 = HighLevelGraph.from_collections("c", dsk3)
else:
dsk2.update(dsk1)
t1 = Tuple(dsk1, [("a", h1), ("a", h2)])
t2 = Tuple(dsk2, ["b"])
t3 = Tuple(dsk3, ["c"])
c1 = clone(t2, seed=1, assume_layers=layers)
c2 = clone(t2, seed=1, assume_layers=layers)
c3 = clone(t2, seed=2, assume_layers=layers)
c4 = clone(c1, seed=1,
assume_layers=layers) # Clone of a clone has different keys
c5 = clone(t2, assume_layers=layers) # Random seed
c6 = clone(t2, assume_layers=layers) # Random seed
c7 = clone(t2, omit=t1, seed=1, assume_layers=layers)
assert c1.__dask_graph__() == c2.__dask_graph__()
assert_no_common_keys(c1, t2, layers=layers)
assert_no_common_keys(c1, c3, layers=layers)
assert_no_common_keys(c1, c4, layers=layers)
assert_no_common_keys(c1, c5, layers=layers)
assert_no_common_keys(c5, c6, layers=layers)
assert_no_common_keys(c7, t2, omit=t1, layers=layers)
assert dask.compute(t2, c1, c2, c3, c4, c5, c6, c7) == ((3, ), ) * 8
# Clone nested; some of the collections in omit are unrelated
out = clone({"x": [t2]}, omit={"y": [t1, t3]}, assume_layers=layers)
assert dask.compute(out) == ({"x": [(3, )]}, )
c8 = out["x"][0]
assert_no_common_keys(c8, t2, omit=t1, layers=layers)
assert_no_common_keys(c8, t3, layers=layers)
def _predict(ms, stack, **kw):
args = OmegaConf.create(kw)
OmegaConf.set_struct(args, True)
pyscilog.log_to_file(args.output_filename + '.log')
pyscilog.enable_memory_logging(level=3)
# number of threads per worker
if args.nthreads is None:
if args.host_address is not None:
raise ValueError(
"You have to specify nthreads when using a distributed scheduler"
)
import multiprocessing
nthreads = multiprocessing.cpu_count()
args.nthreads = nthreads
else:
nthreads = args.nthreads
if args.mem_limit is None:
if args.host_address is not None:
raise ValueError(
"You have to specify mem-limit when using a distributed scheduler"
)
import psutil
mem_limit = int(psutil.virtual_memory()[0] /
1e9) # 100% of memory by default
args.mem_limit = mem_limit
else:
mem_limit = args.mem_limit
nband = args.nband
if args.nworkers is None:
nworkers = nband
args.nworkers = nworkers
else:
nworkers = args.nworkers
if args.nthreads_per_worker is None:
nthreads_per_worker = 1
args.nthreads_per_worker = nthreads_per_worker
else:
nthreads_per_worker = args.nthreads_per_worker
# the number of chunks being read in simultaneously is equal to
# the number of dask threads
nthreads_dask = nworkers * nthreads_per_worker
if args.ngridder_threads is None:
if args.host_address is not None:
ngridder_threads = nthreads // nthreads_per_worker
else:
ngridder_threads = nthreads // nthreads_dask
args.ngridder_threads = ngridder_threads
else:
ngridder_threads = args.ngridder_threads
ms = list(ms)
print('Input Options:', file=log)
for key in kw.keys():
print(' %25s = %s' % (key, args[key]), file=log)
# numpy imports have to happen after this step
from pfb import set_client
set_client(nthreads, mem_limit, nworkers, nthreads_per_worker,
args.host_address, stack, log)
import numpy as np
from pfb.utils.misc import chan_to_band_mapping
import dask
from dask.distributed import performance_report
from dask.graph_manipulation import clone
from daskms import xds_from_storage_ms as xds_from_ms
from daskms import xds_from_storage_table as xds_from_table
from daskms.utils import dataset_type
mstype = dataset_type(ms[0])
if mstype == 'casa':
from daskms import xds_to_table
elif mstype == 'zarr':
from daskms.experimental.zarr import xds_to_zarr as xds_to_table
import dask.array as da
from africanus.constants import c as lightspeed
from africanus.gridding.wgridder.dask import model as im2vis
from pfb.utils.fits import load_fits
from pfb.utils.misc import restore_corrs, plan_row_chunk
from astropy.io import fits
# always returns 4D
# gridder expects freq axis
model = np.atleast_3d(load_fits(args.model).squeeze())
nband, nx, ny = model.shape
hdr = fits.getheader(args.model)
cell_d = np.abs(hdr['CDELT1'])
cell_rad = np.deg2rad(cell_d)
# chan <-> band mapping
freqs, freq_bin_idx, freq_bin_counts, freq_out, band_mapping, chan_chunks = chan_to_band_mapping(
ms, nband=nband)
# degridder memory budget
max_chan_chunk = 0
for ims in ms:
for spw in freqs[ims]:
counts = freq_bin_counts[ims][spw].compute()
max_chan_chunk = np.maximum(max_chan_chunk, counts.max())
# assumes number of correlations are the same across MS/SPW
xds = xds_from_ms(ms[0])
ncorr = xds[0].dims['corr']
nrow = xds[0].dims['row']
if args.output_type is not None:
output_type = np.dtype(args.output_type)
else:
output_type = np.result_type(np.dtype(args.real_type), np.complex64)
data_bytes = output_type.itemsize
bytes_per_row = max_chan_chunk * ncorr * data_bytes
memory_per_row = bytes_per_row # model
memory_per_row += 3 * 8 # uvw
if mstype == 'zarr':
if args.model_column in xds[0].keys():
model_chunks = getattr(xds[0], args.model_column).data.chunks
else:
model_chunks = xds[0].DATA.data.chunks
print('Chunking model same as data')
# get approx image size
# this is not a conservative estimate when multiple SPW's map to a single
# imaging band
pixel_bytes = np.dtype(args.output_type).itemsize
band_size = nx * ny * pixel_bytes
if args.host_address is None:
# full image on single node
row_chunk = plan_row_chunk(mem_limit / nworkers, band_size, nrow,
memory_per_row, nthreads_per_worker)
else:
# single band per node
row_chunk = plan_row_chunk(mem_limit, band_size, nrow, memory_per_row,
nthreads_per_worker)
if args.row_chunks is not None:
row_chunk = int(args.row_chunks)
if row_chunk == -1:
row_chunk = nrow
print(
"nrows = %i, row chunks set to %i for a total of %i chunks per node" %
(nrow, row_chunk, int(np.ceil(nrow / row_chunk))),
file=log)
chunks = {}
for ims in ms:
chunks[ims] = [] # xds_from_ms expects a list per ds
for spw in freqs[ims]:
chunks[ims].append({
'row': row_chunk,
'chan': chan_chunks[ims][spw]['chan']
})
model = da.from_array(model.astype(args.real_type),
chunks=(1, nx, ny),
name=False)
writes = []
radec = None # assumes we are only imaging field 0 of first MS
for ims in ms:
xds = xds_from_ms(ims, chunks=chunks[ims], columns=('UVW'))
# subtables
ddids = xds_from_table(ims + "::DATA_DESCRIPTION")
fields = xds_from_table(ims + "::FIELD")
spws = xds_from_table(ims + "::SPECTRAL_WINDOW")
pols = xds_from_table(ims + "::POLARIZATION")
# subtable data
ddids = dask.compute(ddids)[0]
fields = dask.compute(fields)[0]
spws = dask.compute(spws)[0]
pols = dask.compute(pols)[0]
out_data = []
for ds in xds:
field = fields[ds.FIELD_ID]
radec = field.PHASE_DIR.data.squeeze()
# check fields match
if radec is None:
radec = field.PHASE_DIR.data.squeeze()
if not np.array_equal(radec, field.PHASE_DIR.data.squeeze()):
continue
spw = ds.DATA_DESC_ID # this is not correct, need to use spw
uvw = clone(ds.UVW.data)
bands = band_mapping[ims][spw]
model = model[list(bands), :, :]
vis = im2vis(uvw,
freqs[ims][spw],
model,
freq_bin_idx[ims][spw],
freq_bin_counts[ims][spw],
cell_rad,
nthreads=ngridder_threads,
epsilon=args.epsilon,
do_wstacking=args.wstack)
model_vis = restore_corrs(vis, ncorr)
if mstype == 'zarr':
model_vis = model_vis.rechunk(model_chunks)
uvw = uvw.rechunk((model_chunks[0], 3))
out_ds = ds.assign(
**{
args.model_column: (("row", "chan", "corr"), model_vis),
'UVW': (("row", "three"), uvw)
})
# out_ds = ds.assign(**{args.model_column: (("row", "chan", "corr"), model_vis)})
out_data.append(out_ds)
writes.append(xds_to_table(out_data, ims, columns=[args.model_column]))
dask.visualize(*writes,
filename=args.output_filename + '_predict_graph.pdf',
optimize_graph=False,
collapse_outputs=True)
if not args.mock:
with performance_report(filename=args.output_filename +
'_predict_per.html'):
dask.compute(writes, optimize_graph=False)
print("All done here.", file=log)
def _psf(**kw):
args = OmegaConf.create(kw)
from omegaconf import ListConfig
if not isinstance(args.ms, list) and not isinstance(args.ms, ListConfig):
args.ms = [args.ms]
OmegaConf.set_struct(args, True)
import numpy as np
from pfb.utils.misc import chan_to_band_mapping
import dask
# from dask.distributed import performance_report
from dask.graph_manipulation import clone
from daskms import xds_from_storage_ms as xds_from_ms
from daskms import xds_from_storage_table as xds_from_table
from daskms import Dataset
from daskms.experimental.zarr import xds_to_zarr
import dask.array as da
from africanus.constants import c as lightspeed
from africanus.gridding.wgridder.dask import dirty as vis2im
from ducc0.fft import good_size
from pfb.utils.misc import stitch_images, plan_row_chunk
from pfb.utils.fits import set_wcs, save_fits
# chan <-> band mapping
ms = args.ms
nband = args.nband
freqs, freq_bin_idx, freq_bin_counts, freq_out, band_mapping, chan_chunks = chan_to_band_mapping(
ms, nband=nband)
# gridder memory budget
max_chan_chunk = 0
max_freq = 0
for ims in args.ms:
for spw in freqs[ims]:
counts = freq_bin_counts[ims][spw].compute()
freq = freqs[ims][spw].compute()
max_chan_chunk = np.maximum(max_chan_chunk, counts.max())
max_freq = np.maximum(max_freq, freq.max())
# assumes measurement sets have the same columns,
# number of correlations etc.
xds = xds_from_ms(args.ms[0])
ncorr = xds[0].dims['corr']
nrow = xds[0].dims['row']
# we still have to cater for complex valued data because we cast
# the weights to complex but we not longer need to factor the
# weight column into our memory budget
data_bytes = getattr(xds[0], args.data_column).data.itemsize
bytes_per_row = max_chan_chunk * ncorr * data_bytes
memory_per_row = bytes_per_row
# flags (uint8 or bool)
memory_per_row += bytes_per_row / 8
# UVW
memory_per_row += xds[0].UVW.data.itemsize * 3
# ANTENNA1/2
memory_per_row += xds[0].ANTENNA1.data.itemsize * 2
# TIME
memory_per_row += xds[0].TIME.data.itemsize
# data column is not actually read into memory just used to infer
# dtype and chunking
columns = (args.data_column, args.weight_column, args.flag_column, 'UVW',
'ANTENNA1', 'ANTENNA2', 'TIME')
# flag row
if 'FLAG_ROW' in xds[0]:
columns += ('FLAG_ROW', )
memory_per_row += xds[0].FLAG_ROW.data.itemsize
# imaging weights
if args.imaging_weight_column is not None:
columns += (args.imaging_weight_column, )
memory_per_row += bytes_per_row / 2
# Mueller term (complex valued)
if args.mueller_column is not None:
columns += (args.mueller_column, )
memory_per_row += bytes_per_row
# get max uv coords over all fields
uvw = []
u_max = 0.0
v_max = 0.0
for ims in args.ms:
xds = xds_from_ms(ims, columns=('UVW'), chunks={'row': -1})
for ds in xds:
uvw = ds.UVW.data
u_max = da.maximum(u_max, abs(uvw[:, 0]).max())
v_max = da.maximum(v_max, abs(uvw[:, 1]).max())
uv_max = da.maximum(u_max, v_max)
uv_max = uv_max.compute()
del uvw
# image size
cell_N = 1.0 / (2 * uv_max * max_freq / lightspeed)
if args.cell_size is not None:
cell_size = args.cell_size
cell_rad = cell_size * np.pi / 60 / 60 / 180
if cell_N / cell_rad < 1:
raise ValueError(
"Requested cell size too small. "
"Super resolution factor = ", cell_N / cell_rad)
print("Super resolution factor = %f" % (cell_N / cell_rad), file=log)
else:
cell_rad = cell_N / args.super_resolution_factor
cell_size = cell_rad * 60 * 60 * 180 / np.pi
print("Cell size set to %5.5e arcseconds" % cell_size, file=log)
if args.nx is None:
fov = args.field_of_view * 3600
npix = int(args.psf_oversize * fov / cell_size)
if npix % 2:
npix += 1
nx = npix
ny = npix
else:
nx = args.nx
ny = args.ny if args.ny is not None else nx
print("PSF size set to (%i, %i, %i)" % (nband, nx, ny), file=log)
# get approx image size
# this is not a conservative estimate when multiple SPW's map to a single
# imaging band
pixel_bytes = np.dtype(args.output_type).itemsize
band_size = nx * ny * pixel_bytes
if args.host_address is None:
# full image on single node
row_chunk = plan_row_chunk(args.mem_limit / args.nworkers, band_size,
nrow, memory_per_row,
args.nthreads_per_worker)
else:
# single band per node
row_chunk = plan_row_chunk(args.mem_limit, band_size, nrow,
memory_per_row, args.nthreads_per_worker)
if args.row_chunks is not None:
row_chunk = int(args.row_chunks)
if row_chunk == -1:
row_chunk = nrow
print(
"nrows = %i, row chunks set to %i for a total of %i chunks per node" %
(nrow, row_chunk, int(np.ceil(nrow / row_chunk))),
file=log)
chunks = {}
for ims in args.ms:
chunks[ims] = [] # xds_from_ms expects a list per ds
for spw in freqs[ims]:
chunks[ims].append({
'row': row_chunk,
'chan': chan_chunks[ims][spw]['chan']
})
psfs = []
radec = None # assumes we are only imaging field 0 of first MS
out_datasets = []
for ims in args.ms:
xds = xds_from_ms(ims, chunks=chunks[ims], columns=columns)
# subtables
ddids = xds_from_table(ims + "::DATA_DESCRIPTION")
fields = xds_from_table(ims + "::FIELD")
spws = xds_from_table(ims + "::SPECTRAL_WINDOW")
pols = xds_from_table(ims + "::POLARIZATION")
# subtable data
ddids = dask.compute(ddids)[0]
fields = dask.compute(fields)[0]
spws = dask.compute(spws)[0]
pols = dask.compute(pols)[0]
for ds in xds:
field = fields[ds.FIELD_ID]
# check fields match
if radec is None:
radec = field.PHASE_DIR.data.squeeze()
if not np.array_equal(radec, field.PHASE_DIR.data.squeeze()):
continue
# this is not correct, need to use spw
spw = ds.DATA_DESC_ID
uvw = clone(ds.UVW.data)
data_type = getattr(ds, args.data_column).data.dtype
data_shape = getattr(ds, args.data_column).data.shape
data_chunks = getattr(ds, args.data_column).data.chunks
weights = getattr(ds, args.weight_column).data
if len(weights.shape) < 3:
weights = da.broadcast_to(weights[:, None, :],
data_shape,
chunks=data_chunks)
if args.imaging_weight_column is not None:
imaging_weights = getattr(ds, args.imaging_weight_column).data
if len(imaging_weights.shape) < 3:
imaging_weights = da.broadcast_to(imaging_weights[:,
None, :],
data_shape,
chunks=data_chunks)
weightsxx = imaging_weights[:, :, 0] * weights[:, :, 0]
weightsyy = imaging_weights[:, :, -1] * weights[:, :, -1]
else:
weightsxx = weights[:, :, 0]
weightsyy = weights[:, :, -1]
# apply mueller term
if args.mueller_column is not None:
mueller = getattr(ds, args.mueller_column).data
weightsxx *= da.absolute(mueller[:, :, 0])**2
weightsyy *= da.absolute(mueller[:, :, -1])**2
# weighted sum corr to Stokes I
weights = weightsxx + weightsyy
# MS may contain auto-correlations
if 'FLAG_ROW' in xds[0]:
frow = ds.FLAG_ROW.data | (ds.ANTENNA1.data
== ds.ANTENNA2.data)
else:
frow = (ds.ANTENNA1.data == ds.ANTENNA2.data)
# only keep data where both corrs are unflagged
flag = getattr(ds, args.flag_column).data
flagxx = flag[:, :, 0]
flagyy = flag[:, :, -1]
# ducc0 uses uint8 mask not flag
mask = ~da.logical_or((flagxx | flagyy), frow[:, None])
psf = vis2im(uvw,
freqs[ims][spw],
weights.astype(data_type),
freq_bin_idx[ims][spw],
freq_bin_counts[ims][spw],
nx,
ny,
cell_rad,
flag=mask.astype(np.uint8),
nthreads=args.nvthreads,
epsilon=args.epsilon,
do_wstacking=args.wstack,
double_accum=args.double_accum)
psfs.append(psf)
data_vars = {
'FIELD_ID': (('row', ),
da.full_like(ds.TIME.data,
ds.FIELD_ID,
chunks=args.row_out_chunk)),
'DATA_DESC_ID': (('row', ),
da.full_like(ds.TIME.data,
ds.DATA_DESC_ID,
chunks=args.row_out_chunk)),
'WEIGHT':
(('row', 'chan'), weights.rechunk({0: args.row_out_chunk
})), # why no 'f4'?
'UVW': (('row', 'uvw'), uvw.rechunk({0: args.row_out_chunk}))
}
coords = {'chan': (('chan', ), freqs[ims][spw])}
out_ds = Dataset(data_vars, coords)
out_datasets.append(out_ds)
writes = xds_to_zarr(out_datasets,
args.output_filename + '.zarr',
columns='ALL')
# dask.visualize(writes, filename=args.output_filename + '_psf_writes_graph.pdf', optimize_graph=False)
# dask.visualize(psfs, filename=args.output_filename + '_psf_graph.pdf', optimize_graph=False)
if not args.mock:
# psfs = dask.compute(psfs, writes, optimize_graph=False)[0]
# with performance_report(filename=args.output_filename + '_psf_per.html'):
psfs = dask.compute(psfs, writes, optimize_graph=False)[0]
psf = stitch_images(psfs, nband, band_mapping)
hdr = set_wcs(cell_size / 3600, cell_size / 3600, nx, ny, radec,
freq_out)
save_fits(args.output_filename + '_psf.fits',
psf,
hdr,
dtype=args.output_type)
psf_mfs = np.sum(psf, axis=0)
wsum = psf_mfs.max()
psf_mfs /= wsum
hdr_mfs = set_wcs(cell_size / 3600, cell_size / 3600, nx, ny, radec,
np.mean(freq_out))
save_fits(args.output_filename + '_psf_mfs.fits',
psf_mfs,
hdr_mfs,
dtype=args.output_type)
print("All done here.", file=log)
def test_bind_clone_collections(func):
@delayed
def double(x):
return x * 2
# dask.delayed
d1 = double(2)
d2 = double(d1)
# dask.array
a1 = da.ones((10, 10), chunks=5)
a2 = a1 + 1
a3 = a2.T
# dask.bag
b1 = db.from_sequence([1, 2], npartitions=2)
# b1's tasks are not callable, so we need an extra step to properly test bind
b2 = b1.map(lambda x: x * 2)
b3 = b2.map(lambda x: x + 1)
b4 = b3.min()
# dask.dataframe
df = pd.DataFrame({"x": list(range(10))})
ddf1 = dd.from_pandas(df, npartitions=2)
# ddf1's tasks are not callable, so we need an extra step to properly test bind
ddf2 = ddf1.map_partitions(lambda x: x * 2)
ddf3 = ddf2.map_partitions(lambda x: x + 1)
ddf4 = ddf3["x"] # dd.Series
ddf5 = ddf4.min() # dd.Scalar
cnt = NodeCounter()
if func is bind:
parent = da.ones((10, 10), chunks=5).map_blocks(cnt.f)
cnt.n = 0
d2c, a3c, b3c, b4c, ddf3c, ddf4c, ddf5c = bind(
children=(d2, a3, b3, b4, ddf3, ddf4, ddf5),
parents=parent,
omit=(d1, a1, b2, ddf2),
seed=0,
)
else:
d2c, a3c, b3c, b4c, ddf3c, ddf4c, ddf5c = clone(
d2,
a3,
b3,
b4,
ddf3,
ddf4,
ddf5,
omit=(d1, a1, b2, ddf2),
seed=0,
)
assert_did_not_materialize(d2c, d2)
assert_did_not_materialize(a3c, a3)
assert_did_not_materialize(b3c, b3)
assert_did_not_materialize(b4c, b4)
assert_did_not_materialize(ddf3c, ddf3)
assert_did_not_materialize(ddf4c, ddf4)
assert_did_not_materialize(ddf5c, ddf5)
assert_no_common_keys(d2c, d2, omit=d1, layers=True)
assert_no_common_keys(a3c, a3, omit=a1, layers=True)
assert_no_common_keys(b3c, b3, omit=b2, layers=True)
assert_no_common_keys(ddf3c, ddf3, omit=ddf2, layers=True)
assert_no_common_keys(ddf4c, ddf4, omit=ddf2, layers=True)
assert_no_common_keys(ddf5c, ddf5, omit=ddf2, layers=True)
assert d2.compute() == d2c.compute()
assert cnt.n == 4 or func is clone
da.utils.assert_eq(a3c, a3)
assert cnt.n == 8 or func is clone
db.utils.assert_eq(b3c, b3)
assert cnt.n == 12 or func is clone
db.utils.assert_eq(b4c, b4)
assert cnt.n == 16 or func is clone
dd.utils.assert_eq(ddf3c, ddf3)
assert cnt.n == 24 or func is clone # dd.utils.assert_eq calls compute() twice
dd.utils.assert_eq(ddf4c, ddf4)
assert cnt.n == 32 or func is clone # dd.utils.assert_eq calls compute() twice
dd.utils.assert_eq(ddf5c, ddf5)
assert cnt.n == 36 or func is clone
def _residual(ms, stack, **kw):
args = OmegaConf.create(kw)
OmegaConf.set_struct(args, True)
pyscilog.log_to_file(args.output_filename + '.log')
pyscilog.enable_memory_logging(level=3)
# number of threads per worker
if args.nthreads is None:
if args.host_address is not None:
raise ValueError(
"You have to specify nthreads when using a distributed scheduler"
)
import multiprocessing
nthreads = multiprocessing.cpu_count()
args.nthreads = nthreads
else:
nthreads = args.nthreads
# configure memory limit
if args.mem_limit is None:
if args.host_address is not None:
raise ValueError(
"You have to specify mem-limit when using a distributed scheduler"
)
import psutil
mem_limit = int(psutil.virtual_memory()[0] /
1e9) # 100% of memory by default
args.mem_limit = mem_limit
else:
mem_limit = args.mem_limit
nband = args.nband
if args.nworkers is None:
nworkers = nband
args.nworkers = nworkers
else:
nworkers = args.nworkers
if args.nthreads_per_worker is None:
nthreads_per_worker = 1
args.nthreads_per_worker = nthreads_per_worker
else:
nthreads_per_worker = args.nthreads_per_worker
# the number of chunks being read in simultaneously is equal to
# the number of dask threads
nthreads_dask = nworkers * nthreads_per_worker
if args.ngridder_threads is None:
if args.host_address is not None:
ngridder_threads = nthreads // nthreads_per_worker
else:
ngridder_threads = nthreads // nthreads_dask
args.ngridder_threads = ngridder_threads
else:
ngridder_threads = args.ngridder_threads
ms = list(ms)
print('Input Options:', file=log)
for key in kw.keys():
print(' %25s = %s' % (key, args[key]), file=log)
# numpy imports have to happen after this step
from pfb import set_client
set_client(nthreads, mem_limit, nworkers, nthreads_per_worker,
args.host_address, stack, log)
import numpy as np
from pfb.utils.misc import chan_to_band_mapping
import dask
from dask.graph_manipulation import clone
from dask.distributed import performance_report
from daskms import xds_from_storage_ms as xds_from_ms
from daskms import xds_from_storage_table as xds_from_table
import dask.array as da
from africanus.constants import c as lightspeed
from africanus.gridding.wgridder.dask import residual as im2residim
from ducc0.fft import good_size
from pfb.utils.misc import stitch_images, plan_row_chunk
from pfb.utils.fits import set_wcs, save_fits
# chan <-> band mapping
freqs, freq_bin_idx, freq_bin_counts, freq_out, band_mapping, chan_chunks = chan_to_band_mapping(
ms, nband=nband)
# gridder memory budget
max_chan_chunk = 0
max_freq = 0
for ims in ms:
for spw in freqs[ims]:
counts = freq_bin_counts[ims][spw].compute()
freq = freqs[ims][spw].compute()
max_chan_chunk = np.maximum(max_chan_chunk, counts.max())
max_freq = np.maximum(max_freq, freq.max())
# assumes measurement sets have the same columns,
# number of correlations etc.
xds = xds_from_ms(ms[0])
ncorr = xds[0].dims['corr']
nrow = xds[0].dims['row']
data_bytes = getattr(xds[0], args.data_column).data.itemsize
bytes_per_row = max_chan_chunk * ncorr * data_bytes
memory_per_row = bytes_per_row
# real valued weights
wdims = getattr(xds[0], args.weight_column).data.ndim
if wdims == 2: # WEIGHT
memory_per_row += ncorr * data_bytes / 2
else: # WEIGHT_SPECTRUM
memory_per_row += bytes_per_row / 2
# flags (uint8 or bool)
memory_per_row += np.dtype(np.uint8).itemsize * max_chan_chunk * ncorr
# UVW
memory_per_row += xds[0].UVW.data.itemsize * 3
# ANTENNA1/2
memory_per_row += xds[0].ANTENNA1.data.itemsize * 2
columns = (args.data_column, args.weight_column, args.flag_column, 'UVW',
'ANTENNA1', 'ANTENNA2')
# flag row
if 'FLAG_ROW' in xds[0]:
columns += ('FLAG_ROW', )
memory_per_row += xds[0].FLAG_ROW.data.itemsize
# imaging weights
if args.imaging_weight_column is not None:
columns += (args.imaging_weight_column, )
memory_per_row += bytes_per_row / 2
# Mueller term (complex valued)
if args.mueller_column is not None:
columns += (args.mueller_column, )
memory_per_row += bytes_per_row
# get max uv coords over all fields
uvw = []
u_max = 0.0
v_max = 0.0
for ims in ms:
xds = xds_from_ms(ims, columns=('UVW'), chunks={'row': -1})
for ds in xds:
uvw = ds.UVW.data
u_max = da.maximum(u_max, abs(uvw[:, 0]).max())
v_max = da.maximum(v_max, abs(uvw[:, 1]).max())
uv_max = da.maximum(u_max, v_max)
uv_max = uv_max.compute()
del uvw
# image size
cell_N = 1.0 / (2 * uv_max * max_freq / lightspeed)
if args.cell_size is not None:
cell_size = args.cell_size
cell_rad = cell_size * np.pi / 60 / 60 / 180
if cell_N / cell_rad < 1:
raise ValueError(
"Requested cell size too small. "
"Super resolution factor = ", cell_N / cell_rad)
print("Super resolution factor = %f" % (cell_N / cell_rad), file=log)
else:
cell_rad = cell_N / args.super_resolution_factor
cell_size = cell_rad * 60 * 60 * 180 / np.pi
print("Cell size set to %5.5e arcseconds" % cell_size, file=log)
if args.nx is None:
fov = args.field_of_view * 3600
npix = int(fov / cell_size)
if npix % 2:
npix += 1
nx = good_size(npix)
ny = good_size(npix)
else:
nx = args.nx
ny = args.ny if args.ny is not None else nx
print("Image size set to (%i, %i, %i)" % (nband, nx, ny), file=log)
# get approx image size
# this is not a conservative estimate when multiple SPW's map to a single
# imaging band
pixel_bytes = np.dtype(args.output_type).itemsize
band_size = nx * ny * pixel_bytes
if args.host_address is None:
# full image on single node
row_chunk = plan_row_chunk(mem_limit / nworkers, band_size, nrow,
memory_per_row, nthreads_per_worker)
else:
# single band per node
row_chunk = plan_row_chunk(mem_limit, band_size, nrow, memory_per_row,
nthreads_per_worker)
if args.row_chunks is not None:
row_chunk = int(args.row_chunks)
if row_chunk == -1:
row_chunk = nrow
print(
"nrows = %i, row chunks set to %i for a total of %i chunks per node" %
(nrow, row_chunk, int(np.ceil(nrow / row_chunk))),
file=log)
chunks = {}
for ims in ms:
chunks[ims] = [] # xds_from_ms expects a list per ds
for spw in freqs[ims]:
chunks[ims].append({
'row': row_chunk,
'chan': chan_chunks[ims][spw]['chan']
})
dirties = []
radec = None # assumes we are only imaging field 0 of first MS
for ims in ms:
xds = xds_from_ms(ims, chunks=chunks[ims], columns=columns)
# subtables
ddids = xds_from_table(ims + "::DATA_DESCRIPTION")
fields = xds_from_table(ims + "::FIELD")
spws = xds_from_table(ims + "::SPECTRAL_WINDOW")
pols = xds_from_table(ims + "::POLARIZATION")
# subtable data
ddids = dask.compute(ddids)[0]
fields = dask.compute(fields)[0]
spws = dask.compute(spws)[0]
pols = dask.compute(pols)[0]
for ds in xds:
field = fields[ds.FIELD_ID]
# check fields match
if radec is None:
radec = field.PHASE_DIR.data.squeeze()
if not np.array_equal(radec, field.PHASE_DIR.data.squeeze()):
continue
# this is not correct, need to use spw
spw = ds.DATA_DESC_ID
uvw = clone(ds.UVW.data)
data = getattr(ds, args.data_column).data
dataxx = data[:, :, 0]
datayy = data[:, :, -1]
weights = getattr(ds, args.weight_column).data
if len(weights.shape) < 3:
weights = da.broadcast_to(weights[:, None, :],
data.shape,
chunks=data.chunks)
if args.imaging_weight_column is not None:
imaging_weights = getattr(ds, args.imaging_weight_column).data
if len(imaging_weights.shape) < 3:
imaging_weights = da.broadcast_to(imaging_weights[:,
None, :],
data.shape,
chunks=data.chunks)
weightsxx = imaging_weights[:, :, 0] * weights[:, :, 0]
weightsyy = imaging_weights[:, :, -1] * weights[:, :, -1]
else:
weightsxx = weights[:, :, 0]
weightsyy = weights[:, :, -1]
# apply adjoint of mueller term.
# Phases modify data amplitudes modify weights.
if args.mueller_column is not None:
mueller = getattr(ds, args.mueller_column).data
dataxx *= da.exp(-1j * da.angle(mueller[:, :, 0]))
datayy *= da.exp(-1j * da.angle(mueller[:, :, -1]))
weightsxx *= da.absolute(mueller[:, :, 0])
weightsyy *= da.absolute(mueller[:, :, -1])
# weighted sum corr to Stokes I
weights = weightsxx + weightsyy
data = (weightsxx * dataxx + weightsyy * datayy)
# TODO - turn off this stupid warning
data = da.where(weights, data / weights, 0.0j)
# MS may contain auto-correlations
if 'FLAG_ROW' in xds[0]:
frow = ds.FLAG_ROW.data | (ds.ANTENNA1.data
== ds.ANTENNA2.data)
else:
frow = (ds.ANTENNA1.data == ds.ANTENNA2.data)
# only keep data where both corrs are unflagged
flag = getattr(ds, args.flag_column).data
flagxx = flag[:, :, 0]
flagyy = flag[:, :, -1]
# ducc0 uses uint8 mask not flag
mask = ~da.logical_or((flagxx | flagyy), frow[:, None])
dirty = vis2im(uvw,
freqs[ims][spw],
data,
freq_bin_idx[ims][spw],
freq_bin_counts[ims][spw],
nx,
ny,
cell_rad,
weights=weights,
flag=mask.astype(np.uint8),
nthreads=ngridder_threads,
epsilon=args.epsilon,
do_wstacking=args.wstack,
double_accum=args.double_accum)
dirties.append(dirty)
# dask.visualize(dirties, filename=args.output_filename + '_graph.pdf', optimize_graph=False)
if not args.mock:
# result = dask.compute(dirties, wsum, optimize_graph=False)
with performance_report(filename=args.output_filename + '_per.html'):
result = dask.compute(dirties, optimize_graph=False)
dirties = result[0]
dirty = stitch_images(dirties, nband, band_mapping)
hdr = set_wcs(cell_size / 3600, cell_size / 3600, nx, ny, radec,
freq_out)
save_fits(args.output_filename + '_dirty.fits',
dirty,
hdr,
dtype=args.output_type)
print("All done here.", file=log)
