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 _restore(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 astropy.io import fits
mhdr = fits.getheader(args.model)
from pfb.utils.fits import load_fits
model = load_fits(args.model).squeeze() # drop Stokes axis
# check images compatible
rhdr = fits.getheader(args.residual)
from pfb.utils.fits import compare_headers
compare_headers(mhdr, rhdr)
residual = load_fits(args.residual).squeeze()
# fit restoring psf
from pfb.utils.misc import fitcleanbeam
psf = load_fits(args.psf, dtype=args.real_type).squeeze()
nband, nx_psf, ny_psf = psf.shape
wsums = np.amax(psf.reshape(args.nband, nx_psf, ny_psf), axis=1)
wsum = np.sum(wsums)
psf /= wsum
psf_mfs = np.sum(psf, axis=0)
# fit restoring psf
GaussPar = fitcleanbeam(psf_mfs[None], level=0.5, pixsize=1.0)
GaussPars = fitcleanbeam(psf, level=0.5, pixsize=1.0)
cpsf_mfs = np.zeros(psf_mfs.shape, dtype=args.real_type)
cpsf = np.zeros(psf.shape, dtype=args.real_type)
lpsf = np.arange(-R.nx_psf / 2, R.nx_psf / 2)
mpsf = np.arange(-R.ny_psf / 2, R.ny_psf / 2)
xx, yy = np.meshgrid(lpsf, mpsf, indexing='ij')
cpsf_mfs = Gaussian2D(xx, yy, GaussPar[0], normalise=False)
for v in range(args.nband):
cpsf[v] = Gaussian2D(xx, yy, GaussPars[v], normalise=False)
from pfb.utils.fits import add_beampars
GaussPar = list(GaussPar[0])
GaussPar[0] *= args.cell_size / 3600
GaussPar[1] *= args.cell_size / 3600
GaussPar = tuple(GaussPar)
hdr_psf_mfs = add_beampars(hdr_psf_mfs, GaussPar)
save_fits(args.outfile + '_cpsf_mfs.fits', cpsf_mfs, hdr_psf_mfs)
save_fits(args.outfile + '_psf_mfs.fits', psf_mfs, hdr_psf_mfs)
if args.beam is not None:
bhdr = fits.getheader(args.beam)
compare_headers(mhdr, bhdr)
beam = load_fits(args.beam).squeeze()
model = np.where(beam > args.pb_min, model / beam, 0.0)
nband, nx, ny = model.shape
guassparf = ()
if nband > 1:
for b in range(nband):
guassparf += (rhdr['BMAJ' + str(b)], rhdr['BMIN' + str(b)],
rhdr['BPA' + str(b)])
else:
guassparf += (rhdr['BMAJ'], rhdr['BMIN'], rhdr['BPA'])
# if args.convolve_residuals:
cellx = np.abs(mhdr['CDELT1'])
celly = np.abs(mhdr['CDELT2'])
from pfb.utils.restoration import restore_image
def _main(dest=sys.stdout):
from pfb.parser import create_parser
args = create_parser().parse_args()
if not args.nthreads:
import multiprocessing
args.nthreads = multiprocessing.cpu_count()
if not args.mem_limit:
import psutil
args.mem_limit = int(psutil.virtual_memory()[0] /
1e9) # 100% of memory by default
import numpy as np
import numba
import numexpr
import dask
import dask.array as da
from daskms import xds_from_ms, xds_from_table
from astropy.io import fits
from pfb.utils.fits import (set_wcs, load_fits, save_fits, compare_headers,
data_from_header)
from pfb.utils.restoration import fitcleanbeam
from pfb.utils.misc import Gaussian2D
from pfb.operators.gridder import Gridder
from pfb.operators.psf import PSF
from pfb.deconv.sara import sara
from pfb.deconv.clean import clean
from pfb.deconv.spotless import spotless
from pfb.deconv.nnls import nnls
from pfb.opt.pcg import pcg
if not isinstance(args.ms, list):
args.ms = [args.ms]
pyscilog.log_to_file(args.outfile + '.log')
pyscilog.enable_memory_logging(level=3)
GD = vars(args)
print('Input Options:', file=log)
for key in GD.keys():
print(' %25s = %s' % (key, GD[key]), file=log)
# get max uv coords over all fields
uvw = []
u_max = 0.0
v_max = 0.0
all_freqs = []
for ims in args.ms:
xds = xds_from_ms(ims,
group_cols=('FIELD_ID', 'DATA_DESC_ID'),
columns=('UVW'),
chunks={'row': args.row_chunks})
spws = xds_from_table(ims + "::SPECTRAL_WINDOW", group_cols="__row__")
spws = dask.compute(spws)[0]
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)
spw = spws[ds.DATA_DESC_ID]
tmp_freq = spw.CHAN_FREQ.data.squeeze()
all_freqs.append(list([tmp_freq]))
uv_max = u_max.compute()
del uvw
# get Nyquist cell size
from africanus.constants import c as lightspeed
all_freqs = dask.compute(all_freqs)
freq = np.unique(all_freqs)
cell_N = 1.0 / (2 * uv_max * freq.max() / lightspeed)
if args.cell_size is not None:
cell_rad = args.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=dest)
else:
cell_rad = cell_N / args.super_resolution_factor
args.cell_size = cell_rad * 60 * 60 * 180 / np.pi
print("Cell size set to %5.5e arcseconds" % args.cell_size, file=dest)
if args.nx is None or args.ny is None:
from ducc0.fft import good_size
fov = args.fov * 3600
npix = int(fov / args.cell_size)
if npix % 2:
npix += 1
args.nx = good_size(npix)
args.ny = good_size(npix)
if args.nband is None:
args.nband = freq.size
print("Image size set to (%i, %i, %i)" % (args.nband, args.nx, args.ny),
file=dest)
# mask
if args.mask is not None:
mask_array = load_fits(args.mask, dtype=args.real_type).squeeze()
if mask_array.shape != (args.nx, args.ny):
raise ValueError("Mask has incorrect shape.")
# add freq axis
mask_array = mask_array[None]
def mask(x):
return mask_array * x
else:
mask_array = None
def mask(x):
return x
# init gridder
R = Gridder(
args.ms,
args.nx,
args.ny,
args.cell_size,
nband=args.nband,
nthreads=args.nthreads,
do_wstacking=args.do_wstacking,
row_chunks=args.row_chunks,
psf_oversize=args.psf_oversize,
data_column=args.data_column,
epsilon=args.epsilon,
weight_column=args.weight_column,
imaging_weight_column=args.imaging_weight_column,
model_column=args.model_column,
flag_column=args.flag_column,
weighting=args.weighting,
robust=args.robust,
mem_limit=int(
0.8 * args.mem_limit)) # assumes gridding accounts for 80% memory
freq_out = R.freq_out
radec = R.radec
print("PSF size set to (%i, %i, %i)" % (args.nband, R.nx_psf, R.ny_psf),
file=dest)
# get headers
hdr = set_wcs(args.cell_size / 3600, args.cell_size / 3600, args.nx,
args.ny, radec, freq_out)
hdr_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600, args.nx,
args.ny, radec, np.mean(freq_out))
hdr_psf = set_wcs(args.cell_size / 3600, args.cell_size / 3600, R.nx_psf,
R.ny_psf, radec, freq_out)
hdr_psf_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600,
R.nx_psf, R.ny_psf, radec, np.mean(freq_out))
# psf
if args.psf is not None:
try:
compare_headers(hdr_psf, fits.getheader(args.psf))
psf = load_fits(args.psf, dtype=args.real_type).squeeze()
except BaseException:
raise
psf = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf, hdr_psf)
else:
psf = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf, hdr_psf)
# Normalising by wsum (so that the PSF always sums to 1) results in the
# most intuitive sig_21 values and by far the least bookkeeping.
# However, we won't save the cubes that way as it destroys information
# about the noise in image space. Note only the MFS images will have the
# usual units of Jy/beam.
wsums = np.amax(psf.reshape(args.nband, R.nx_psf * R.ny_psf), axis=1)
wsum = np.sum(wsums)
psf /= wsum
psf_mfs = np.sum(psf, axis=0)
# fit restoring psf
GaussPar = fitcleanbeam(psf_mfs[None], level=0.5, pixsize=1.0)
GaussPars = fitcleanbeam(psf, level=0.5, pixsize=1.0)
cpsf_mfs = np.zeros(psf_mfs.shape, dtype=args.real_type)
cpsf = np.zeros(psf.shape, dtype=args.real_type)
lpsf = np.arange(-R.nx_psf / 2, R.nx_psf / 2)
mpsf = np.arange(-R.ny_psf / 2, R.ny_psf / 2)
xx, yy = np.meshgrid(lpsf, mpsf, indexing='ij')
cpsf_mfs = Gaussian2D(xx, yy, GaussPar[0], normalise=False)
for v in range(args.nband):
cpsf[v] = Gaussian2D(xx, yy, GaussPars[v], normalise=False)
from pfb.utils.fits import add_beampars
GaussPar = list(GaussPar[0])
GaussPar[0] *= args.cell_size / 3600
GaussPar[1] *= args.cell_size / 3600
GaussPar = tuple(GaussPar)
hdr_psf_mfs = add_beampars(hdr_psf_mfs, GaussPar)
save_fits(args.outfile + '_cpsf_mfs.fits', cpsf_mfs, hdr_psf_mfs)
save_fits(args.outfile + '_psf_mfs.fits', psf_mfs, hdr_psf_mfs)
GaussPars = list(GaussPars)
for b in range(args.nband):
GaussPars[b] = list(GaussPars[b])
GaussPars[b][0] *= args.cell_size / 3600
GaussPars[b][1] *= args.cell_size / 3600
GaussPars[b] = tuple(GaussPars[b])
GaussPars = tuple(GaussPars)
hdr_psf = add_beampars(hdr_psf, GaussPar, GaussPars)
save_fits(args.outfile + '_cpsf.fits', cpsf, hdr_psf)
# dirty
if args.dirty is not None:
try:
compare_headers(hdr, fits.getheader(args.dirty))
dirty = load_fits(args.dirty).squeeze()
except BaseException:
raise
dirty = R.make_dirty()
save_fits(args.outfile + '_dirty.fits', dirty, hdr)
else:
dirty = R.make_dirty()
save_fits(args.outfile + '_dirty.fits', dirty, hdr)
dirty /= wsum
dirty_mfs = np.sum(dirty, axis=0)
save_fits(args.outfile + '_dirty_mfs.fits', dirty_mfs, hdr_mfs)
quit()
# initial model and residual
if args.x0 is not None:
try:
compare_headers(hdr, fits.getheader(args.x0))
model = load_fits(args.x0, dtype=args.real_type).squeeze()
if args.first_residual is not None:
try:
compare_headers(hdr, fits.getheader(args.first_residual))
residual = load_fits(args.first_residual,
dtype=args.real_type).squeeze()
except BaseException:
residual = R.make_residual(model)
save_fits(args.outfile + '_first_residual.fits', residual,
hdr)
else:
residual = R.make_residual(model)
save_fits(args.outfile + '_first_residual.fits', residual, hdr)
residual /= wsum
except BaseException:
model = np.zeros((args.nband, args.nx, args.ny))
residual = dirty.copy()
else:
model = np.zeros((args.nband, args.nx, args.ny))
residual = dirty.copy()
residual_mfs = np.sum(residual, axis=0)
save_fits(args.outfile + '_first_residual_mfs.fits', residual_mfs, hdr_mfs)
# smooth beam
if args.beam_model is not None:
if args.beam_model[-5:] == '.fits':
beam_image = load_fits(args.beam_model,
dtype=args.real_type).squeeze()
if beam_image.shape != (args.nband, args.nx, args.ny):
raise ValueError("Beam has incorrect shape")
elif args.beam_model == "JimBeam":
from katbeam import JimBeam
if args.band.lower() == 'l':
beam = JimBeam('MKAT-AA-L-JIM-2020')
else:
beam = JimBeam('MKAT-AA-UHF-JIM-2020')
beam_image = np.zeros((args.nband, args.nx, args.ny),
dtype=args.real_type)
l_coord, ref_l = data_from_header(hdr, axis=1)
l_coord -= ref_l
m_coord, ref_m = data_from_header(hdr, axis=2)
m_coord -= ref_m
xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij')
for v in range(args.nband):
beam_image[v] = beam.I(xx, yy, freq_out[v])
def beam(x):
return beam_image * x
else:
beam_image = None
def beam(x):
return x
if args.init_nnls:
print("Initialising with NNLS", file=log)
model = nnls(psf,
model,
residual,
mask=mask_array,
beam_image=beam_image,
hdr=hdr,
hdr_mfs=hdr_mfs,
outfile=args.outfile,
maxit=1,
nthreads=args.nthreads)
residual = R.make_residual(beam(mask(model))) / wsum
residual_mfs = np.sum(residual, axis=0)
# deconvolve
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
redo_dirty = False
print("Peak of initial residual is %f and rms is %f" % (rmax, rms),
file=dest)
for i in range(0, args.maxit):
# run minor cycle of choice
modelp = model.copy()
if args.deconv_mode == 'sara':
model = sara(psf,
model,
residual,
mask=mask_array,
beam_image=beam_image,
hessian=R.convolve,
wsum=wsum,
adapt_sig21=args.adapt_sig21,
hdr=hdr,
hdr_mfs=hdr_mfs,
outfile=args.outfile,
cpsf=cpsf,
nthreads=args.nthreads,
sig_21=args.sig_21,
sigma_frac=args.sigma_frac,
maxit=args.minormaxit,
tol=args.minortol,
gamma=args.gamma,
psi_levels=args.psi_levels,
psi_basis=args.psi_basis,
pdtol=args.pdtol,
pdmaxit=args.pdmaxit,
pdverbose=args.pdverbose,
positivity=args.positivity,
cgtol=args.cgtol,
cgminit=args.cgminit,
cgmaxit=args.cgmaxit,
cgverbose=args.cgverbose,
pmtol=args.pmtol,
pmmaxit=args.pmmaxit,
pmverbose=args.pmverbose)
elif args.deconv_mode == 'clean':
model = clean(psf,
model,
residual,
mask=mask_array,
beam=beam_image,
nthreads=args.nthreads,
maxit=args.minormaxit,
gamma=args.gamma,
peak_factor=args.peak_factor,
threshold=args.threshold,
hbgamma=args.hbgamma,
hbpf=args.hbpf,
hbmaxit=args.hbmaxit,
hbverbose=args.hbverbose)
elif args.deconv_mode == 'spotless':
model = spotless(psf,
model,
residual,
mask=mask_array,
beam_image=beam_image,
hessian=R.convolve,
wsum=wsum,
adapt_sig21=args.adapt_sig21,
cpsf=cpsf_mfs,
hdr=hdr,
hdr_mfs=hdr_mfs,
outfile=args.outfile,
sig_21=args.sig_21,
sigma_frac=args.sigma_frac,
nthreads=args.nthreads,
gamma=args.gamma,
peak_factor=args.peak_factor,
maxit=args.minormaxit,
tol=args.minortol,
threshold=args.threshold,
positivity=args.positivity,
hbgamma=args.hbgamma,
hbpf=args.hbpf,
hbmaxit=args.hbmaxit,
hbverbose=args.hbverbose,
pdtol=args.pdtol,
pdmaxit=args.pdmaxit,
pdverbose=args.pdverbose,
cgtol=args.cgtol,
cgminit=args.cgminit,
cgmaxit=args.cgmaxit,
cgverbose=args.cgverbose,
pmtol=args.pmtol,
pmmaxit=args.pmmaxit,
pmverbose=args.pmverbose)
else:
raise ValueError("Unknown deconvolution mode ", args.deconv_mode)
# get residual
if redo_dirty:
# Need to do this if weights or Jones has changed
# (eg. if we change robustness factor, reweight or calibrate)
psf = R.make_psf()
wsums = np.amax(psf.reshape(args.nband, R.nx_psf * R.ny_psf),
axis=1)
wsum = np.sum(wsums)
psf /= wsum
dirty = R.make_dirty() / wsum
# compute in image space
# residual = dirty - R.convolve(beam(mask(model))) / wsum
residual = R.make_residual(beam(mask(model))) / wsum
residual_mfs = np.sum(residual, axis=0)
# save current iteration
model_mfs = np.mean(model, axis=0)
save_fits(args.outfile + '_major' + str(i + 1) + '_model_mfs.fits',
model_mfs, hdr_mfs)
save_fits(args.outfile + '_major' + str(i + 1) + '_model.fits', model,
hdr)
save_fits(args.outfile + '_major' + str(i + 1) + '_residual_mfs.fits',
residual_mfs, hdr_mfs)
save_fits(args.outfile + '_major' + str(i + 1) + '_residual.fits',
residual * wsum, hdr)
# check stopping criteria
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
eps = np.linalg.norm(model - modelp) / np.linalg.norm(model)
print("At iteration %i peak of residual is %f, rms is %f, current "
"eps is %f" % (i + 1, rmax, rms, eps),
file=dest)
if eps < args.tol:
break
if args.mop_flux:
print("Mopping flux", file=dest)
# vague Gaussian prior on x
def hess(x):
return mask(beam(R.convolve(mask(beam(x))))) / wsum + 1e-6 * x
def M(x):
return x / 1e-6 # preconditioner
x = pcg(hess,
mask(beam(residual)),
np.zeros(residual.shape, dtype=residual.dtype),
M=M,
tol=0.1 * args.cgtol,
maxit=args.cgmaxit,
minit=args.cgminit,
verbosity=args.cgverbose)
model += x
# residual = dirty - R.convolve(beam(mask(model))) / wsum
residual = R.make_residual(beam(mask(model))) / wsum
save_fits(args.outfile + '_mopped_model.fits', model, hdr)
save_fits(args.outfile + '_mopped_residual.fits', residual, hdr)
model_mfs = np.mean(model, axis=0)
save_fits(args.outfile + '_mopped_model_mfs.fits', model_mfs, hdr_mfs)
residual_mfs = np.sum(residual, axis=0)
save_fits(args.outfile + '_mopped_residual_mfs.fits', residual_mfs,
hdr_mfs)
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
print("After mopping flux peak of residual is %f, rms is %f" %
(rmax, rms),
file=dest)
# if args.interp_model:
# nband = args.nband
# order = args.spectral_poly_order
# phi.trim_fat(model)
# I = np.argwhere(phi.mask).squeeze()
# Ix = I[:, 0]
# Iy = I[:, 1]
# npix = I.shape[0]
# # get components
# beta = model[:, Ix, Iy]
# # fit integrated polynomial to model components
# # we are given frequencies at bin centers, convert to bin edges
# ref_freq = np.mean(freq_out)
# delta_freq = freq_out[1] - freq_out[0]
# wlow = (freq_out - delta_freq/2.0)/ref_freq
# whigh = (freq_out + delta_freq/2.0)/ref_freq
# wdiff = whigh - wlow
# # set design matrix for each component
# Xdesign = np.zeros([freq_out.size, args.spectral_poly_order])
# for i in range(1, args.spectral_poly_order+1):
# Xdesign[:, i-1] = (whigh**i - wlow**i)/(i*wdiff)
# weights = psf_max[:, None]
# dirty_comps = Xdesign.T.dot(weights*beta)
# hess_comps = Xdesign.T.dot(weights*Xdesign)
# comps = np.linalg.solve(hess_comps, dirty_comps)
# np.savez(args.outfile + "spectral_comps", comps=comps, ref_freq=ref_freq, mask=np.any(model, axis=0))
if args.write_model:
print("Writing model", file=dest)
R.write_model(model)
if args.make_restored:
print("Making restored", file=dest)
cpsfo = PSF(cpsf, residual.shape, nthreads=args.nthreads)
restored = cpsfo.convolve(model)
# residual needs to be in Jy/beam before adding to convolved model
wsums = np.amax(psf.reshape(-1, R.nx_psf * R.ny_psf), axis=1)
restored += residual / wsums[:, None, None]
save_fits(args.outfile + '_restored.fits', restored, hdr)
restored_mfs = np.mean(restored, axis=0)
save_fits(args.outfile + '_restored_mfs.fits', restored_mfs, hdr_mfs)
residual_mfs = np.sum(residual, axis=0)
def _forward(**kw):
args = OmegaConf.create(kw)
OmegaConf.set_struct(args, True)
import numpy as np
import numexpr as ne
import dask
import dask.array as da
from dask.distributed import performance_report
from pfb.utils.fits import load_fits, set_wcs, save_fits, data_from_header
from pfb.opt.hogbom import hogbom
from astropy.io import fits
print("Loading residual", file=log)
residual = load_fits(args.residual, dtype=args.output_type).squeeze()
nband, nx, ny = residual.shape
hdr = fits.getheader(args.residual)
print("Loading psf", file=log)
psf = load_fits(args.psf, dtype=args.output_type).squeeze()
_, nx_psf, ny_psf = psf.shape
hdr_psf = fits.getheader(args.psf)
wsums = np.amax(psf.reshape(-1, nx_psf*ny_psf), axis=1)
wsum = np.sum(wsums)
psf /= wsum
psf_mfs = np.sum(psf, axis=0)
assert (psf_mfs.max() - 1.0) < 1e-4
residual /= wsum
residual_mfs = np.sum(residual, axis=0)
# get info required to set WCS
ra = np.deg2rad(hdr['CRVAL1'])
dec = np.deg2rad(hdr['CRVAL2'])
radec = [ra, dec]
cell_deg = np.abs(hdr['CDELT1'])
if cell_deg != np.abs(hdr['CDELT2']):
raise NotImplementedError('cell sizes have to be equal')
cell_rad = np.deg2rad(cell_deg)
l_coord, ref_l = data_from_header(hdr, axis=1)
l_coord -= ref_l
m_coord, ref_m = data_from_header(hdr, axis=2)
m_coord -= ref_m
freq_out, ref_freq = data_from_header(hdr, axis=3)
hdr_mfs = set_wcs(cell_deg, cell_deg, nx, ny, radec, ref_freq)
save_fits(args.output_filename + '_residual_mfs.fits', residual_mfs, hdr_mfs,
dtype=args.output_type)
rms = np.std(residual_mfs)
rmax = np.abs(residual_mfs).max()
print("Initial peak residual = %f, rms = %f" % (rmax, rms), file=log)
# load beam
if args.beam_model is not None:
if args.beam_model.endswith('.fits'): # beam already interpolated
bhdr = fits.getheader(args.beam_model)
l_coord_beam, ref_lb = data_from_header(bhdr, axis=1)
l_coord_beam -= ref_lb
if not np.array_equal(l_coord_beam, l_coord):
raise ValueError("l coordinates of beam model do not match those of image. Use power_beam_maker to interpolate to fits header.")
m_coord_beam, ref_mb = data_from_header(bhdr, axis=2)
m_coord_beam -= ref_mb
if not np.array_equal(m_coord_beam, m_coord):
raise ValueError("m coordinates of beam model do not match those of image. Use power_beam_maker to interpolate to fits header.")
freq_beam, _ = data_from_header(bhdr, axis=freq_axis)
if not np.array_equal(freq_out, freq_beam):
raise ValueError("Freqs of beam model do not match those of image. Use power_beam_maker to interpolate to fits header.")
beam_image = load_fits(args.beam_model, dtype=args.output_type).squeeze()
elif args.beam_model.lower() == "jimbeam":
from katbeam import JimBeam
if args.band.lower() == 'l':
beam = JimBeam('MKAT-AA-L-JIM-2020')
elif args.band.lower() == 'uhf':
beam = JimBeam('MKAT-AA-UHF-JIM-2020')
else:
raise ValueError("Unkown band %s"%args.band[i])
xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij')
beam_image = np.zeros(residual.shape, dtype=args.output_type)
for v in range(freq_out.size):
# freq must be in MHz
beam_image[v] = beam.I(xx, yy, freq_out[v]/1e6).astype(args.output_type)
else:
beam_image = np.ones((nband, nx, ny), dtype=args.output_type)
if args.mask is not None:
mask = load_fits(args.mask).squeeze()
assert mask.shape == (nx, ny)
beam_image *= mask[None, :, :]
beam_image = da.from_array(beam_image, chunks=(1, -1, -1))
# if weight table is provided we use the vis space Hessian approximation
if args.weight_table is not None:
print("Solving for update using vis space approximation", file=log)
normfact = wsum
from pfb.utils.misc import plan_row_chunk
from daskms.experimental.zarr import xds_from_zarr
xds = xds_from_zarr(args.weight_table)[0]
nrow = xds.row.size
freq = xds.chan.data
nchan = freq.size
# bin edges
fmin = freq.min()
fmax = freq.max()
fbins = np.linspace(fmin, fmax, nband + 1)
# chan <-> band mapping
band_mapping = {}
chan_chunks = {}
freq_bin_idx = {}
freq_bin_counts = {}
band_map = np.zeros(freq.size, dtype=np.int32)
for band in range(nband):
indl = freq >= fbins[band]
indu = freq < fbins[band + 1] + 1e-6
band_map = np.where(indl & indu, band, band_map)
# to dask arrays
bands, bin_counts = np.unique(band_map, return_counts=True)
band_mapping = tuple(bands)
chan_chunks = {'chan': tuple(bin_counts)}
freq = da.from_array(freq, chunks=tuple(bin_counts))
bin_idx = np.append(np.array([0]), np.cumsum(bin_counts))[0:-1]
freq_bin_idx = da.from_array(bin_idx, chunks=1)
freq_bin_counts = da.from_array(bin_counts, chunks=1)
max_chan_chunk = bin_counts.max()
bin_counts = tuple(bin_counts)
# the first factor of 3 accounts for the intermediate visibilities
# produced in Hessian (i.e. complex data + real weights)
memory_per_row = (3 * max_chan_chunk * xds.WEIGHT.data.itemsize +
3 * xds.UVW.data.itemsize)
# get approx image size
pixel_bytes = np.dtype(args.output_type).itemsize
band_size = nx * ny * pixel_bytes
if args.host_address is None:
# nworker bands 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)
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)
residual = da.from_array(residual, chunks=(1, -1, -1))
x0 = da.zeros((nband, nx, ny), chunks=(1, -1, -1), dtype=residual.dtype)
xds = xds_from_zarr(args.weight_table, chunks={'row': -1, #row_chunk,
'chan': bin_counts})[0]
from pfb.opt.pcg import pcg_wgt
model = pcg_wgt(xds.UVW.data,
xds.WEIGHT.data.astype(args.output_type),
residual,
x0,
beam_image,
freq,
freq_bin_idx,
freq_bin_counts,
cell_rad,
args.wstack,
args.epsilon,
args.double_accum,
args.nvthreads,
args.sigmainv,
wsum,
args.cg_tol,
args.cg_maxit,
args.cg_minit,
args.cg_verbose,
args.cg_report_freq,
args.backtrack).compute()
else: # we use the image space approximation
print("Solving for update using image space approximation", file=log)
normfact = 1.0
from pfb.operators.psf import hessian
from ducc0.fft import r2c
iFs = np.fft.ifftshift
npad_xl = (nx_psf - nx)//2
npad_xr = nx_psf - nx - npad_xl
npad_yl = (ny_psf - ny)//2
npad_yr = ny_psf - ny - npad_yl
padding = ((0, 0), (npad_xl, npad_xr), (npad_yl, npad_yr))
unpad_x = slice(npad_xl, -npad_xr)
unpad_y = slice(npad_yl, -npad_yr)
lastsize = ny + np.sum(padding[-1])
psf_pad = iFs(psf, axes=(1, 2))
psfhat = r2c(psf_pad, axes=(1, 2), forward=True,
nthreads=nthreads, inorm=0)
psfhat = da.from_array(psfhat, chunks=(1, -1, -1))
residual = da.from_array(residual, chunks=(1, -1, -1))
x0 = da.zeros((nband, nx, ny), chunks=(1, -1, -1))
from pfb.opt.pcg import pcg_psf
model = pcg_psf(psfhat,
residual,
x0,
beam_image,
args.sigmainv,
args.nvthreads,
padding,
unpad_x,
unpad_y,
lastsize,
args.cg_tol,
args.cg_maxit,
args.cg_minit,
args.cg_verbose,
args.cg_report_freq,
args.backtrack).compute()
print("Saving results", file=log)
save_fits(args.output_filename + '_update.fits', model, hdr)
model_mfs = np.mean(model, axis=0)
save_fits(args.output_filename + '_update_mfs.fits', model_mfs, hdr_mfs)
print("All done here.", file=log)
def nnls(**kw):
'''
Minor cycle implementing non-negative least squares
'''
args = OmegaConf.create(kw)
pyscilog.log_to_file(args.output_filename + '.log')
pyscilog.enable_memory_logging(level=3)
print('Input Options:', file=log)
for key in kw.keys():
print(' %25s = %s' % (key, kw[key]), file=log)
from pfb.utils.fits import load_fits
from astropy.io import fits
import numpy as np
def resid_func(x, dirty, psfo):
"""
Returns the unattenuated residual
"""
residual = dirty - psfo.convolve(x)
residual_mfs = np.sum(residual, axis=0)
return residual, residual_mfs
def value_and_grad(x, dirty, psfo):
model_conv = psfo.convolve(x)
return np.vdot(x, model_conv - 2 * dirty), 2 * (model_conv - dirty)
def prox(x):
x[x < args.min_value] = 0.0
return x
dirty = load_fits(args.dirty).squeeze()
nband, nx, ny = dirty.shape
hdr = fits.getheader(args.dirty)
psf = load_fits(args.psf).squeeze()
_, nx_psf, ny_psf = psf.shape
hdr_psf = fits.getheader(args.psf)
wsums = np.amax(psf.reshape(-1, nx_psf * ny_psf), axis=1)
wsum = np.sum(wsums)
psf /= wsum
psf_mfs = np.sum(psf, axis=0)
assert (psf_mfs.max() - 1.0) < 1e-4
dirty /= wsum
dirty_mfs = np.sum(dirty, axis=0)
from pfb.operators.psf import PSF
psfo = PSF(psf, dirty.shape, nthreads=args.nthreads)
from pfb.opt.power_method import power_method
beta, betavec = power_method(psfo.convolve,
dirty.shape,
tol=args.pm_tol,
maxit=args.pm_maxit,
verbosity=args.pm_verbose,
report_freq=args.pm_report_freq)
fprime = partial(value_and_grad, dirty=dirty, psfo=psfo)
from pfb.opt.fista import fista
if args.x0 is None:
x0 = np.zeros_like(dirty)
else:
x0 = load_fits(args.x0, dtype=dirty.dtype).squeeze()
model = fista(x0,
beta,
fprime,
prox,
tol=args.fista_tol,
maxit=args.fista_maxit,
verbosity=args.fista_verbose,
report_freq=args.fista_report_freq)
residual, residual_mfs = resid_func(model, dirty, psfo)
from pfb.utils.fits import save_fits
save_fits(args.output_filename + '_model.fits', model, hdr)
save_fits(args.output_filename + '_residual.fits', residual, hdr)
def _spifit(**kw):
args = OmegaConf.create(kw)
OmegaConf.set_struct(args, True)
import dask.array as da
import numpy as np
from astropy.io import fits
from africanus.model.spi.dask import fit_spi_components
from pfb.utils.fits import load_fits, save_fits, data_from_header, set_wcs
from pfb.utils.misc import convolve2gaussres
# get max gausspars
gaussparf = None
if args.psf_pars is None:
if args.residual is None:
ppsource = args.image
else:
ppsource = args.residual
for image in ppsource:
try:
pphdr = fits.getheader(image)
except Exception as e:
raise e
if 'BMAJ0' in pphdr.keys():
emaj = pphdr['BMAJ0']
emin = pphdr['BMIN0']
pa = pphdr['BPA0']
gausspars = [emaj, emin, pa]
freq_idx0 = 0
elif 'BMAJ1' in pphdr.keys():
emaj = pphdr['BMAJ1']
emin = pphdr['BMIN1']
pa = pphdr['BPA1']
gausspars = [emaj, emin, pa]
freq_idx0 = 1
elif 'BMAJ' in pphdr.keys():
emaj = pphdr['BMAJ']
emin = pphdr['BMIN']
pa = pphdr['BPA']
gausspars = [emaj, emin, pa]
freq_idx0 = 0
else:
raise ValueError("No beam parameters found in residual."
"You will have to provide them manually.")
if gaussparf is None:
gaussparf = gausspars
else:
# we need to take the max in both directions
gaussparf[0] = np.maximum(gaussparf[0], gausspars[0])
gaussparf[1] = np.maximum(gaussparf[1], gausspars[1])
else:
freq_idx0 = 0 # assumption
gaussparf = list(args.psf_pars)
if args.circ_psf:
e = np.maximum(gaussparf[0], gaussparf[1])
gaussparf[0] = e
gaussparf[1] = e
gaussparf[2] = 0.0
gaussparf = tuple(gaussparf)
print("Using emaj = %3.2e, emin = %3.2e, PA = %3.2e \n" % gaussparf, file=log)
# get required data products
image_dict = {}
for i in range(len(args.image)):
image_dict[i] = {}
# load model image
model = load_fits(args.image[i], dtype=args.out_dtype).squeeze()
mhdr = fits.getheader(args.image[i])
if model.ndim < 3:
model = model[None, :, :]
l_coord, ref_l = data_from_header(mhdr, axis=1)
l_coord -= ref_l
m_coord, ref_m = data_from_header(mhdr, axis=2)
m_coord -= ref_m
if mhdr["CTYPE4"].lower() == 'freq':
freq_axis = 4
stokes_axis = 3
elif mhdr["CTYPE3"].lower() == 'freq':
freq_axis = 3
stokes_axis = 4
else:
raise ValueError("Freq axis must be 3rd or 4th")
freqs, ref_freq = data_from_header(mhdr, axis=freq_axis)
image_dict[i]['freqs'] = freqs
nband = freqs.size
npix_l = l_coord.size
npix_m = m_coord.size
xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij')
# load beam
if args.beam_model is not None:
bhdr = fits.getheader(args.beam_model[i])
l_coord_beam, ref_lb = data_from_header(bhdr, axis=1)
l_coord_beam -= ref_lb
if not np.array_equal(l_coord_beam, l_coord):
raise ValueError("l coordinates of beam model do not match "
"those of image. Use binterp to make "
"compatible beam images")
m_coord_beam, ref_mb = data_from_header(bhdr, axis=2)
m_coord_beam -= ref_mb
if not np.array_equal(m_coord_beam, m_coord):
raise ValueError("m coordinates of beam model do not match "
"those of image. Use binterp to make "
"compatible beam images")
freqs_beam, _ = data_from_header(bhdr, axis=freq_axis)
if not np.array_equal(freqs, freqs_beam):
raise ValueError("Freq coordinates of beam model do not match "
"those of image. Use binterp to make "
"compatible beam images")
beam_image = load_fits(args.beam_model[i],
dtype=args.out_dtype).squeeze()
if beam_image.ndim < 3:
beam_image = beam_image[None, :, :]
else:
beam_image = np.ones(model.shape, dtype=args.out_dtype)
image_dict[i]['beam'] = beam_image
if not args.dont_convolve:
print("Convolving model %i"%i, file=log)
# convolve model to desired resolution
model, gausskern = convolve2gaussres(model, xx, yy, gaussparf,
args.nthreads, None,
args.padding_frac)
image_dict[i]['model'] = model
# add in residuals and set threshold
if args.residual is not None:
msg = "of residual do not match those of model"
rhdr = fits.getheader(args.residual[i])
l_res, ref_lb = data_from_header(rhdr, axis=1)
l_res -= ref_lb
if not np.array_equal(l_res, l_coord):
raise ValueError("l coordinates " + msg)
m_res, ref_mb = data_from_header(rhdr, axis=2)
m_res -= ref_mb
if not np.array_equal(m_res, m_coord):
raise ValueError("m coordinates " + msg)
freqs_res, _ = data_from_header(rhdr, axis=freq_axis)
if not np.array_equal(freqs, freqs_res):
raise ValueError("Freqs " + msg)
resid = load_fits(args.residual[i],
dtype=args.out_dtype).squeeze()
if resid.ndim < 3:
resid = resid[None, :, :]
# convolve residual to same resolution as model
gausspari = ()
for b in range(nband):
key = 'BMAJ' + str(b + freq_idx0)
if key in rhdr.keys():
emaj = rhdr[key]
emin = rhdr[key]
pa = rhdr[key]
gausspari += ((emaj, emin, pa),)
elif 'BMAJ' in rhdr.keys():
emaj = rhdr['BMAJ']
emin = rhdr['BMIN']
pa = rhdr['BPA']
gausspari += ((emaj, emin, pa),)
else:
print("Can't find Gausspars in residual header, "
"unable to add residuals back in", file=log)
gausspari = None
break
if gausspari is not None and args.add_convolved_residuals:
print("Convolving residuals %i"%i, file=log)
resid, _ = convolve2gaussres(resid, xx, yy, gaussparf,
args.nthreads, gausspari,
args.padding_frac,
norm_kernel=False)
model += resid
print("Convolved residuals added to convolved model %i"%i,
file=log)
image_dict[i]['resid'] = resid
else:
image_dict[i]['resid'] = None
# concatenate images along frequency here
freqs = []
model = []
beam_image = []
resid = []
for i in image_dict.keys():
freqs.append(image_dict[i]['freqs'])
model.append(image_dict[i]['model'])
beam_image.append(image_dict[i]['beam'])
resid.append(image_dict[i]['resid'])
freqs = np.concatenate(freqs, axis=0)
Isort = np.argsort(freqs)
freqs = freqs[Isort]
model = np.concatenate(model, axis=0)
model = model[Isort]
# create header
cell_deg = mhdr['CDELT1']
ra = np.deg2rad(mhdr['CRVAL1'])
dec = np.deg2rad(mhdr['CRVAL2'])
radec = [ra, dec]
nband, nx, ny = model.shape
hdr = set_wcs(cell_deg, cell_deg, nx, ny, radec, freqs)
for i in range(1, nband+1):
hdr['BMAJ' + str(i)] = gaussparf[0]
hdr['BMIN' + str(i)] = gaussparf[1]
hdr['BPA' + str(i)] = gaussparf[2]
if args.ref_freq is None:
ref_freq = np.mean(freqs)
else:
ref_freq = args.ref_freq
hdr_mfs = set_wcs(cell_deg, cell_deg, nx, ny, radec, ref_freq)
hdr_mfs['BMAJ'] = gaussparf[0]
hdr_mfs['BMIN'] = gaussparf[1]
hdr_mfs['BPA'] = gaussparf[2]
# save convolved model
if 'm' in args.products:
name = args.output_filename + '.convolved_model.fits'
save_fits(name, model, hdr, dtype=args.out_dtype)
print("Wrote convolved model to %s" % name, file=log)
beam_image = np.concatenate(beam_image, axis=0)
beam_image = beam_image[Isort]
if 'b' in args.products:
name = args.output_filename + '.power_beam.fits'
save_fits(name, beam_image, hdr, dtype=args.out_dtype)
print("Wrote average power beam to %s" % name, file=log)
if resid[0] is not None:
resid = np.concatenate(resid, axis=0)
resid = resid[Isort]
if 'r' in args.products:
name = args.output_filename + '.convolved_residual.fits'
save_fits(name, resid, hdr, dtype=args.out_dtype)
print("Wrote convolved residuals to %s" % name, file=log)
# get threshold
counts = np.sum(resid != 0)
rms = np.sqrt(np.sum(resid**2)/counts)
rms_cube = np.std(resid.reshape(nband, npix_l*npix_m), axis=1).ravel()
threshold = args.threshold * rms
else:
print("No residual provided. Setting threshold i.t.o dynamic range. "
"Max dynamic range is %i " % args.maxdr, file=log)
threshold = model.max()/args.maxdr
rms_cube = None
print("Threshold set to %f Jy. \n" % threshold, file=log)
# beam cut off
beam_min = np.amin(beam_image, axis=0)
model = np.where(beam_min[None] > args.pb_min, model, 0.0)
# get pixels above threshold
minimage = np.amin(model, axis=0)
maskindices = np.argwhere(minimage > threshold)
nanindices = np.argwhere(minimage <= threshold)
if not maskindices.size:
raise ValueError("No components found above threshold. "
"Try lowering your threshold."
"Max of convolved model is %3.2e" % model.max())
fitcube = model[:, maskindices[:, 0], maskindices[:, 1]].T
beam_comps = beam_image[:, maskindices[:, 0], maskindices[:, 1]].T
# set weights for fit
if rms_cube is not None:
print("Using RMS in each imaging band to determine weights.", file=log)
weights = np.where(rms_cube > 0, 1.0/rms_cube**2, 0.0)
# normalise
weights /= weights.max()
else:
if args.band_weights is not None:
weights = np.array(args.band_weights)
try:
assert weights.size == nband
except Exception as e:
raise ValueError("Inconsistent weighst provided.")
print("Using provided channel weights.", file=log)
else:
print("No residual or channel weights provided. Using equal weights.", file=log)
weights = np.ones(nband, dtype=np.float64)
ncomps, _ = fitcube.shape
fitcube = da.from_array(fitcube.astype(np.float64),
chunks=(ncomps//args.nthreads, nband))
beam_comps = da.from_array(beam_comps.astype(np.float64),
chunks=(ncomps//args.nthreads, nband))
weights = da.from_array(weights.astype(np.float64), chunks=(nband))
freqsdask = da.from_array(freqs.astype(np.float64), chunks=(nband))
print("Fitting %i components" % ncomps, file=log)
alpha, alpha_err, Iref, i0_err = fit_spi_components(fitcube, weights, freqsdask,
np.float64(ref_freq), beam=beam_comps).compute()
print("Done. Writing output.", file=log)
alphamap = np.zeros(model[0].shape, dtype=model.dtype)
alphamap[...] = np.nan
alpha_err_map = np.zeros(model[0].shape, dtype=model.dtype)
alpha_err_map[...] = np.nan
i0map = np.zeros(model[0].shape, dtype=model.dtype)
i0map[...] = np.nan
i0_err_map = np.zeros(model[0].shape, dtype=model.dtype)
i0_err_map[...] = np.nan
alphamap[maskindices[:, 0], maskindices[:, 1]] = alpha
alpha_err_map[maskindices[:, 0], maskindices[:, 1]] = alpha_err
i0map[maskindices[:, 0], maskindices[:, 1]] = Iref
i0_err_map[maskindices[:, 0], maskindices[:, 1]] = i0_err
if 'I' in args.products:
# get the reconstructed cube
Irec_cube = i0map[None, :, :] * \
(freqs[:, None, None]/ref_freq)**alphamap[None, :, :]
name = args.output_filename + '.Irec_cube.fits'
save_fits(name, Irec_cube, hdr, dtype=args.out_dtype)
print("Wrote reconstructed cube to %s" % name, file=log)
# save alpha map
if 'a' in args.products:
name = args.output_filename + '.alpha.fits'
save_fits(name, alphamap, hdr_mfs, dtype=args.out_dtype)
print("Wrote alpha map to %s" % name, file=log)
# save alpha error map
if 'e' in args.products:
name = args.output_filename + '.alpha_err.fits'
save_fits(name, alpha_err_map, mhdr, dtype=args.out_dtype)
print("Wrote alpha error map to %s" % name, file=log)
# save I0 map
if 'i' in args.products:
name = args.output_filename + '.I0.fits'
save_fits(name, i0map, mhdr, dtype=args.out_dtype)
print("Wrote I0 map to %s" % name, file=log)
# save I0 error map
if 'k' in args.products:
name = args.output_filename + '.I0_err.fits'
save_fits(name, i0_err_map, mhdr, dtype=args.out_dtype)
print("Wrote I0 error map to %s" % name, file=log)
print("All done here", file=log)
def test_forwardmodel(do_beam, do_gains, tmp_path_factory):
test_dir = tmp_path_factory.mktemp("test_pfb")
packratt.get('/test/ms/2021-06-24/elwood/test_ascii_1h60.0s.MS.tar',
str(test_dir))
import numpy as np
np.random.seed(420)
from numpy.testing import assert_allclose
from pyrap.tables import table
ms = table(str(test_dir / 'test_ascii_1h60.0s.MS'), readonly=False)
spw = table(str(test_dir / 'test_ascii_1h60.0s.MS::SPECTRAL_WINDOW'))
utime = np.unique(ms.getcol('TIME'))
freq = spw.getcol('CHAN_FREQ').squeeze()
freq0 = np.mean(freq)
ntime = utime.size
nchan = freq.size
nant = np.maximum(
ms.getcol('ANTENNA1').max(),
ms.getcol('ANTENNA2').max()) + 1
ncorr = ms.getcol('FLAG').shape[-1]
uvw = ms.getcol('UVW')
nrow = uvw.shape[0]
u_max = abs(uvw[:, 0]).max()
v_max = abs(uvw[:, 1]).max()
uv_max = np.maximum(u_max, v_max)
# image size
from africanus.constants import c as lightspeed
cell_N = 1.0 / (2 * uv_max * freq.max() / lightspeed)
srf = 2.0
cell_rad = cell_N / srf
cell_size = cell_rad * 180 / np.pi
print("Cell size set to %5.5e arcseconds" % cell_size)
fov = 2
npix = int(fov / cell_size)
if npix % 2:
npix += 1
nx = npix
ny = npix
print("Image size set to (%i, %i, %i)" % (nchan, nx, ny))
# model
model = np.zeros((nchan, nx, ny), dtype=np.float64)
nsource = 10
Ix = np.random.randint(0, npix, nsource)
Iy = np.random.randint(0, npix, nsource)
alpha = -0.7 + 0.1 * np.random.randn(nsource)
I0 = 1.0 + np.abs(np.random.randn(nsource))
for i in range(nsource):
model[:, Ix[i], Iy[i]] = I0[i] * (freq / freq0)**alpha[i]
if do_beam:
# primary beam
from katbeam import JimBeam
beam = JimBeam('MKAT-AA-L-JIM-2020')
l_coord = -np.arange(-(nx // 2), nx // 2) * cell_size
m_coord = np.arange(-(ny // 2), ny // 2) * cell_size
xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij')
pbeam = np.zeros((nchan, nx, ny), dtype=np.float64)
for i in range(nchan):
pbeam[i] = beam.I(xx, yy, freq[i] / 1e6) # freq in MHz
model_att = pbeam * model
bm = 'JimBeam'
else:
model_att = model
bm = None
# model vis
from ducc0.wgridder import dirty2ms
model_vis = np.zeros((nrow, nchan, ncorr), dtype=np.complex128)
for c in range(nchan):
model_vis[:, c:c + 1, 0] = dirty2ms(uvw,
freq[c:c + 1],
model_att[c],
pixsize_x=cell_rad,
pixsize_y=cell_rad,
epsilon=1e-8,
do_wstacking=True,
nthreads=8)
model_vis[:, c, -1] = model_vis[:, c, 0]
ms.putcol('MODEL_DATA', model_vis.astype(np.complex64))
if do_gains:
t = (utime - utime.min()) / (utime.max() - utime.min())
nu = 2.5 * (freq / freq0 - 1.0)
from africanus.gps.utils import abs_diff
tt = abs_diff(t, t)
lt = 0.25
Kt = 0.1 * np.exp(-tt**2 / (2 * lt**2))
Lt = np.linalg.cholesky(Kt + 1e-10 * np.eye(ntime))
vv = abs_diff(nu, nu)
lv = 0.1
Kv = 0.1 * np.exp(-vv**2 / (2 * lv**2))
Lv = np.linalg.cholesky(Kv + 1e-10 * np.eye(nchan))
L = (Lt, Lv)
from pfb.utils.misc import kron_matvec
jones = np.zeros((ntime, nant, nchan, 1, ncorr), dtype=np.complex128)
for p in range(nant):
for c in [0, -1]: # for now only diagonal
xi_amp = np.random.randn(ntime, nchan)
amp = np.exp(-nu[None, :]**2 +
kron_matvec(L, xi_amp).reshape(ntime, nchan))
xi_phase = np.random.randn(ntime, nchan)
phase = kron_matvec(L, xi_phase).reshape(ntime, nchan)
jones[:, p, :, 0, c] = amp * np.exp(1.0j * phase)
# corrupted vis
model_vis = model_vis.reshape(nrow, nchan, 1, 2, 2)
from africanus.calibration.utils import chunkify_rows
time = ms.getcol('TIME')
row_chunks, tbin_idx, tbin_counts = chunkify_rows(time, ntime)
ant1 = ms.getcol('ANTENNA1')
ant2 = ms.getcol('ANTENNA2')
from africanus.calibration.utils import corrupt_vis
vis = corrupt_vis(tbin_idx, tbin_counts, ant1, ant2, jones,
model_vis).reshape(nrow, nchan, ncorr)
model_vis[:, :, 0, 0, 0] = 1.0 + 0j
model_vis[:, :, 0, -1, -1] = 1.0 + 0j
muellercol = corrupt_vis(tbin_idx, tbin_counts, ant1, ant2, jones,
model_vis).reshape(nrow, nchan, ncorr)
ms.putcol('DATA', vis.astype(np.complex64))
ms.putcol('CORRECTED_DATA', muellercol.astype(np.complex64))
ms.close()
mcol = 'CORRECTED_DATA'
else:
ms.putcol('DATA', model_vis.astype(np.complex64))
mcol = None
from pfb.workers.grid.dirty import _dirty
_dirty(ms=str(test_dir / 'test_ascii_1h60.0s.MS'),
data_column="DATA",
weight_column='WEIGHT',
imaging_weight_column=None,
flag_column='FLAG',
mueller_column=mcol,
row_chunks=None,
epsilon=1e-5,
wstack=True,
mock=False,
double_accum=True,
output_filename=str(test_dir / 'test'),
nband=nchan,
field_of_view=fov,
super_resolution_factor=srf,
cell_size=None,
nx=None,
ny=None,
output_type='f4',
nworkers=1,
nthreads_per_worker=1,
nvthreads=8,
mem_limit=8,
nthreads=8,
host_address=None)
from pfb.workers.grid.psf import _psf
_psf(ms=str(test_dir / 'test_ascii_1h60.0s.MS'),
data_column="DATA",
weight_column='WEIGHT',
imaging_weight_column=None,
flag_column='FLAG',
mueller_column=mcol,
row_out_chunk=-1,
row_chunks=None,
epsilon=1e-5,
wstack=True,
mock=False,
psf_oversize=2,
double_accum=True,
output_filename=str(test_dir / 'test'),
nband=nchan,
field_of_view=fov,
super_resolution_factor=srf,
cell_size=None,
nx=None,
ny=None,
output_type='f4',
nworkers=1,
nthreads_per_worker=1,
nvthreads=8,
mem_limit=8,
nthreads=8,
host_address=None)
# solve for model using pcg and mask
mask = np.any(model, axis=0)
from astropy.io import fits
from pfb.utils.fits import save_fits
hdr = fits.getheader(str(test_dir / 'test_dirty.fits'))
save_fits(str(test_dir / 'test_model.fits'), model, hdr)
save_fits(str(test_dir / 'test_mask.fits'), mask, hdr)
from pfb.workers.deconv.forward import _forward
_forward(residual=str(test_dir / 'test_dirty.fits'),
psf=str(test_dir / 'test_psf.fits'),
mask=str(test_dir / 'test_mask.fits'),
beam_model=bm,
band='L',
weight_table=str(test_dir / 'test.zarr'),
output_filename=str(test_dir / 'test'),
nband=nchan,
output_type='f4',
epsilon=1e-5,
sigmainv=0.0,
wstack=True,
double_accum=True,
cg_tol=1e-6,
cg_minit=10,
cg_maxit=100,
cg_verbose=0,
cg_report_freq=10,
backtrack=False,
nworkers=1,
nthreads_per_worker=1,
nvthreads=1,
mem_limit=8,
nthreads=1,
host_address=None)
# get inferred model
from pfb.utils.fits import load_fits
model_inferred = load_fits(str(test_dir / 'test_update.fits')).squeeze()
for i in range(nsource):
if do_beam:
beam = pbeam[:, Ix[i], Iy[i]]
assert_allclose(
0.0,
beam *
(model_inferred[:, Ix[i], Iy[i]] - model[:, Ix[i], Iy[i]]),
atol=1e-4)
else:
assert_allclose(0.0,
model_inferred[:, Ix[i], Iy[i]] -
model[:, Ix[i], Iy[i]],
atol=1e-4)
def _clean(**kw):
args = OmegaConf.create(kw)
OmegaConf.set_struct(args, True)
import numpy as np
import numexpr as ne
import dask
import dask.array as da
from dask.distributed import performance_report
from pfb.utils.fits import load_fits, set_wcs, save_fits, data_from_header
from pfb.opt.hogbom import hogbom
from astropy.io import fits
print("Loading dirty", file=log)
dirty = load_fits(args.dirty, dtype=args.output_type).squeeze()
nband, nx, ny = dirty.shape
hdr = fits.getheader(args.dirty)
print("Loading psf", file=log)
psf = load_fits(args.psf, dtype=args.output_type).squeeze()
_, nx_psf, ny_psf = psf.shape
hdr_psf = fits.getheader(args.psf)
wsums = np.amax(psf.reshape(-1, nx_psf * ny_psf), axis=1)
wsum = np.sum(wsums)
psf /= wsum
psf_mfs = np.sum(psf, axis=0)
assert (psf_mfs.max() - 1.0) < 1e-4
dirty /= wsum
dirty_mfs = np.sum(dirty, axis=0)
# get info required to set WCS
ra = np.deg2rad(hdr['CRVAL1'])
dec = np.deg2rad(hdr['CRVAL2'])
radec = [ra, dec]
cell_deg = np.abs(hdr['CDELT1'])
if cell_deg != np.abs(hdr['CDELT2']):
raise NotImplementedError('cell sizes have to be equal')
cell_rad = np.deg2rad(cell_deg)
freq_out, ref_freq = data_from_header(hdr, axis=3)
hdr_mfs = set_wcs(cell_deg, cell_deg, nx, ny, radec, ref_freq)
save_fits(args.output_filename + '_dirty_mfs.fits',
dirty_mfs,
hdr_mfs,
dtype=args.output_type)
# set up Hessian approximation
if args.weight_table is not None:
normfact = wsum
from africanus.gridding.wgridder.dask import hessian
from pfb.utils.misc import plan_row_chunk
from daskms.experimental.zarr import xds_from_zarr
xds = xds_from_zarr(args.weight_table)[0]
nrow = xds.row.size
freqs = xds.chan.data
nchan = freqs.size
# bin edges
fmin = freqs.min()
fmax = freqs.max()
fbins = np.linspace(fmin, fmax, nband + 1)
# chan <-> band mapping
band_mapping = {}
chan_chunks = {}
freq_bin_idx = {}
freq_bin_counts = {}
band_map = np.zeros(freqs.size, dtype=np.int32)
for band in range(nband):
indl = freqs >= fbins[band]
indu = freqs < fbins[band + 1] + 1e-6
band_map = np.where(indl & indu, band, band_map)
# to dask arrays
bands, bin_counts = np.unique(band_map, return_counts=True)
band_mapping = tuple(bands)
chan_chunks = {'chan': tuple(bin_counts)}
freqs = da.from_array(freqs, chunks=tuple(bin_counts))
bin_idx = np.append(np.array([0]), np.cumsum(bin_counts))[0:-1]
freq_bin_idx = da.from_array(bin_idx, chunks=1)
freq_bin_counts = da.from_array(bin_counts, chunks=1)
max_chan_chunk = bin_counts.max()
bin_counts = tuple(bin_counts)
# the first factor of 3 accounts for the intermediate visibilities
# produced in Hessian (i.e. complex data + real weights)
memory_per_row = (3 * max_chan_chunk * xds.WEIGHT.data.itemsize +
3 * xds.UVW.data.itemsize)
# get approx image size
pixel_bytes = np.dtype(args.output_type).itemsize
band_size = nx * ny * pixel_bytes
if args.host_address is None:
# nworker bands 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)
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)
def convolver(x):
model = da.from_array(x, chunks=(1, nx, ny), name=False)
xds = xds_from_zarr(args.weight_table,
chunks={
'row': row_chunk,
'chan': bin_counts
})[0]
convolvedim = hessian(xds.UVW.data,
freqs,
model,
freq_bin_idx,
freq_bin_counts,
cell_rad,
weights=xds.WEIGHT.data.astype(
args.output_type),
nthreads=args.nvthreads,
epsilon=args.epsilon,
do_wstacking=args.wstack,
double_accum=args.double_accum)
return convolvedim
else:
normfact = 1.0
from pfb.operators.psf import hessian
from ducc0.fft import r2c
iFs = np.fft.ifftshift
npad_xl = (nx_psf - nx) // 2
npad_xr = nx_psf - nx - npad_xl
npad_yl = (ny_psf - ny) // 2
npad_yr = ny_psf - ny - npad_yl
padding = ((0, 0), (npad_xl, npad_xr), (npad_yl, npad_yr))
unpad_x = slice(npad_xl, -npad_xr)
unpad_y = slice(npad_yl, -npad_yr)
lastsize = ny + np.sum(padding[-1])
psf_pad = iFs(psf, axes=(1, 2))
psfhat = r2c(psf_pad,
axes=(1, 2),
forward=True,
nthreads=nthreads,
inorm=0)
psfhat = da.from_array(psfhat, chunks=(1, -1, -1))
def convolver(x):
model = da.from_array(x, chunks=(1, nx, ny), name=False)
convolvedim = hessian(model, psfhat, padding, nvthreads, unpad_x,
unpad_y, lastsize)
return convolvedim
# psfo = PSF(psf, dirty.shape, nthreads=args.nthreads)
# def convolver(x): return psfo.convolve(x)
rms = np.std(dirty_mfs)
rmax = np.abs(dirty_mfs).max()
print("Iter %i: peak residual = %f, rms = %f" % (0, rmax, rms), file=log)
residual = dirty.copy()
residual_mfs = dirty_mfs.copy()
model = np.zeros_like(residual)
for k in range(args.nmiter):
print("Running Hogbom", file=log)
x = hogbom(residual,
psf,
gamma=args.hb_gamma,
pf=args.hb_peak_factor,
maxit=args.hb_maxit,
verbosity=args.hb_verbose,
report_freq=args.hb_report_freq)
model += x
print("Getting residual", file=log)
convimage = convolver(model)
dask.visualize(convimage,
filename=args.output_filename + '_hessian' + str(k) +
'_graph.pdf',
optimize_graph=False)
with performance_report(filename=args.output_filename + '_hessian' +
str(k) + '_per.html'):
convimage = dask.compute(convimage, optimize_graph=False)[0]
ne.evaluate('dirty - convimage/normfact',
out=residual,
casting='same_kind')
ne.evaluate('sum(residual, axis=0)',
out=residual_mfs,
casting='same_kind')
rms = np.std(residual_mfs)
rmax = np.abs(residual_mfs).max()
print("Iter %i: peak residual = %f, rms = %f" % (k + 1, rmax, rms),
file=log)
print("Saving results", file=log)
save_fits(args.output_filename + '_model.fits', model, hdr)
model_mfs = np.mean(model, axis=0)
save_fits(args.output_filename + '_model_mfs.fits', model_mfs, hdr_mfs)
save_fits(args.output_filename + '_residual.fits',
residual * wsums[:, None, None], hdr)
save_fits(args.output_filename + '_residual.fits', residual_mfs, hdr_mfs)
print("All done here.", file=log)
