def test_L_beam_image():
beam = JimBeam('MKAT-AA-L-JIM-2020')
return showbeam(beam, 1420, 'VV', 5.)
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 test_sample_beam_values(name, pol, x, y, freqMHz, value):
beam = JimBeam(name)
pattern = getattr(beam, pol)
assert pattern(x, y, freqMHz) == pytest.approx(value)
def test_UHF_beam_image():
beam = JimBeam('MKAT-AA-UHF-JIM-2020')
return showbeam(beam, 800, 'HH', 10.)
def test_unknown_model_name():
with pytest.raises(ValueError):
JimBeam('MKAT-AA-UHF-JIM-2012')
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 main():
# -------------------------------------------------
# Options and some error checking
parser = OptionParser(usage='%prog [options] input_fits')
parser.add_option('--band',
dest='band',
help='Select [U]HF or [L]-band (default = L-band)',
default='L')
parser.add_option(
'--freq',
dest='freq',
help=
'Frequency in MHz at which to evaluate beam model (default = get from input FITS header)',
default='')
parser.add_option(
'--pbcut',
dest='pbcut',
help=
'Primary beam gain level beyond which to blank output images (default = 0.3)',
default=0.3)
parser.add_option(
'--noavg',
dest='azavg',
help=
'Do not azimuthally-average the primary beam pattern (default = do this)',
default=True,
action='store_false')
parser.add_option(
'--nopbcorfits',
dest='savepbcor',
help=
'Do not save primary beam corrected image (default = save corrected image)',
action='store_false',
default=True)
parser.add_option(
'--nopbfits',
dest='savepb',
help='Do not save primary beam image (default = save PB image)',
action='store_false',
default=True)
parser.add_option(
'--nowtfits',
dest='savewt',
help='Do not save weight image (default = save weight image)',
action='store_false',
default=True)
parser.add_option(
'--pbcorname',
dest='pbcor_fits',
help=
'Filename for primary beam corrected image (default = based on input image)',
default='')
parser.add_option(
'--pbname',
dest='pb_fits',
help='Filename for primary beam image (default = based on input image)',
default='')
parser.add_option(
'--wtname',
dest='wt_fits',
help='Filename for weight image (default = based on input image)',
default='')
parser.add_option(
'--overwrite',
'-f',
dest='overwrite',
help='Overwrite any existing output files (default = do not overwrite)',
action='store_true',
default=False)
(options, args) = parser.parse_args()
# Input FITS file
if len(args) != 1:
msg('Please provide a FITS image')
sys.exit()
else:
input_fits = args[0].rstrip('/')
# MeerKAT band
band = options.band[0].lower()
if band not in ['l', 'u']:
msg('Please check requested band')
sys.exit()
# Frequency
freq = options.freq
# Primary beam cut level
pbcut = float(options.pbcut)
# Azimuthal averaging
azavg = options.azavg
if azavg:
try:
from skued import azimuthal_average as aa
except:
msg('scikit-ued not found, azimuthal averaging is not available.')
msg('Try: pip install scikit-ued')
azavg = False
# Output files
savepbcor = options.savepbcor
savepb = options.savepb
savewt = options.savewt
if [savepbcor, savepb, savewt] == [False, False, False]:
msg('Nothing to do, please check your options')
sys.exit()
# Generate output names if not provided
pbcor_fits = options.pbcor_fits
pb_fits = options.pb_fits
wt_fits = options.wt_fits
if pbcor_fits == '':
pbcor_fits = input_fits.replace('.fits', '.pbcor.fits')
check_name(input_fits, pbcor_fits)
if pb_fits == '':
pb_fits = input_fits.replace('.fits', '.pb.fits')
check_name(input_fits, pb_fits)
if wt_fits == '':
wt_fits = input_fits.replace('.fits', '.wt.fits')
check_name(input_fits, wt_fits)
# Bail out if some files will be overwritten
overwrite = options.overwrite
if not overwrite:
file_check = []
if savepbcor:
file_check.append(check_file(pbcor_fits))
if savepb:
file_check.append(check_file(pb_fits))
if savewt:
file_check.append(check_file(wt_fits))
if True in file_check:
sys.exit()
pol = 'I' # hardwired Stokes I beam for now
# -------------------------------------------------
# Set up band
if band == 'l':
beam_model = 'MKAT-AA-L-JIM-2020'
band = 'L-band'
elif band == 'u':
beam_model = 'MKAT-AA-UHF-JIM-2020'
band = 'UHF'
msg('Band is ' + band)
msg('Beam model is ' + beam_model)
beam = JimBeam(beam_model)
# Get header info
msg('Reading FITS image')
msg(' <--- ' + input_fits)
nx, ny, dx, dy, fitsfreq = get_header(input_fits)
if nx != ny or abs(dx) != abs(dy):
msg('Can only handle square images / pixels')
sys.exit()
extent = nx * dx # degrees
if freq == '':
freq = fitsfreq / 1e6
else:
freq = float(freq)
msg('Evaluating beam at ' + str(round(freq, 4)) + ' MHz')
interval = numpy.linspace(-extent / 2.0, extent / 2.0, nx)
xx, yy = numpy.meshgrid(interval, interval)
beam_image = beam.I(xx, yy, freq)
msg('Masking beam beyond the ' + str(pbcut) + ' level')
mask = beam_image < pbcut
beam_image[mask] = numpy.nan
if azavg:
msg('Azimuthally averaging the beam pattern')
x0 = int(nx / 2)
y0 = int(ny / 2)
radius, average = aa(beam_image, center=(x0, y0))
# This can probably be sped up...
for y in range(0, ny):
for x in range(0, nx):
val = (((float(y) - y0)**2.0) + ((float(x) - x0)**2.0))**0.5
beam_image[y][x] = average[int(val)]
if savepbcor:
msg('Correcting image')
input_image = get_image(input_fits)
pbcor_image = input_image / beam_image
msg('Writing primary beam corrected image')
msg(' ---> ' + pbcor_fits)
copyfile(input_fits, pbcor_fits)
flush_fits(pbcor_image, pbcor_fits)
if savepb:
msg('Writing primary beam image')
msg(' ---> ' + pb_fits)
copyfile(input_fits, pb_fits)
flush_fits(beam_image, pb_fits)
if savewt:
msg('Writing weight (pb^2) image')
msg(' ---> ' + wt_fits)
copyfile(input_fits, wt_fits)
flush_fits(beam_image**2.0, wt_fits)
msg('Done')
def _binterp(**kw):
args = OmegaConf.create(kw)
OmegaConf.set_struct(args, True)
from pfb.utils.fits import save_fits
import dask
import dask.array as da
import numpy as np
from numba import jit
from astropy.io import fits
import warnings
from africanus.rime import parallactic_angles
from pfb.utils.fits import load_fits, save_fits, data_from_header
from daskms import xds_from_ms, xds_from_table
if args.ms is None:
if args.beam_model.lower() == 'jimbeam':
for image in args.image:
mhdr = fits.getheader(image)
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")
freq, ref_freq = data_from_header(mhdr, axis=freq_axis)
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((freq.size, l_coord.size, m_coord.size),
dtype=args.out_dtype)
for v in range(freq.size):
# freq must be in MHz
beam_image[v] = beam.I(xx, yy, freq[v] / 1e6).astype(
args.out_dtype)
if args.output_dir in image:
idx = len(args.output_dir)
iname = image[idx::]
outname = iname + '.' + args.postfix
else:
outname = image + '.' + args.postfix
beam_image = np.expand_dims(beam_image,
axis=3 - stokes_axis + 1)
save_fits(args.output_dir + outname,
beam_image,
mhdr,
dtype=args.out_dtype)
else:
raise NotImplementedError("Not there yet, sorry")
print("All done here.", file=log)
# @jit(nopython=True, nogil=True, cache=True)
# def _unflagged_counts(flags, time_idx, out):
# for i in range(time_idx.size):
# ilow = time_idx[i]
# ihigh = time_idx[i+1]
# out[i] = np.sum(~flags[ilow:ihigh])
# return out
# def extract_dde_info(args, freqs):
# """
# Computes paralactic angles, antenna scaling and pointing information
# required for beam interpolation.
# """
# # get ms info required to compute paralactic angles and weighted sum
# nband = freqs.size
# if args.ms is not None:
# utimes = []
# unflag_counts = []
# ant_pos = None
# phase_dir = None
# for ms_name in args.ms:
# # get antenna positions
# ant = xds_from_table(ms_name + '::ANTENNA')[0].compute()
# if ant_pos is None:
# ant_pos = ant['POSITION'].data
# else: # check all are the same
# tmp = ant['POSITION']
# if not np.array_equal(ant_pos, tmp):
# raise ValueError(
# "Antenna positions not the same across measurement sets")
# # get phase center for field
# field = xds_from_table(ms_name + '::FIELD')[0].compute()
# if phase_dir is None:
# phase_dir = field['PHASE_DIR'][args.field].data.squeeze()
# else:
# tmp = field['PHASE_DIR'][args.field].data.squeeze()
# if not np.array_equal(phase_dir, tmp):
# raise ValueError(
# 'Phase direction not the same across measurement sets')
# # get unique times and count flags
# xds = xds_from_ms(ms_name, columns=["TIME", "FLAG_ROW"], group_cols=[
# "FIELD_ID"])[args.field]
# utime, time_idx = np.unique(
# xds.TIME.data.compute(), return_index=True)
# ntime = utime.size
# # extract subset of times
# if args.sparsify_time > 1:
# I = np.arange(0, ntime, args.sparsify_time)
# utime = utime[I]
# time_idx = time_idx[I]
# ntime = utime.size
# utimes.append(utime)
# flags = xds.FLAG_ROW.data.compute()
# unflag_count = _unflagged_counts(flags.astype(
# np.int32), time_idx, np.zeros(ntime, dtype=np.int32))
# unflag_counts.append(unflag_count)
# utimes = np.concatenate(utimes)
# unflag_counts = np.concatenate(unflag_counts)
# ntimes = utimes.size
# # compute paralactic angles
# parangles = parallactic_angles(utimes, ant_pos, phase_dir)
# # mean over antanna nant -> 1
# parangles = np.mean(parangles, axis=1, keepdims=True)
# nant = 1
# # beam_cube_dde requirements
# ant_scale = np.ones((nant, nband, 2), dtype=np.float64)
# point_errs = np.zeros((ntimes, nant, nband, 2), dtype=np.float64)
# return (parangles,
# da.from_array(ant_scale, chunks=ant_scale.shape),
# point_errs,
# unflag_counts,
# True)
# else:
# ntimes = 1
# nant = 1
# parangles = np.zeros((ntimes, nant,), dtype=np.float64)
# ant_scale = np.ones((nant, nband, 2), dtype=np.float64)
# point_errs = np.zeros((ntimes, nant, nband, 2), dtype=np.float64)
# unflag_counts = np.array([1])
# return (parangles, ant_scale, point_errs, unflag_counts, False)
# def make_power_beam(args, lm_source, freqs, use_dask):
# print("Loading fits beam patterns from %s" % args.beam_model)
# from glob import glob
# paths = glob(args.beam_model + '**_**.fits')
# beam_hdr = None
# if args.corr_type == 'linear':
# corr1 = 'XX'
# corr2 = 'YY'
# elif args.corr_type == 'circular':
# corr1 = 'LL'
# corr2 = 'RR'
# else:
# raise KeyError(
# "Unknown corr_type supplied. Only 'linear' or 'circular' supported")
# for path in paths:
# if corr1.lower() in path[-10::]:
# if 're' in path[-7::]:
# corr1_re = load_fits(path)
# if beam_hdr is None:
# beam_hdr = fits.getheader(path)
# elif 'im' in path[-7::]:
# corr1_im = load_fits(path)
# else:
# raise NotImplementedError("Only re/im patterns supported")
# elif corr2.lower() in path[-10::]:
# if 're' in path[-7::]:
# corr2_re = load_fits(path)
# elif 'im' in path[-7::]:
# corr2_im = load_fits(path)
# else:
# raise NotImplementedError("Only re/im patterns supported")
# # get power beam
# beam_amp = (corr1_re**2 + corr1_im**2 + corr2_re**2 + corr2_im**2)/2.0
# # get cube in correct shape for interpolation code
# beam_amp = np.ascontiguousarray(np.transpose(beam_amp, (1, 2, 0))
# [:, :, :, None, None])
# # get cube info
# if beam_hdr['CUNIT1'].lower() != "deg":
# raise ValueError("Beam image units must be in degrees")
# npix_l = beam_hdr['NAXIS1']
# refpix_l = beam_hdr['CRPIX1']
# delta_l = beam_hdr['CDELT1']
# l_min = (1 - refpix_l)*delta_l
# l_max = (1 + npix_l - refpix_l)*delta_l
# if beam_hdr['CUNIT2'].lower() != "deg":
# raise ValueError("Beam image units must be in degrees")
# npix_m = beam_hdr['NAXIS2']
# refpix_m = beam_hdr['CRPIX2']
# delta_m = beam_hdr['CDELT2']
# m_min = (1 - refpix_m)*delta_m
# m_max = (1 + npix_m - refpix_m)*delta_m
# if (l_min > lm_source[:, 0].min() or m_min > lm_source[:, 1].min() or
# l_max < lm_source[:, 0].max() or m_max < lm_source[:, 1].max()):
# raise ValueError("The supplied beam is not large enough")
# beam_extents = np.array([[l_min, l_max], [m_min, m_max]])
# # get frequencies
# if beam_hdr["CTYPE3"].lower() != 'freq':
# raise ValueError(
# "Cubes are assumed to be in format [nchan, nx, ny]")
# nchan = beam_hdr['NAXIS3']
# refpix = beam_hdr['CRPIX3']
# delta = beam_hdr['CDELT3'] # assumes units are Hz
# freq0 = beam_hdr['CRVAL3']
# bfreqs = freq0 + np.arange(1 - refpix, 1 + nchan - refpix) * delta
# if bfreqs[0] > freqs[0] or bfreqs[-1] < freqs[-1]:
# warnings.warn("The supplied beam does not have sufficient "
# "bandwidth. Beam frequencies:")
# with np.printoptions(precision=2):
# print(bfreqs)
# if use_dask:
# return (da.from_array(beam_amp, chunks=beam_amp.shape),
# da.from_array(beam_extents, chunks=beam_extents.shape),
# da.from_array(bfreqs, bfreqs.shape))
# else:
# return beam_amp, beam_extents, bfreqs
# def interpolate_beam(ll, mm, freqs, args):
# """
# Interpolate beam to image coordinates and optionally compute average
# over time if MS is provoded
# """
# nband = freqs.size
# print("Interpolating beam")
# parangles, ant_scale, point_errs, unflag_counts, use_dask = extract_dde_info(
# args, freqs)
# lm_source = np.vstack((ll.ravel(), mm.ravel())).T
# beam_amp, beam_extents, bfreqs = make_power_beam(
# args, lm_source, freqs, use_dask)
# # interpolate beam
# if use_dask:
# from africanus.rime.dask import beam_cube_dde
# lm_source = da.from_array(lm_source, chunks=lm_source.shape)
# freqs = da.from_array(freqs, chunks=freqs.shape)
# # compute ncpu images at a time to avoid memory errors
# ntimes = parangles.shape[0]
# I = np.arange(0, ntimes, args.ncpu)
# nchunks = I.size
# I = np.append(I, ntimes)
# beam_image = np.zeros((ll.size, 1, nband), dtype=beam_amp.dtype)
# for i in range(nchunks):
# ilow = I[i]
# ihigh = I[i+1]
# part_parangles = da.from_array(
# parangles[ilow:ihigh], chunks=(1, 1))
# part_point_errs = da.from_array(
# point_errs[ilow:ihigh], chunks=(1, 1, freqs.size, 2))
# # interpolate and remove redundant axes
# part_beam_image = beam_cube_dde(beam_amp, beam_extents, bfreqs,
# lm_source, part_parangles, part_point_errs,
# ant_scale, freqs).compute()[:, :, 0, :, 0, 0]
# # weighted sum over time
# beam_image += np.sum(part_beam_image *
# unflag_counts[None, ilow:ihigh, None], axis=1, keepdims=True)
# # normalise by sum of weights
# beam_image /= np.sum(unflag_counts)
# # remove time axis
# beam_image = beam_image[:, 0, :]
# else:
# from africanus.rime.fast_beam_cubes import beam_cube_dde
# beam_image = beam_cube_dde(beam_amp, beam_extents, bfreqs,
# lm_source, parangles, point_errs,
# ant_scale, freqs).squeeze()
# # swap source and freq axes and reshape to image shape
# beam_source = np.transpose(beam_image, axes=(1, 0))
# return beam_source.squeeze().reshape((freqs.size, *ll.shape))
# def main(args):
# # get coord info
# hdr = fits.getheader(args.image)
# 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
# if hdr["CTYPE4"].lower() == 'freq':
# freq_axis = 4
# elif hdr["CTYPE3"].lower() == 'freq':
# freq_axis = 3
# else:
# raise ValueError("Freq axis must be 3rd or 4th")
# freqs, ref_freq = data_from_header(hdr, axis=freq_axis)
# xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij')
# # interpolate primary beam to fits header and optionally average over time
# beam_image = interpolate_beam(xx, yy, freqs, args)
# # save power beam
# save_fits(args.output_filename, beam_image, hdr)
# print("Wrote interpolated beam cube to %s \n" % args.output_filename)
# return
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)