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 _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 _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 _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)