# 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',