# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('DIRTY')
@cli.command()
@click.option('-ms', '--ms', required=True,
help='Path to measurement set.')
@click.option('-mc', '--mueller-column', default='CORRECTED_DATA',
help="Column to write Mueller term to.")
@click.option('-gt', '--gain-table', required=True,
help="QuartiCal gain table at the same resolution as MS")
@click.option('-acol', '--acol',
help='Will apply gains to this column if supplied')
@click.option('-c2', '--compareto',
help="Will compare corrupted vis to this column if provided. "
"Mainly useful for testing.")
@click.option('-o', '--output-filename', type=str, required=True,
help="Basename of output.")
@click.option('-ha', '--host-address',
help='Address where the distributed client lives. '
'Will use a local cluster if no address is provided')
@click.option('-nw', '--nworkers', type=int, default=1,
help='Number of workers for the client.')
@click.option('-ntpw', '--nthreads-per-worker', type=int, default=1,
help='Number of dask threads per worker.')
import numpy as np
from functools import partial
import numexpr as ne
import dask
import dask.array as da
from daskms import xds_from_ms, xds_from_table, xds_to_table
import psutil
from africanus.gridding.wgridder.dask import dirty as vis2im
from africanus.gridding.wgridder.dask import model as im2vis
from africanus.gridding.wgridder.dask import residual as im2residim
from africanus.gridding.wgridder.dask import hessian
from africanus.gridding.wgridder.hessian import hessian as hessian_np
import pyscilog
log = pyscilog.get_logger('GRID')
class Gridder(object):
def __init__(self,
ms_name,
nx,
ny,
cell_size,
nband=None,
nthreads=8,
do_wstacking=1,
Stokes='I',
row_chunks=-1,
chan_chunks=32,
optimise_chunks=True,
epsilon=1e-5,
psf_oversize=2.0,
# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('PREDICT')
@cli.command()
@click.argument('ms', nargs=-1)
@click.option('-m', '--model', required=True, help='Path to model.fits')
@click.option('-mc',
'--model-column',
default='MODEL_DATA',
help="Model column to write visibilities to."
"Must be the same across MSs")
@click.option('-rchunk', '--row-chunks', help="Number of rows in a chunk.")
@click.option('-cchunk',
'--chan-chunks',
help="Number of channels in a chunk.")
@click.option('-eps',
'--epsilon',
type=float,
default=1e-5,
help='Gridder accuracy')
@click.option('--wstack/--no-wstack', default=True)
@click.option('--mock/--no-mock', default=False)
@click.option('--double-accum/--no-double-accum', default=True)
@click.option('-o',
# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('PSF')
@cli.command()
@click.option('-ms', '--ms', required=True, help='Path to measurement set.')
@click.option('-dc',
'--data-column',
help="Data column to image."
"Must be the same across MSs")
@click.option('-wc',
'--weight-column',
help="Column containing natural weights."
"Must be the same across MSs")
@click.option('-iwc',
'--imaging-weight-column',
help="Column containing/to write imaging weights to."
"Must be the same across MSs")
@click.option('-eps',
'--epsilon',
type=float,
default=1e-5,
help='Gridder accuracy')
@click.option('--wstack/--no-wstack', default=True)
@click.option('--double-accum/--no-double-accum', default=True)
import numpy as np
import pyscilog
log = pyscilog.get_logger('PD')
def primal_dual(
A,
xbar,
x0,
v0, # initial guess for primal and dual variables
lam, # regulariser strength,
psi, # linear operator in dual domain
weights, # weights for l1 thresholding
L,
prox, # prox of regulariser
nu=1.0, # spectral norms
sigma=None, # step size of dual update
mask=None, # regions where mask is False will be masked
tol=1e-5,
maxit=1000,
positivity=True,
report_freq=10,
axis=1,
gamma=1.0,
verbosity=1):
# initialise
x = x0.copy()
v = v0.copy()
# gradient function
import numpy as np
from scipy.linalg import norm
import pyscilog
log = pyscilog.get_logger('PM')
def power_method(A,
imsize,
b0=None,
tol=1e-5,
maxit=250,
verbosity=1,
report_freq=25):
if b0 is None:
b = np.random.randn(*imsize)
b /= norm(b)
else:
b = b0 / norm(b0)
beta = 1.0
eps = 1.0
k = 0
while eps > tol and k < maxit:
bp = b
b = A(bp)
bnorm = np.linalg.norm(b)
betap = beta
beta = np.vdot(bp, b) / np.vdot(bp, bp)
b /= bnorm
eps = np.linalg.norm(beta - betap) / betap
k += 1
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# flake8: noqa
import sys
import os
from pfb import set_threads
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('PFB')
def main():
_main(dest=log)
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
# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('RESTORE')
@cli.command()
@click.option('-m', '--model', required=True, help="Path to model image cube")
@click.option('-r',
'--residual',
required=True,
help="Path to residual image cube")
@click.option('-p', '--psf', required=True, help="Path to PSF cube")
@click.option('-beam',
'--beam',
required=False,
help="Path to power beam image cube")
@click.option('-o',
'--output-filename',
required=True,
help="Basename of output.")
@click.option('-nthreads',
'--nthreads',
type=int,
default=1,
show_default=True,
help='Number of threads to use')
import numpy as np
import pyscilog
log = pyscilog.get_logger('FISTA')
def back_track_func(x, xp, gradp, likp, L):
df = x - xp
return likp + np.vdot(gradp, df) + L * np.vdot(df, df) / 2
def fista(x0,
L, # spectral norm of measurement operator
fprime, # function returning value and gradient
prox, # function implementing prox of regulariser
tol=1e-3,
maxit=100,
report_freq=50,
verbosity=1):
# start iterations
t = 1.0
x = x0.copy()
y = x0.copy()
eps = 1.0
k = 0
fidp, gradp = fprime(x)
for k in range(maxit):
xp = x.copy()
# gradient update
x = y - gradp / L
# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
from functools import partial
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('FORWARD')
@cli.command()
@click.option('-r', '--residual', required=True,
help="Path to residual.fits")
@click.option('-p', '--psf', required=True,
help="Path to PSF.fits")
@click.option('-mask', '--mask',
help="Path to mask.fits.")
@click.option('-bm', '--beam-model',
help="Path to beam_model.fits or JimBeam.")
@click.option('-band', '--band', default='L',
help='L or UHF band when using JimBeam.')
@click.option('-wt', '--weight-table',
help="Path to weight table produced by psf worker")
@click.option('-o', '--output-filename', type=str, required=True,
help="Basename of output.")
@click.option('-nb', '--nband', type=int, required=True,
help="Number of imaging bands")
@click.option('-otype', '--output-type', default='f4',
help="Data type of output")
@click.option('-eps', '--epsilon', type=float, default=1e-5,
help='Gridder accuracy')
import numpy as np
import numexpr as ne
from pfb.utils.misc import give_edges
import pyscilog
log = pyscilog.get_logger('HOGBOM')
def hogbom(ID,
PSF,
gamma=0.1,
pf=0.1,
maxit=10000,
report_freq=1000,
verbosity=1):
nband, nx, ny = ID.shape
_, nx_psf, ny_psf = PSF.shape
nx0 = nx_psf // 2
ny0 = ny_psf // 2
x = np.zeros((nband, nx, ny), dtype=ID.dtype)
IR = ID.copy()
IRsearch = np.sum(IR, axis=0)**2
pq = IRsearch.argmax()
p = pq // ny
q = pq - p * ny
IRmax = np.sqrt(IRsearch[p, q])
_, nx_psf, ny_psf = PSF.shape
wsums = np.amax(PSF.reshape(-1, nx_psf * ny_psf), axis=1)
tol = pf * IRmax
k = 0
stall_count = 0
while IRmax > tol and k < maxit and stall_count < 5:
# flake8: noqa
from pfb.workers.main import cli
from functools import partial
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('NNLS')
@cli.command()
@click.option('-d', '--dirty', help="Path to dirty.")
@click.option('-p', '--psf', help="Path to PSF")
@click.option('-x0', '--x0', help="Initial model")
@click.option('-mv',
'--min-value',
type=float,
default=0.0,
help="Minimum value below which to threshold")
@click.option('-o',
'--output-filename',
type=str,
required=True,
help="Basename of output.")
@click.option('-nthreads',
'--nthreads',
type=int,
default=0,
help="Total number of threads to use per worker")
@click.option('-ftol',
# flake8: noqa
from contextlib import ExitStack
import click
from omegaconf import OmegaConf
from pfb.workers.main import cli
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('SPIFIT')
@cli.command()
@click.option('-image', '--image', required=True,
help="Path to model or restored image cube.")
@click.option('-resid', "--residual", required=False,
help="Path to residual image cube.")
@click.option('-o', '--output-filename', required=True,
help="Path to output directory + prefix.")
@click.option('-pp', '--psf-pars', nargs=3, type=float,
help="Beam parameters matching FWHM of restoring beam "
"specified as emaj emin pa."
"By default these are taken from the fits header "
"of the residual image.")
@click.option('--circ-psf/--no-circ-psf', default=False)
@click.option('-th', '--threshold', default=10, type=float, show_default=True,
help="Multiple of the rms in the residual to threshold on."
"Only components above threshold*rms will be fit.")
@click.option('-maxdr', '--maxdr', default=100, type=float, show_default=True,
help="Maximum dynamic range used to determine the "
"threshold above which components need to be fit. "
"Only used if residual is not passed in.")
@click.option('-bw', '--band-weights', type=float,
help="Per bands weights to use during the fit")
# flake8: noqa
import os
from contextlib import ExitStack
import click
from omegaconf import OmegaConf
from pfb.workers.main import cli
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('BINTERP')
@cli.command()
@click.option('-image',
'--image',
required=True,
help="Path to model or restored image cube.")
@click.option('-o',
'--output-dir',
help='Output directory. Placed next to -image if not provided')
@click.option('-postfix',
'--postfix',
default='beam.fits',
help="Postfix to append to -image for writing out beams.")
@click.option('-ms',
'--ms',
help="Will interpolate beam onto TIME if provided.")
@click.option('-bm',
'--beam-model',
default=None,
required=True,
help="Fits beam model to use. \n"
import numpy as np
from scipy.stats import expon
from pfb.opt.power_method import power_method
from pfb.opt.pcg import pcg
from pfb.opt.primal_dual import primal_dual
from pfb.operators.psi import DaskPSI
from pfb.operators.psf import PSF
from pfb.prox.prox_21 import prox_21
from pfb.utils.fits import save_fits
from pfb.utils.misc import Gaussian2D
import pyscilog
log = pyscilog.get_logger('SARA')
def resid_func(x, dirty, hessian, mask, beam, wsum):
"""
Returns the unattenuated residual
"""
residual = dirty - hessian(mask(beam(x))) / wsum
residual_mfs = np.sum(residual, axis=0)
residual = residual
return residual, residual_mfs
def sara(psf,
model,
residual,
mask=None,
beam_image=None,
hessian=None,
wsum=1,
import numpy as np
import scipy
from pfb.operators.psf import PSF
from pfb.operators.dirac import Dirac
from pfb.opt.primal_dual import primal_dual
from pfb.opt.pcg import pcg
from pfb.opt.power_method import power_method
from pfb.opt.hogbom import hogbom
from pfb.prox.prox_21m import prox_21m
from pfb.utils.fits import save_fits
from skimage.filters import threshold_mean
import pyscilog
log = pyscilog.get_logger('SPOTLESS')
def make_noise_map(restored_image, boxsize):
# Modified version of Cyril's magic minimum filter
# Plundered from the depths of
# https://github.com/cyriltasse/DDFacet/blob/master/SkyModel/MakeMask.py
box = (boxsize, boxsize)
n = boxsize**2.0
x = np.linspace(-10, 10, 1000)
f = 0.5 * (1.0 + scipy.special.erf(x / np.sqrt(2.0)))
F = 1.0 - (1.0 - f)**n
ratio = np.abs(np.interp(0.5, F, x))
noise = -scipy.ndimage.filters.minimum_filter(restored_image, box) / ratio
negative_mask = noise < 0.0
noise[negative_mask] = 1.0e-10
median_noise = np.median(noise)
median_mask = noise < median_noise
noise[median_mask] = median_noise
import numpy as np
from pfb.operators.psf import PSF
from pfb.opt.hogbom import hogbom
import pyscilog
log = pyscilog.get_logger('CLEAN')
def resid_func(x, dirty, psfo):
residual = dirty - psfo.convolve(x)
return residual
def clean(psf,
model,
residual,
mask=None,
beam=None,
nthreads=0,
maxit=10,
gamma=1.0,
peak_factor=0.01,
threshold=0.0,
hbgamma=0.1,
hbpf=0.1,
hbmaxit=5000,
hbverbose=1): # Hogbom options
if len(residual.shape) > 3:
raise ValueError("Residual must have shape (nband, nx, ny)")
nband, nx, ny = residual.shape
# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('RESIDUAL')
@cli.command()
@click.argument('ms', nargs=-1)
@click.option('-dc',
'--data-column',
help="Data column to image."
"Must be the same across MSs")
@click.option('-wc',
'--weight-column',
help="Column containing natural weights."
"Must be the same across MSs")
@click.option('-iwc',
'--imaging-weight-column',
help="Column containing imaging weights. "
"Must be the same across MSs")
@click.option('-fc',
'--flag-column',
default='FLAG',
help="Column containing data flags."
"Must be the same across MSs")
@click.option('-muc',