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 main(args):
# 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
print("Super resolution factor = ", cell_N / cell_rad)
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)
if args.nx is None or args.ny is None:
fov = args.fov * 3600
npix = int(fov / args.cell_size)
if npix % 2:
npix += 1
args.nx = npix
args.ny = npix
if args.nband is None:
args.nband = freq.size
print("Image size set to (%i, %i, %i)" % (args.nband, args.nx, args.ny))
# 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,
optimise_chunks=True,
data_column=args.data_column,
weight_column=args.weight_column,
imaging_weight_column=args.imaging_weight_column,
model_column=args.model_column,
flag_column=args.flag_column)
freq_out = R.freq_out
radec = R.radec
# 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,
2 * args.nx, 2 * args.ny, radec, freq_out)
hdr_psf_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600,
2 * args.nx, 2 * args.ny, radec, np.mean(freq_out))
# psf
if args.psf is not None:
compare_headers(hdr_psf, fits.getheader(args.psf))
psf_array = load_fits(args.psf)
else:
psf_array = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf_array, hdr_psf)
psf_max = np.amax(psf_array.reshape(args.nband, 4 * args.nx * args.ny),
axis=1)
wsum = np.sum(psf_max)
counts = np.sum(psf_max > 0)
psf_max_mean = wsum / counts
psf_array /= psf_max_mean
psf = PSF(psf_array, args.nthreads)
psf_max = np.amax(psf_array.reshape(args.nband, 4 * args.nx * args.ny),
axis=1)
psf_max[psf_max < 1e-15] = 1e-15
if args.dirty is not None:
compare_headers(hdr, fits.getheader(args.dirty))
dirty = load_fits(args.dirty)
else:
dirty = R.make_dirty()
save_fits(args.outfile + '_dirty.fits', dirty, hdr)
dirty /= psf_max_mean
# mfs residual
wsum = np.sum(psf_max)
dirty_mfs = np.sum(dirty, axis=0) / wsum
rmax = np.abs(dirty_mfs).max()
rms = np.std(dirty_mfs)
save_fits(args.outfile + '_dirty_mfs.fits', dirty_mfs, hdr_mfs)
psf_mfs = np.sum(psf_array, axis=0) / wsum
save_fits(
args.outfile + '_psf_mfs.fits', psf_mfs[args.nx // 2:3 * args.nx // 2,
args.ny // 2:3 * args.ny // 2],
hdr_mfs)
# mask
if args.mask is not None:
mask = load_fits(args.mask, dtype=np.int64)
if mask.shape != (args.nx, args.ny):
raise ValueError("Mask has incorrect shape")
else:
mask = np.ones((args.nx, args.ny), dtype=np.int64)
if args.point_mask is not None:
pmask = load_fits(args.point_mask, dtype=np.bool)
if pmask.shape != (args.nx, args.ny):
raise ValueError("Mask has incorrect shape")
else:
pmask = None
# Reporting
print("At iteration 0 peak of residual is %f and rms is %f" % (rmax, rms))
report_iters = list(np.arange(0, args.maxit, args.report_freq))
if report_iters[-1] != args.maxit - 1:
report_iters.append(args.maxit - 1)
# set up point sources
phi = Dirac(args.nband, args.nx, args.ny, mask=pmask)
dual = np.zeros((args.nband, args.nx, args.ny), dtype=np.float64)
weights_21 = np.where(phi.mask, 1, np.inf)
# preconditioning matrix
def hess(beta):
return phi.hdot(psf.convolve(
phi.dot(beta))) + beta / args.sig_l2**2 # vague prior on beta
# get new spectral norm
L = power_method(hess, dirty.shape, tol=args.pmtol, maxit=args.pmmaxit)
# deconvolve
eps = 1.0
i = 0
residual = dirty.copy()
model = np.zeros(dirty.shape, dtype=dirty.dtype)
for i in range(1, args.maxit):
# find point source candidates
if args.do_clean:
model_tmp = hogbom(mask[None] * residual / psf_max[:, None, None],
psf_array / psf_max[:, None, None],
gamma=args.cgamma,
pf=args.peak_factor)
phi.update_locs(np.any(model_tmp, axis=0))
# get new spectral norm
L = power_method(hess,
model.shape,
tol=args.pmtol,
maxit=args.pmmaxit)
else:
model_tmp = np.zeros_like(residual, dtype=residual.dtype)
# solve for beta updates
x = pcg(hess,
phi.hdot(residual),
phi.hdot(model_tmp),
M=lambda x: x * args.sig_l2**2,
tol=args.cgtol,
maxit=args.cgmaxit,
verbosity=args.cgverbose)
modelp = model.copy()
model += args.gamma * x
# impose sparsity and positivity in point sources
weights_21 = np.where(phi.mask, 1, 1e10) # 1e10 for effective infinity
model, dual = primal_dual(hess,
model,
modelp,
dual,
args.sig_21,
phi,
weights_21,
L,
tol=args.pdtol,
maxit=args.pdmaxit,
axis=0,
positivity=args.positivity,
report_freq=100)
# update Dirac dictionary (remove zero components)
phi.trim_fat(model)
# get residual
residual = R.make_residual(model) / psf_max_mean
# check stopping criteria
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(mask * residual_mfs).max()
rms = np.std(mask * residual_mfs)
eps = np.linalg.norm(model - modelp) / np.linalg.norm(model)
if i in report_iters:
# save current iteration
save_fits(args.outfile + str(i) + '_model.fits', model, hdr)
model_mfs = np.mean(model, axis=0)
save_fits(args.outfile + str(i) + '_model_mfs.fits', model_mfs,
hdr_mfs)
save_fits(args.outfile + str(i) + '_residual.fits',
residual / psf_max[:, None, None], hdr)
save_fits(args.outfile + str(i) + '_residual_mfs.fits',
residual_mfs, hdr_mfs)
print(
"At iteration %i peak of residual is %f, rms is %f, current eps is %f"
% (i, rmax, rms, eps))
if eps < args.tol:
print("We have convergence!")
break
# final iteration with only a positivity constraint on pixel locs
tmp = phi.hdot(model)
x = pcg(hess,
phi.hdot(residual),
np.zeros_like(tmp, dtype=tmp.dtype),
M=lambda x: x * args.sig_l2**2,
tol=args.cgtol,
maxit=args.cgmaxit,
verbosity=args.cgverbose)
modelp = model.copy()
model += args.gamma * x
model, dual = primal_dual(hess,
model,
modelp,
dual,
0.0,
phi,
weights_21,
L,
tol=args.pdtol,
maxit=args.pdmaxit,
axis=0,
report_freq=100)
# get residual
residual = R.make_residual(model) / psf_max_mean
# check stopping criteria
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(mask * residual_mfs).max()
rms = np.std(mask * residual_mfs)
print("At final iteration peak of residual is %f and rms is %f" %
(rmax, rms))
save_fits(args.outfile + '_model.fits', model, hdr)
model_mfs = np.mean(model, axis=0)
save_fits(args.outfile + '_model_mfs.fits', model_mfs, hdr_mfs)
save_fits(args.outfile + '_residual.fits',
residual / psf_max[:, None, None], hdr)
save_fits(args.outfile + '_residual_mfs.fits', residual_mfs, hdr_mfs)
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:
if args.interp_model:
R.write_component_model(comps, ref_freq, phi.mask, args.row_chunks,
args.chan_chunks)
else:
R.write_model(model)
def main(args):
# 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
print("Super resolution factor = ", cell_N / cell_rad)
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)
if args.nx is None or args.ny is None:
fov = args.fov * 3600
npix = int(fov / args.cell_size)
if npix % 2:
npix += 1
args.nx = npix
args.ny = npix
if args.nband is None:
args.nband = freq.size
print("Image size set to (%i, %i, %i)" % (args.nband, args.nx, args.ny))
# 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,
data_column=args.data_column,
weight_column=args.weight_column,
epsilon=args.epsilon,
imaging_weight_column=args.imaging_weight_column,
model_column=args.model_column,
flag_column=args.flag_column)
freq_out = R.freq_out
radec = R.radec
# 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,
2 * args.nx, 2 * args.ny, radec, freq_out)
hdr_psf_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600,
2 * args.nx, 2 * args.ny, radec, np.mean(freq_out))
# psf
if args.psf is not None:
try:
compare_headers(hdr_psf, fits.getheader(args.psf))
psf_array = load_fits(args.psf)
except:
psf_array = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf_array, hdr_psf)
else:
psf_array = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf_array, hdr_psf)
psf_max = np.amax(psf_array.reshape(args.nband, 4 * args.nx * args.ny),
axis=1)
wsum = np.sum(psf_max)
counts = np.sum(psf_max > 0)
psf_max_mean = wsum / counts # normalissation for more intuitive sig_21 values
psf_array /= psf_max_mean
psf = PSF(psf_array, args.nthreads)
psf_max = np.amax(psf_array.reshape(args.nband, 4 * args.nx * args.ny),
axis=1)
wsum = np.sum(psf_max)
psf_max[psf_max < 1e-15] = 1e-15 # LB - is this the right thing to do?
psf_mfs = np.sum(psf_array, axis=0) / wsum
save_fits(
args.outfile + '_psf_mfs.fits', psf_mfs[args.nx // 2:3 * args.nx // 2,
args.ny // 2:3 * args.ny // 2],
hdr_mfs)
# dirty
if args.dirty is not None:
try:
compare_headers(hdr, fits.getheader(args.dirty))
dirty = load_fits(args.dirty)
except:
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_mfs = np.sum(dirty / psf_max_mean, axis=0) / wsum
save_fits(args.outfile + '_dirty_mfs.fits', dirty_mfs, hdr_mfs)
residual = dirty.copy()
model = np.zeros((2, args.nband, args.nx, args.ny))
recompute_residual = False
if args.beta0 is not None:
compare_headers(hdr, fits.getheader(args.beta0))
model[0] = load_fits(args.beta0).squeeze()
recompute_residual = True
if args.alpha0 is not None:
compare_headers(hdr, fits.getheader(args.alpha0))
model[1] = load_fits(args.alpha0).squeeze()
recompute_residual = True
# normalise for more intuitive hypers
residual /= psf_max_mean
residual_mfs = np.sum(residual, axis=0) / wsum
save_fits(args.outfile + '_first_residual_mfs.fits', residual_mfs, hdr_mfs)
# mask
if args.mask is not None:
mask = load_fits(args.mask, dtype=np.int64)[None, :, :]
if mask.shape != (1, args.nx, args.ny):
raise ValueError("Mask has incorrect shape")
else:
mask = np.ones((1, args.nx, args.ny), dtype=np.int64)
# point mask
pmask = load_fits(args.point_mask, dtype=np.bool)[None, :, :]
if pmask.shape != (1, args.nx, args.ny):
raise ValueError("Mask has incorrect shape")
# set up splitting operator
phi = lambda x: x[0] * pmask + x[1] * mask
phih = lambda x: np.concatenate(
((pmask * x)[None], (mask * x)[None]), axis=0)
if recompute_residual:
image = phi(model)
residual = R.make_residual(image) / psf_max_mean
residual_mfs = np.sum(residual, axis=0) / wsum
# Gaussian "prior" used for preconditioning extended emission
A = Gauss(args.sig_l2a, args.nband, args.nx, args.ny, args.nthreads)
# preconditioning matrix
def hess(x):
return phih(psf.convolve(phi(x))) + np.concatenate(
(x[0:1] / args.sig_l2b**2, A.idot(x[1])[None]), axis=0)
# return phih(psf.convolve(phi(x))) + np.concatenate((x[0:1]/args.sig_l2b**2, x[1::]/args.sig_l2a**2), axis=0)
# M_func = lambda x: np.concatenate((x[0:1] * args.sig_l2b**2, x[1::] * args.sig_l2a**2), axis=0)
M_func = lambda x: np.concatenate(
(x[0:1] * args.sig_l2b**2, A.convolve(x[1])[None]), axis=0)
par_shape = phih(dirty).shape
if args.beta is None:
print("Getting spectral norm of update operator")
beta = power_method(hess,
par_shape,
tol=args.pmtol,
maxit=args.pmmaxit)
else:
beta = args.beta
print(" beta = %f " % beta)
# set up wavelet basis
theta = DaskTheta(args.nband, args.nx, args.ny, nthreads=args.nthreads)
nbasis = theta.nbasis
weights_21 = np.ones((theta.nbasis + 1, theta.nmax), dtype=np.float64)
tmp = np.pad(pmask.ravel(), (0, theta.nmax - args.nx * args.ny),
mode='constant')
weights_21[0] = np.where(tmp, args.sig_21b / args.sig_21a, 1e15)
dual = np.zeros((theta.nbasis + 1, args.nband, theta.nmax),
dtype=np.float64)
# Reporting
report_iters = list(np.arange(0, args.maxit, args.report_freq))
if report_iters[-1] != args.maxit - 1:
report_iters.append(args.maxit - 1)
# deconvolve
eps = 1.0
i = 0
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
print("Peak of initial residual is %f and rms is %f" % (rmax, rms))
for i in range(1, args.maxit):
x = pcg(hess,
phih(residual),
np.zeros(par_shape, dtype=np.float64),
M=M_func,
tol=args.cgtol,
maxit=args.cgmaxit,
verbosity=args.cgverbose)
if i in report_iters:
save_fits(args.outfile + str(i) + '_point_update.fits', x[0], hdr)
save_fits(args.outfile + str(i) + '_fluff_update.fits', x[1], hdr)
# update model
modelp = model
model = modelp + args.gamma * x
model, dual = primal_dual(hess,
model,
modelp,
dual,
args.sig_21a,
theta,
weights_21,
beta,
tol=args.pdtol,
maxit=args.pdmaxit,
report_freq=100,
mask=mask,
positivity=args.positivity,
gamma=args.gamma)
# get residual
image = phi(model)
residual = R.make_residual(image) / psf_max_mean
# check stopping criteria
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
eps = np.linalg.norm(model - modelp) / np.linalg.norm(model)
if i in report_iters:
# save current iteration
save_fits(args.outfile + str(i) + '_model.fits', image, hdr)
save_fits(args.outfile + str(i) + '_point.fits', model[0], hdr)
save_fits(args.outfile + str(i) + '_fluff.fits', model[1], hdr)
model_mfs = np.mean(image, axis=0)
save_fits(args.outfile + str(i) + '_model_mfs.fits', model_mfs,
hdr_mfs)
save_fits(args.outfile + str(i) + '_residual.fits', residual, hdr)
save_fits(args.outfile + str(i) + '_residual_mfs.fits',
residual_mfs, hdr_mfs)
print(
"At iteration %i peak of residual is %f, rms is %f, current eps is %f"
% (i, rmax, rms, eps))
if eps < args.tol:
break
if args.interp_model:
nband = args.nband
order = args.spectral_poly_order
mask = np.where(model_mfs > 1e-10, 1, 0)
I = np.argwhere(mask).squeeze()
Ix = I[:, 0]
Iy = I[:, 1]
npix = I.shape[0]
# get components
beta = image[:, 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:
if args.interp_model:
R.write_component_model(comps, ref_freq, mask, args.row_chunks,
args.chan_chunks)
else:
R.write_model(model)
def main(args):
if args.psf_pars is None:
print("Attempting to take psf_pars from residual fits header")
try:
rhdr = fits.getheader(args.residual)
except KeyError:
raise RuntimeError("Either provide a residual with beam "
"information or pass them in using --psf_pars "
"argument")
if 'BMAJ1' in rhdr.keys():
emaj = rhdr['BMAJ1']
emin = rhdr['BMIN1']
pa = rhdr['BPA1']
gaussparf = (emaj, emin, pa)
elif 'BMAJ' in rhdr.keys():
emaj = rhdr['BMAJ']
emin = rhdr['BMIN']
pa = rhdr['BPA']
gaussparf = (emaj, emin, pa)
else:
gaussparf = tuple(args.psf_pars)
if args.circ_psf:
e = (gaussparf[0] + gaussparf[1]) / 2.0
gaussparf[0] = e
gaussparf[1] = e
print("Using emaj = %3.2e, emin = %3.2e, PA = %3.2e \n" % gaussparf)
# load model image
model = load_fits(args.model, dtype=args.out_dtype)
model = model.squeeze()
orig_shape = model.shape
mhdr = fits.getheader(args.model)
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
elif mhdr["CTYPE3"].lower() == 'freq':
freq_axis = 3
else:
raise ValueError("Freq axis must be 3rd or 4th")
mfs_shape = list(orig_shape)
mfs_shape[0] = 1
mfs_shape = tuple(mfs_shape)
freqs, ref_freq = data_from_header(mhdr, axis=freq_axis)
nband = freqs.size
if nband < 2:
raise ValueError("Can't produce alpha map from a single band image")
npix_l = l_coord.size
npix_m = m_coord.size
# update cube psf-pars
for i in range(1, nband + 1):
mhdr['BMAJ' + str(i)] = gaussparf[0]
mhdr['BMIN' + str(i)] = gaussparf[1]
mhdr['BPA' + str(i)] = gaussparf[2]
if args.ref_freq is not None and args.ref_freq != ref_freq:
ref_freq = args.ref_freq
print(
'Provided reference frequency does not match that of fits file. Will overwrite.'
)
print("Cube frequencies:")
with np.printoptions(precision=2):
print(freqs)
print("Reference frequency is %3.2e Hz \n" % ref_freq)
# LB - new header for cubes if ref_freqs differ
new_hdr = set_header_info(mhdr, ref_freq, freq_axis, args, gaussparf)
# save next to model if no outfile is provided
if args.output_filename is None:
# strip .fits from model filename
tmp = args.model[::-1]
idx = tmp.find('.')
outfile = args.model[0:-(idx + 1)]
else:
outfile = args.output_filename
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)
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."
)
freqs_beam, _ = data_from_header(bhdr, axis=freq_axis)
if not np.array_equal(freqs, freqs_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.out_dtype).reshape(model.shape)
else:
beam_image = np.ones(model.shape, dtype=args.out_dtype)
# do beam correction LB - TODO: use forward model instead
beammin = np.amin(beam_image, axis=0)[None, :, :]
model = np.where(beammin >= args.pb_min, model / beam_image, 0.0)
if not args.dont_convolve:
print("Computing clean beam")
# convolve model to desired resolution
model, gausskern = convolve2gaussres(model, xx, yy, gaussparf,
args.ncpu, None,
args.padding_frac)
# save clean beam
if 'c' in args.products:
name = outfile + '.clean_psf.fits'
save_fits(name,
gausskern.reshape(mfs_shape),
new_hdr,
dtype=args.out_dtype)
print("Wrote clean psf to %s \n" % name)
# save convolved model
if 'm' in args.products:
name = outfile + '.convolved_model.fits'
save_fits(name,
model.reshape(orig_shape),
new_hdr,
dtype=args.out_dtype)
print("Wrote convolved model to %s \n" % name)
# add in residuals and set threshold
if args.residual is not None:
resid = load_fits(args.residual, dtype=args.out_dtype).squeeze()
rhdr = fits.getheader(args.residual)
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 of residual do not match those of model")
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 of residual do not match those of model")
freqs_res, _ = data_from_header(rhdr, axis=freq_axis)
if not np.array_equal(freqs, freqs_res):
raise ValueError("Freqs of residual do not match those of model")
# convolve residual to same resolution as model
gausspari = ()
for i in range(1, nband + 1):
key = 'BMAJ' + str(i)
if key in rhdr.keys():
emaj = rhdr[key]
emin = rhdr['BMIN' + str(i)]
pa = rhdr['BPA' + str(i)]
gausspari += ((emaj, emin, pa), )
else:
print(
"Can't find Gausspars in residual header, unable to add residuals back in"
)
gausspari = None
break
if gausspari is not None and args.add_convolved_residuals:
resid, _ = convolve2gaussres(resid,
xx,
yy,
gaussparf,
args.ncpu,
gausspari,
args.padding_frac,
norm_kernel=True)
model += resid
print("Convolved residuals added to convolved model")
if 'c' in args.products:
name = outfile + '.convolved_residual.fits'
save_fits(name, resid.reshape(orig_shape), rhdr)
print("Wrote convolved residuals to %s" % name)
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
print("Setting cutoff threshold as %i times the rms "
"of the residual " % args.threshold)
del resid
else:
print("No residual provided. Setting threshold i.t.o dynamic range. "
"Max dynamic range is %i " % args.maxDR)
threshold = model.max() / args.maxDR
rms_cube = None
print("Threshold set to %f Jy. \n" % threshold)
# get pixels above threshold
minimage = np.amin(model, axis=0)
maskindices = 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
# set weights for fit
if rms_cube is not None:
print("Using RMS in each imaging band to determine weights. \n")
weights = np.where(rms_cube > 0, 1.0 / rms_cube**2, 0.0)
# normalise
weights /= weights.max()
else:
if args.channel_weights is not None:
weights = np.array(args.channel_weights)
print("Using provided channel weights \n")
else:
print(
"No residual or channel weights provided. Using equal weights. \n"
)
weights = np.ones(nband, dtype=np.float64)
ncomps, _ = fitcube.shape
fitcube = da.from_array(fitcube.astype(np.float64),
chunks=(ncomps // args.ncpu, 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)
alpha, alpha_err, Iref, i0_err = fit_spi_components(
fitcube, weights, freqsdask, np.float64(ref_freq)).compute()
print("Done. Writing output. \n")
alphamap = np.zeros(model[0].shape, dtype=model.dtype)
alpha_err_map = np.zeros(model[0].shape, dtype=model.dtype)
i0map = np.zeros(model[0].shape, dtype=model.dtype)
i0_err_map = np.zeros(model[0].shape, dtype=model.dtype)
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 = outfile + '.Irec_cube.fits'
save_fits(name,
Irec_cube.reshape(orig_shape),
mhdr,
dtype=args.out_dtype)
print("Wrote reconstructed cube to %s" % name)
# save alpha map
if 'a' in args.products:
name = outfile + '.alpha.fits'
save_fits(name,
alphamap.reshape(mfs_shape),
mhdr,
dtype=args.out_dtype)
print("Wrote alpha map to %s" % name)
# save alpha error map
if 'e' in args.products:
name = outfile + '.alpha_err.fits'
save_fits(name,
alpha_err_map.reshape(mfs_shape),
mhdr,
dtype=args.out_dtype)
print("Wrote alpha error map to %s" % name)
# save I0 map
if 'i' in args.products:
name = outfile + '.I0.fits'
save_fits(name, i0map.reshape(mfs_shape), mhdr, dtype=args.out_dtype)
print("Wrote I0 map to %s" % name)
# save I0 error map
if 'k' in args.products:
name = outfile + '.I0_err.fits'
save_fits(name,
i0_err_map.reshape(mfs_shape),
mhdr,
dtype=args.out_dtype)
print("Wrote I0 error map to %s" % name)
print(' \n ')
print("All done here")
def main(args):
# read coords from fits file
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)
nchan = freqs.size
gausspari = ()
if freqs.size > 1:
for i in range(1, nchan + 1):
key = 'BMAJ' + str(i)
if key in hdr.keys():
emaj = hdr[key]
emin = hdr['BMIN' + str(i)]
pa = hdr['BPA' + str(i)]
gausspari += ((emaj, emin, pa), )
else:
if 'BMAJ' in hdr.keys():
emaj = hdr['BMAJ']
emin = hdr['BMIN']
pa = hdr['BPA']
# using key of 1 for consistency with fits standard
gausspari = ((emaj, emin, pa), )
if len(gausspari) == 0 and args.psf_pars is None:
raise ValueError("No psf parameters in fits file and none passed in.")
if len(gausspari) == 0:
print(
"No psf parameters in fits file. Convolving model to resolution specified by psf-pars."
)
gaussparf = tuple(args.psf_pars)
else:
if args.psf_pars is None:
gaussparf = gausspari[0]
else:
gaussparf = tuple(args.psf_pars)
if args.circ_psf:
e = (gaussparf[0] + gaussparf[1]) / 2.0
gaussparf[0] = e
gaussparf[1] = e
print("Using emaj = %3.2e, emin = %3.2e, PA = %3.2e \n" % gaussparf)
# update header
if freqs.size > 1:
for i in range(1, nchan + 1):
hdr['BMAJ' + str(i)] = gaussparf[0]
hdr['BMIN' + str(i)] = gaussparf[1]
hdr['BPA' + str(i)] = gaussparf[2]
else:
hdr['BMAJ'] = gaussparf[0]
hdr['BMIN'] = gaussparf[1]
hdr['BPA'] = gaussparf[2]
# coodinate grid
xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij')
# convolve image
imagei = load_fits(args.image, dtype=np.float32).squeeze()
print(imagei.shape)
image, gausskern = convolve2gaussres(imagei, xx, yy, gaussparf, args.ncpu,
gausspari, args.padding_frac)
# load beam and correct
if args.beam_model is not None:
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."
)
freqs_beam, _ = data_from_header(bhdr, axis=freq_axis)
if not np.array_equal(freqs, freqs_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=np.float32).squeeze()
image = np.where(beam_image >= args.pb_min, image / beam_image, 0.0)
# save next to model if no outfile is provided
if args.output_filename is None:
# strip .fits from model filename
tmp = args.model[::-1]
idx = tmp.find('.')
outfile = args.model[0:-idx]
else:
outfile = args.output_filename
# save images
name = outfile + '.clean_psf.fits'
save_fits(name, gausskern, hdr)
print("Wrote clean psf to %s \n" % name)
name = outfile + '.convolved.fits'
save_fits(name, image, hdr)
print("Wrote convolved model to %s \n" % name)
print("All done here")
def main(args):
# 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
print("Super resolution factor = ", cell_N / cell_rad)
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)
if args.nx is None or args.ny is None:
fov = args.fov * 3600
npix = int(fov / args.cell_size)
if npix % 2:
npix += 1
args.nx = npix
args.ny = npix
if args.nband is None:
args.nband = freq.size
print("Image size set to (%i, %i, %i)" % (args.nband, args.nx, args.ny))
# 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,
data_column=args.data_column,
weight_column=args.weight_column,
epsilon=args.epsilon,
imaging_weight_column=args.imaging_weight_column,
model_column=args.model_column,
flag_column=args.flag_column)
freq_out = R.freq_out
radec = R.radec
# 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,
2 * args.nx, 2 * args.ny, radec, freq_out)
hdr_psf_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600,
2 * args.nx, 2 * args.ny, radec, np.mean(freq_out))
# psf
if args.psf is not None:
try:
compare_headers(hdr_psf, fits.getheader(args.psf))
psf_array = load_fits(args.psf)
except:
psf_array = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf_array, hdr_psf)
else:
psf_array = R.make_psf()
save_fits(args.outfile + '_psf.fits', psf_array, hdr_psf)
psf_max = np.amax(psf_array.reshape(args.nband, 4 * args.nx * args.ny),
axis=1)
wsum = np.sum(psf_max)
counts = np.sum(psf_max > 0)
psf_max_mean = wsum / counts # normalissation for more intuitive sig_21 values
psf_array /= psf_max_mean
psf = PSF(psf_array, args.nthreads)
psf_max = np.amax(psf_array.reshape(args.nband, 4 * args.nx * args.ny),
axis=1)
wsum = np.sum(psf_max)
psf_max[psf_max < 1e-15] = 1e-15 # LB - is this the right thing to do?
psf_mfs = np.sum(psf_array, axis=0) / wsum
save_fits(
args.outfile + '_psf_mfs.fits', psf_mfs[args.nx // 2:3 * args.nx // 2,
args.ny // 2:3 * args.ny // 2],
hdr_mfs)
# dirty
if args.dirty is not None:
try:
compare_headers(hdr, fits.getheader(args.dirty))
dirty = load_fits(args.dirty)
except:
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_mfs = np.sum(dirty / psf_max_mean, axis=0) / wsum
save_fits(args.outfile + '_dirty_mfs.fits', dirty_mfs, hdr_mfs)
if args.x0 is not None:
try:
compare_headers(hdr, fits.getheader(args.x0))
model = load_fits(args.x0, dtype=np.float64)
if args.first_residual is not None:
try:
compare_headers(hdr, fits.getheader(args.first_residual))
residual = load_fits(args.first_residual, dtype=np.float64)
except:
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)
except:
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()
# normalise for more intuitive hypers
residual /= psf_max_mean
residual_mfs = np.sum(residual, axis=0) / wsum
save_fits(args.outfile + '_first_residual_mfs.fits', residual_mfs, hdr_mfs)
# mask
if args.mask is not None:
mask = load_fits(args.mask, dtype=np.int64)[None, :, :]
if mask.shape != (1, args.nx, args.ny):
raise ValueError("Mask has incorrect shape")
else:
mask = np.ones((1, args.nx, args.ny), dtype=np.int64)
# preconditioning matrix
def hess(x):
return mask * psf.convolve(mask * x) + x / args.sig_l2**2
if args.beta is None:
print("Getting spectral norm of update operator")
beta = power_method(hess,
dirty.shape,
tol=args.pmtol,
maxit=args.pmmaxit)
else:
beta = args.beta
print(" beta = %f " % beta)
# set up wavelet basis
if args.psi_basis is None:
print("Using Dirac + db1-4 dictionary")
psi = DaskPSI(args.nband,
args.nx,
args.ny,
nlevels=args.psi_levels,
nthreads=args.nthreads)
# psi = PSI(args.nband, args.nx, args.ny, nlevels=args.psi_levels)
else:
if not isinstance(args.psi_basis, list):
args.psi_basis = list(args.psi_basis)
print("Using ", args.psi_basis, " dictionary")
psi = DaskPSI(args.nband,
args.nx,
args.ny,
nlevels=args.psi_levels,
nthreads=args.nthreads,
bases=args.psi_basis)
# psi = PSI(args.nband, args.nx, args.ny, nlevels=args.psi_levels, bases=args.psi_basis)
nbasis = psi.nbasis
weights_21 = np.ones((psi.nbasis, psi.nmax), dtype=np.float64)
dual = np.zeros((psi.nbasis, args.nband, psi.nmax), dtype=np.float64)
# Reweighting
if args.reweight_iters is not None:
if not isinstance(args.reweight_iters, list):
reweight_iters = [args.reweight_iters]
else:
reweight_iters = list(args.reweight_iters)
else:
reweight_iters = list(
np.arange(args.reweight_start, args.reweight_end,
args.reweight_freq))
reweight_iters.append(args.reweight_end)
# Reporting
report_iters = list(np.arange(0, args.maxit, args.report_freq))
if report_iters[-1] != args.maxit - 1:
report_iters.append(args.maxit - 1)
# deconvolve
eps = 1.0
i = 0
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
M = lambda x: x * args.sig_l2**2 # preconditioner
print("Peak of initial residual is %f and rms is %f" % (rmax, rms))
for i in range(1, args.maxit):
x = pcg(hess,
mask * residual,
np.zeros(dirty.shape, dtype=np.float64),
M=M,
tol=args.cgtol,
maxit=args.cgmaxit,
minit=args.cgminit,
verbosity=args.cgverbose)
if i in report_iters:
save_fits(args.outfile + str(i) + '_update.fits', x, hdr)
# update model
modelp = model
model = modelp + args.gamma * x
model, dual = primal_dual(hess,
model,
modelp,
dual,
args.sig_21,
psi,
weights_21,
beta,
tol=args.pdtol,
maxit=args.pdmaxit,
report_freq=100,
mask=mask,
positivity=args.positivity)
# reweighting
if i in reweight_iters:
v = psi.hdot(model)
l2_norm = norm(v, axis=1)
l2_norm = np.where(l2_norm < args.sig_21 * weights_21, 0.0,
l2_norm)
for m in range(psi.nbasis):
indnz = l2_norm[m].nonzero()
alpha = np.percentile(l2_norm[m, indnz].flatten(),
args.reweight_alpha_percent)
alpha = np.maximum(alpha, args.reweight_alpha_min)
print("Reweighting - ", m, alpha)
weights_21[m] = alpha / (l2_norm[m] + alpha)
args.reweight_alpha_percent *= args.reweight_alpha_ff
# print(" reweight alpha percent = ", args.reweight_alpha_percent)
# get residual
residual = R.make_residual(model) / psf_max_mean
# check stopping criteria
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
eps = np.linalg.norm(model - modelp) / np.linalg.norm(model)
if i in report_iters:
# save current iteration
save_fits(args.outfile + str(i) + '_model.fits', model, hdr)
model_mfs = np.mean(model, axis=0)
save_fits(args.outfile + str(i) + '_model_mfs.fits', model_mfs,
hdr_mfs)
save_fits(args.outfile + str(i) + '_residual.fits', residual, hdr)
save_fits(args.outfile + str(i) + '_residual_mfs.fits',
residual_mfs, hdr_mfs)
print(
"At iteration %i peak of residual is %f, rms is %f, current eps is %f"
% (i, rmax, rms, eps))
if args.write_model:
R.write_model(model)
if args.make_restored:
x = pcg(hess,
residual,
np.zeros(dirty.shape, dtype=np.float64),
M=M,
tol=args.cgtol,
maxit=args.cgmaxit)
restored = model + x
# get residual
residual = R.make_residual(restored) / psf_max_mean
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
print("After restoring peak of residual is %f and rms is %f" %
(rmax, rms))
# save current iteration
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)
save_fits(args.outfile + '_restored_residual.fits', residual, hdr)
save_fits(args.outfile + '_restored_residual_mfs.fits', residual_mfs,
hdr_mfs)
def main(args):
# load dirty and psf
dirty = load_fits(args.dirty)
real_type = dirty.dtype
hdr = fits.getheader(args.dirty)
freq = data_from_header(hdr, axis=3)
l_coord = data_from_header(hdr, axis=1)
m_coord = data_from_header(hdr, axis=2)
nband, nx, ny = dirty.shape
psf_array = load_fits(args.psf)
hdr_psf = fits.getheader(args.psf)
try:
assert np.array_equal(freq, data_from_header(hdr_psf, axis=3))
except:
raise ValueError("Fits frequency axes dont match")
print("Image size is (%i, %i, %i)" % (nband, nx, ny))
psf_max = np.amax(psf_array.reshape(nband, 4 * nx * ny), axis=1)
wsum = np.sum(psf_max)
psf_max[psf_max < 1e-15] = 1e-15
dirty_mfs = np.sum(dirty, axis=0) / wsum
rmax = np.abs(dirty_mfs).max()
rms = np.std(dirty_mfs)
print("Peak of dirty is %f and rms is %f" % (rmax, rms))
psf_mfs = np.sum(psf_array, axis=0) / wsum
# set operators
psf = PSF(psf_array, args.ncpu, sigma0=args.sig_l2)
K = Prior(args.sig_l2, nband, nx, ny, nthreads=args.ncpu)
# def hess(x):
# return psf.convolve(x) + K.iconvolve(x)
# get Lipschitz constant
if args.beta is None:
from pfb.opt import power_method
beta = power_method(psf.hess,
dirty.shape,
tol=args.pmtol,
maxit=args.pmmaxit)
else:
beta = args.beta
print("beta = ", beta)
# Reweighting
if args.reweight_iters is not None:
reweight_iters = args.reweight_iters
else:
reweight_iters = list(
np.arange(args.reweight_start, args.reweight_end,
args.reweight_freq))
reweight_iters.append(args.reweight_end)
# Reporting
report_iters = list(np.arange(0, args.maxit, args.report_freq))
if report_iters[-1] != args.maxit - 1:
report_iters.append(args.maxit - 1)
# set up wavelet basis
if args.use_psi:
nband, nx, ny = dirty.shape
psi = PSI(nband, nx, ny, nlevels=args.psi_levels)
nbasis = psi.nbasis
weights_21 = np.ones((psi.nbasis, psi.nmax), dtype=real_type)
else:
psi = None
weights_21 = np.ones(nx * ny, dtype=real_type)
# initalise model
if args.x0 is None:
model = np.zeros(dirty.shape, dtype=real_type)
dual = np.zeros((psi.nbasis, nband, psi.nmax), dtype=real_type)
residual = dirty
else:
compare_headers(hdr, fits.getheader(args.x0))
model = load_fits(args.x0).astype(real_type)
dual = np.zeros((psi.nbasis, nband, psi.nmax), dtype=real_type)
residual = dirty - psf.convolve(model)
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
print("At iteration 0 peak of residual is %f and rms is %f" % (rmax, rms))
# deconvolve
for k in range(args.maxit):
x = pcg(psf.hess,
residual,
np.zeros(dirty.shape, dtype=real_type),
M=K.dot,
tol=args.cgtol,
maxit=args.cgmaxit,
verbosity=args.cgverbose)
modelp = model.copy()
model = modelp + args.gamma * x
if args.use_psi:
model, dual = simple_pd(psf.hess,
model,
modelp,
dual,
args.sig_21,
psi,
weights_21,
beta,
tol=args.pdtol,
maxit=args.pdmaxit,
report_freq=10)
else:
model = prox_21(model,
args.sig_21,
weights_21,
psi=psi,
positivity=True)
# convergence check
normx = norm(model)
if np.isnan(normx) or normx == 0.0:
normx = 1.0
eps = norm(model - modelp) / normx
if eps < args.tol:
break
# reweighting
if k in reweight_iters:
if psi is None:
l2norm = norm(model.reshape(nband, npix), axis=0)
weights_21 = 1.0 / (l2norm + alpha)
else:
v = psi.hdot(model)
l2_norm = norm(v, axis=1)
for m in range(psi.nbasis):
indnz = l2_norm[m].nonzero()
alpha = np.percentile(l2_norm[m, indnz].flatten(),
args.reweight_alpha_percent)
alpha = np.maximum(alpha, args.reweight_alpha_min)
# alpha = args.reweight_alpha_min
print("Reweighting - ", m, alpha)
weights_21[m] = 1.0 / (l2_norm[m] + alpha)
args.reweight_alpha_percent *= args.reweight_alpha_ff
print(" reweight alpha percent = ",
args.reweight_alpha_percent)
# get residual
residual = dirty - psf.convolve(model)
# check stopping criteria
residual_mfs = np.sum(residual, axis=0) / wsum
rmax = np.abs(residual_mfs).max()
rms = np.std(residual_mfs)
# reporting
if k in report_iters:
save_fits(args.outfile + str(k + 1) + '_model.fits',
model,
hdr,
dtype=real_type)
model_mfs = np.mean(model, axis=0)
save_fits(args.outfile + str(k + 1) + '_model_mfs.fits', model_mfs,
hdr)
save_fits(args.outfile + str(k + 1) + '_update.fits', x, hdr)
save_fits(args.outfile + str(k + 1) + '_residual.fits',
residual,
hdr,
dtype=real_type)
save_fits(args.outfile + str(k + 1) + '_residual_mfs.fits',
residual_mfs, hdr)
print(
"At iteration %i peak of residual is %f, rms is %f, current eps is %f"
% (k + 1, rmax, rms, eps))
# save final results
save_fits(args.outfile + '_model.fits', model, hdr, dtype=real_type)
residual = dirty - psf.convolve(model)
save_fits(args.outfile + '_residual.fits',
residual / psf_max[:, None, None], hdr)