Deconstructing an imager

[landmanbester]

Workers

The idea is to create a flexible imaging framework where different steps (eg. gridding, deconvolution or image restoration etc.) are defined as workers and strung together using stimela recipes. pfb-clean uses click to create a cli for each worker so that they can also be run individually. The main (by which I mean sort of mature enough for testing) worker functions are defined under the command group pfb and can be listed using

$ pfb --help
Usage: pfb [OPTIONS] COMMAND [ARGS]...

Options:
  --help  Show this message and exit.

Commands:
  clean    Single-scale clean.
  degrid   Predict model visibilities to measurement sets.
  fwdbwd   Forward backward steps
  grid     Compute imaging weights and create a dirty image, psf from xds.
  init     Create a dirty image, psf and weights from a list of...
  restore  Create restored images.

Help for each worker (eg. the clean worker) can be found by invoking the help for that worker

$ pfb clean --help
Options:
  -eps, --epsilon FLOAT           Gridder accuracy  [default: 1e-07]
  --wstack / --no-wstack          Perform w-correction via improved w-stacking
                                  [default: wstack]
  --double-accum / --no-double-accum
                                  Accumulate onto grid using double precision.
                                  Only has an affect when using single
                                  precision.  [default: double-accum]
  -ha, --host-address TEXT        Address where the distributed client lives.
                                  Will use a local cluster if no address is
                                  provided and scheduler is set to
                                  distributed.  [default: <UNSET>]
  -nw, --nworkers INTEGER         Number of worker processes.  [default:
                                  <UNSET>]
  -ntpw, --nthreads-per-worker INTEGER
                                  Number of threads per worker process.
                                  [default: <UNSET>]
  -nvthreads, --nvthreads INTEGER
                                  Number of threads used to scale vertically
                                  (eg. for FFTs).  [default: <UNSET>]
  -nthreads, --nthreads INTEGER   Total number of threads to use between all
                                  processes. Will use all available resources
                                  by default  [default: <UNSET>]
  -scheduler, --scheduler [threads|single-threaded|sync|distributed]
                                  Which dask scheduler to use.  [default:
                                  threads]
  -rname, --residual-name TEXT    Name of residual in xds  [default: RESIDUAL]
  -mask, --mask TEXT              Either path to mask.fits or set to mds to
                                  use \ the mask contained in the mds.
                                  [default: <UNSET>]
  -nb, --nband INTEGER            Number of imaging bands  [default: <UNSET>;
                                  required]
  -o, --output-filename BASENAME  Basename of output  [default: <UNSET>;
                                  required]
  --postfix TEXT                  Can be used to specify a custom name for the
                                  image space data \ products
  -p, --product TEXT              Imaging products to produce. Options are I,
                                  Q, U, V. \ Only single Stokes products are
                                  currently supported  [default: I]
  -rc, --row-chunk INTEGER        Number of rows in a chunk in output dataset
                                  [default: -1]
  --algo [hogbom|clark]           Which minor cycle to use  [default: clark]
  --update-mask / --no-update-mask
                                  Will create a mask from clean components and
                                  write it to the mds  [default: no-update-
                                  mask]
  --dirosion INTEGER              Perform dilation followed by erosion with
                                  structure element of this width  [default:
                                  1]
  --fits-mfs / --no-fits-mfs      Output MFS fits files  [default: fits-mfs]
  --fits-cubes / --no-fits-cubes  Output fits cubes  [default: no-fits-cubes]
  --do-residual / --no-do-residual
                                  Make final residual  [default: no-do-
                                  residual]
  --mop-flux / --no-mop-flux      Run PCG based flux mop between major
                                  iterations  [default: mop-flux]
  -nmiter, --nmiter INTEGER       Number of major iterations  [default: 5]
  -th, --threshold FLOAT          Absolute threshold at which to stop
                                  cleaning. By default it is set automatically
                                  using sigmathreshold parameter  [default:
                                  <UNSET>]
  --sigmathreshold FLOAT          Multiple of the rms at which to terminate
                                  the minor cycle  [default: 2.0]
  -gamma, --gamma FLOAT           Minor loop gain  [default: 0.1]
  -pf, --peak-factor FLOAT        Peak factor  [default: 0.15]
  -spf, --sub-peak-factor FLOAT   Peak factor of sub-minor loop  [default:
                                  0.75]
  -hmaxit, --hogbom-maxit INTEGER
                                  Maximum number of peak finding iterations
                                  between major cycles  [default: 5000]
  -cmaxit, --clark-maxit INTEGER  Maximum number of PSF convolutions between
                                  major cycles  [default: 50]
  -smaxit, --sub-maxit INTEGER    Maximum number of iterations for the sub-
                                  minor cycle  [default: 1000]
  -verb, --verbose INTEGER        Verbosity level. Set to 2 for maximum
                                  verbosity, 0 for silence  [default: 1]
  -rf, --report-freq INTEGER      Report frequency for minor cycles  [default:
                                  10]
  -cgtol, --cg-tol FLOAT          Tolreance of conjugate gradient algorithm
                                  [default: 0.001]
  -cgmaxit, --cg-maxit INTEGER    Maximum iterations for conjugate gradient
                                  algorithm  [default: 100]
  -cgminit, --cg-minit INTEGER    Minimum iterations for conjugate gradient
                                  algorithm  [default: 1]
  -cgverb, --cg-verbose INTEGER   Verbosity of conjugate gradient algorithm.
                                  Set to > 1 for \ debugging, 0 for silence
                                  [default: 1]
  -cgrf, --cg-report-freq INTEGER
                                  Report frequency of conjugate gradient
                                  algorithm  [default: 100]
  --backtrack / --no-backtrack    Ensure residual decreases at every iteration
                                  [default: backtrack]
  --help                          Show this message and exit.

This currently serves as the main interface for documenting functionality.

Stimela

To string workers together, we use the internally defined stimela cabs and execute them through a stimela recipe. A simple example recipe implementing single scale clean can be found at ssclean. Note that it is possible to expose only a subset of the full worker inputs through the stimela recipe inputs. This is an attempt to avoid bombarding users with parameters that they may not need (or need not worry about). For example, help for the ssclean recipe can be viewed by running (compare this to the parameters of each worker through the cli interface)

$ stimela help ssclean.yaml
2022-08-21 15:50:39 STIMELA INFO: starting
2022-08-21 15:50:39 STIMELA INFO: loaded full configuration from cache
2022-08-21 15:50:39 STIMELA INFO: available backends: native
2022-08-21 15:50:39 STIMELA INFO: default backend is native
2022-08-21 15:50:39 STIMELA INFO: saving config dependencies to ./stimela.config.deps
2022-08-21 15:50:39 STIMELA INFO: loading recipe/config ssclean.yaml
⠦  stimela ⠦ starting 0:00:00  stimela help ssclean.yaml
└── Recipe: ssclean
    ├── Description:
    │   └── Recipe to perform single scale clean with pfb-clean
    ├── Required inputs:
    │   └── ms               MS
    │       nband            int  Number of imaging bands
    │       output_filename  str  Basename of ouput products
    ├── Optional inputs:
    │   └── data_column              str    Which column to image [default: CORRECTED_DATA]
    │       weight_column            str    [default: WEIGHT_SPECTRUM]
    │       model_column             str    [default: MODEL_DATA]
    │       field_of_view            float  Field of view in degrees
    │       super_resolution_factor  float  Will over-sample Nyquist by this factor at max frequency [default: 2]
    │       nmiter                   int    Number of major iterations [default: 10]
    │       sigmathreshold           float  Multiple of the rms at which to terminate the minor cycle [default: 2.0]
    │       gamma                    float  Minor loop gain [default: 0.1]
    │       peak_factor              float  Peak factor i.e. stop minor cycle once threshold has reached max(residual) * peak_factor [default: 0.15]
    │       outputs                  str    Output products. (m)odel, (r)esidual, (i)mage, (c)lean beam. Captitals for cubes, lower case for MFS images. [default: mMrRiI]
    │       robustness               float  Robustness factor for Briggs weighting. None means natural. A value of less than -2 implies uniform [default: 0]
    │       nthreads                 int    Total number of threads to use between all processes. Will use all available resources by default
    ├── Obscure inputs: omitting 24
    ├── Hidden inputs: omitting 4
    └── Steps:
        └── init     Precompute visibilities for output products. In this case simply the Stokes I visibilities.
            grid     Flip visibilties into image space data products
            clean    Deconvolve with single scale modified Clark clean algorithm
            restore  Restore image and render data products to fits files

The recipe can be executed using e.g.

$ stimela run ssclean.yaml ssclean ms=MS/point_gauss_nb.MS_p0 nband=4 output_filename=out/test field_of_view=1.5 super_resolution_factor=2 nthreads=8 data_column=DATA

It is also possible to only execute a subset of the steps by using the -s option e.g.

stimela run -s clean: ssclean.yaml ssclean ms=MS/point_gauss_nb.MS_p0 nband=4 output_filename=out/test field_of_view=1.5 super_resolution_factor=2 nthreads=8 data_column=DATA

will execute everything from the clean step onwards.