Set min counts to fraction of median (#100)

* set low counts to med/level

* remove codex-africanus from setup.py

* add option to write out FFT of residual in restore worker

* opts.opts typo

* use c2c and take fftshift when writing gridded residuals in restore

* upgrade streamjoy requirement and set crf=5 to improve mp4 output quality

* drop empty frames

* pop using index

* make gif optimisation optional

* set frames to max

* Make crf an input parameter to control mp4 quality

* add meta info to map_blocks funcs

* make it possible to ignore dual in dds2cubes even when dual exists

* depend on QuartiCal v0.2.3-fix-ver

* change clean -> klean and fluxmop -> fluxmop

* depend on QuartiCal >= 0.2.3

* depend on QuartiCal >= 0.2.3 (save changes)

* start renaming pfb-clean -> pfb-imaging

* comment debugging code

---------

Co-authored-by: landmanbester <lbester@ska.ac.za>

.github/workflows/ci.yml

@@ -1,4 +1,4 @@
-name: pfb-clean Workflow
+name: pfb-imaging Workflow
 
 on:
   push:
@@ -31,7 +31,7 @@ jobs:
       # - name: Pin setuptools
       #   run: python -m pip install setuptools==65.5
 
-      - name: Install pfb-clean
+      - name: Install pfb-imaging
         run: python -m pip install .[testing]
 
       - name: Run tests

Dockerfile

@@ -12,6 +12,6 @@ RUN apt -y update && \
 
 
 RUN python -m pip install -U pip setuptools wheel && \
-    python -m pip install -U pfb-clean@git+https://github.com/ratt-ru/pfb-clean@main && \
+    python -m pip install -U pfb-imaging@git+https://github.com/ratt-ru/pfb-imaging@main && \
     python -m pip install numpy==1.22 && \
     python -m pip cache purge

README.rst

@@ -1,10 +1,10 @@
-pfb-clean
+pfb-imaging
 =========
-Preconditioned forward-backward clean algorithm.
+Radio interferometric imaging suite base on the pre-conditioned forward-backward algorithm.
 
 Install the package by cloning and running
 
-:code:`$ pip install -e pfb-clean/`
+:code:`$ pip install -e pfb-imaging/`
 
 Note casacore needs to be installed on the system for this to work.
 
@@ -19,14 +19,6 @@ no binary mode eg
 
 :code:`$ pip install -e ducc`
 
-You may also have to make numba aware of the tbb layer by doing
-
-:code:`$ pip install tbb`
-
-:code:`$ export LD_LIBRARY_PATH=/path/to/venv/lib`
-
-see eg. https://github.com/ratt-ru/QuartiCal/issues/268
-
 If you find any of this useful please cite (for now)
 
 https://arxiv.org/abs/2101.08072

pfb/parser/restore.yaml

@@ -6,31 +6,37 @@ inputs:
     dtype: str
     abbreviation: mname
     default: MODEL
-    info: 'Name of model in mds'
+    info:
+      Name of model in dds
   residual-name:
     dtype: str
     abbreviation: rname
     default: RESIDUAL
-    info: 'Name of residual in dds'
+    info:
+      Name of residual in dds
   nband:
     dtype: int
     required: true
     abbreviation: nb
-    info: 'Number of imaging bands'
+    info:
+      Number of imaging bands
   postfix:
     dtype: str
     default: 'main'
-    info: 'Can be used to specify a custom name for the image space data \
-          products'
+    info:
+      Can be used to specify a custom name for the image space data products
   outputs:
     dtype: str
     default: mMrRiI
-    info: 'Output products. (m)odel, (r)esidual, (i)mage, (c)lean beam. \
-          Captitals correspond to cubes.'
+    info:
+      Output products (m)odel, (r)esidual, (i)mage, (c)lean beam, (d)irty,
+      (f)ft_residuals (amplitude and phase will be produced).
+      Use captitals to produce corresponding cubes.
   overwrite:
     dtype: bool
     default: false
-    info: Allow overwrite of output xds
+    info:
+      Allow overwriting fits files
 
 outputs:
     {}

pfb/parser/smoovie.yaml

@@ -152,6 +152,17 @@ inputs:
       The format to write movie out in.
       gifs are usually better quality but can get quite large.
       mp4 quality not currently great with streamjoy.
+  optimize:
+    dtype: bool
+    default: false
+    info:
+      To try and optimize the resulting gif.
+      Only possible if gifsicle is installed.
+  crf:
+    dtype: int
+    default: 12
+    info:
+      Constant rate factor for controlling mp4 output quality.
 
 outputs:
     {}

pfb/parser/spotless.yaml

@@ -96,7 +96,13 @@ inputs:
     dtype: float
     default: 0.5
     info:
-      Lower the initail rmsfactor by this amount
+      Lower the initial rmsfactor by this amount
+  diverge-count:
+    dtype: int
+    default: 5
+    info:
+      Will terminate the algorithm if the rms increases this many times.
+      Set to > niter to disable this check.
 
 outputs:
   {}

pfb/parser/uncabbedcabs.yml

@@ -25,14 +25,14 @@ pfb.degrid:
   _include:
     - (.)degrid.yaml
 
-pfb.clean:
-  command: pfb.workers.clean.clean
+pfb.klean:
+  command: pfb.workers.klean.klean
   flavour: python
   policies:
     pass_missing_as_none: true
 
   _include:
-    - (.)clean.yaml
+    - (.)klean.yaml
 
 pfb.restore:
   command: pfb.workers.restore.restore
@@ -52,14 +52,14 @@ pfb.fwdbwd:
   _include:
     - (.)fwdbwd.yaml
 
-pfb.forward:
-  command: pfb.workers.forward.forward
+pfb.fluxmop:
+  command: pfb.workers.fluxmop.fluxmop
   flavour: python
   policies:
     pass_missing_as_none: true
 
   _include:
-    - (.)forward.yaml
+    - (.)fluxmop.yaml
 
 pfb.spotless:
   command: pfb.workers.spotless.spotless

pfb/utils/fits.py

@@ -68,7 +68,7 @@ def set_wcs(cell_x, cell_y, nx, ny, radec, freq,
 
     header = w.to_header()
     header['RESTFRQ'] = fmean
-    header['ORIGIN'] = 'pfb-clean'
+    header['ORIGIN'] = 'pfb-imaging'
     header['BTYPE'] = 'Intensity'
     header['BUNIT'] = unit
     header['SPECSYS'] = 'TOPOCENT'

pfb/utils/misc.py

@@ -503,13 +503,14 @@ def _restore_corrs(vis, ncorr):
     return model_vis
 
 
-# model to fit
 @jax.jit
 def psf_errorsq(x, data, xy):
+    '''
+    Returns sum of square error for best fit Gaussian to data
+    '''
     emaj, emin, pa = x
     Smin = jnp.minimum(emaj, emin)
     Smaj = jnp.maximum(emaj, emin)
-    # print(emaj, emin, pa)
     A = jnp.array([[1. / Smin ** 2, 0],
                     [0, 1. / Smaj ** 2]])
 
@@ -519,7 +520,7 @@ def psf_errorsq(x, data, xy):
     B = jnp.dot(jnp.dot(R.T, A), R)
     Q = jnp.einsum('nb,bc,cn->n', xy.T, B, xy)
     # GaussPar should corresponds to FWHM
-    fwhm_conv = 2 * jnp.sqrt(2 * np.log(2))
+    fwhm_conv = 2 * jnp.sqrt(2 * jnp.log(2))
     model = jnp.exp(-fwhm_conv * Q)
     res = data - model
     return jnp.vdot(res, res)
@@ -660,7 +661,7 @@ def init_mask(mask, model, output_type, log):
     return mask
 
 
-def dds2cubes(dds, nband, apparent=False):
+def dds2cubes(dds, nband, apparent=False, dual=True):
     real_type = dds[0].DIRTY.dtype
     complex_type = np.result_type(real_type, np.complex64)
     nx, ny = dds[0].DIRTY.shape
@@ -686,7 +687,7 @@ def dds2cubes(dds, nband, apparent=False):
         psfhat = None
     mean_beam = [da.zeros((nx, ny), chunks=(-1, -1),
                             dtype=real_type) for _ in range(nband)]
-    if 'DUAL' in dds[0]:
+    if dual and 'DUAL' in dds[0]:
         nbasis, nymax, nxmax = dds[0].DUAL.shape
         dual = [da.zeros((nbasis, nymax, nxmax), chunks=(-1, -1, -1),
                             dtype=real_type) for _ in range(nband)]
@@ -707,7 +708,7 @@ def dds2cubes(dds, nband, apparent=False):
             psfhat[b] += ds.PSFHAT.data
         if 'MODEL' in ds:
             model[b] = ds.MODEL.data
-        if 'DUAL' in ds:
+        if dual and 'DUAL' in ds:
             dual[b] = ds.DUAL.data
         mean_beam[b] += ds.BEAM.data * ds.WSUM.data[0]
         wsums[b] += ds.WSUM.data[0]
@@ -720,7 +721,7 @@ def dds2cubes(dds, nband, apparent=False):
     if 'PSF' in ds:
         psf = da.stack(psf)/wsum
         psfhat = da.stack(psfhat)/wsum
-    if 'DUAL' in ds:
+    if dual and 'DUAL' in ds:
         dual = da.stack(dual)
     for b in range(nband):
         if wsums[b]:
@@ -1441,3 +1442,72 @@ def combine_columns(x, y, dc, dc1, dc2):
                 out=x,
                 casting='same_kind')
     return x
+
+
+# def fft_interp(image, cellxi, cellyi, nxo, nyo,
+#                cellxo, cellyo, shiftx, shifty):
+#     '''
+#     Use non-uniform fft to interpolate image in a flux conservative way
+
+#     image   - input image
+#     cellxi  - input x cell-size
+#     cellyi  - input y cell-size
+#     nxo     - number of x pixels in output
+#     nyo     - number of y pixels in output
+#     cellxo  - output x cell size
+#     cellyo  - output y cell size
+#     shiftx  - shift x coordinate by this amount
+#     shifty  - shift y coordinate by this amount
+
+#     All sizes are assumed to be in radians.
+#         '''
+#     import matplotlib.pyplot as plt
+#     from scipy.fft import fftn, ifftn
+#     Fs = np.fft.fftshift
+#     iFs = np.fft.ifftshift
+#     # basic
+#     nx, ny = image.shape
+#     imhat = Fs(fftn(image))
+
+#     imabs = np.abs(imhat)
+#     imphase = np.angle(imhat) - 1.0
+#     # imphase = np.roll(imphase, nx//2, axis=0)
+#     imshift = ifftn(iFs(imabs*np.exp(1j*imphase))).real
+
+#     impad = np.pad(imhat, ((nx//2, nx//2), (ny//2, ny//2)), mode='constant')
+#     imo = ifftn(iFs(impad)).real
+
+#     print(np.sum(image) - np.sum(imo))
+
+#     plt.figure(1)
+#     plt.imshow(image/image.max(), vmin=0, vmax=1, interpolation=None)
+#     plt.colorbar()
+#     plt.figure(2)
+#     plt.imshow(imo/imo.max(), vmin=0, vmax=1, interpolation=None)
+#     plt.colorbar()
+#     plt.figure(3)
+#     plt.imshow(imshift/imshift.max() - image/image.max(), vmin=0, vmax=1, interpolation=None)
+#     plt.colorbar()
+
+#     plt.show()
+
+    # # coordinates on input grid
+    # nx, ny = image.shapeimhat
+    # x = np.arange(-(nx//2), nx//2) * cellxi
+    # y = np.arange(-(ny//2), ny//2) * cellyi
+    # xx, yy = np.meshgrid(x, y, indexing='ij')
+
+    # # frequencies on output grid
+    # celluo = 1/(nxo*cellxo)
+    # cellvo = 1/(nyo*cellyo)
+    # uo = np.arange(-(nxo//2), nxo//2) * celluo/nxo
+    # vo = np.arange(-(nyo//2), nyo//2) * cellvo/nyo
+
+    # uu, vv = np.meshgrid(uo, vo, indexing='ij')
+    # uv = np.vstack((uo, vo)).T
+
+
+    # res1 = finufft.nufft2d3(xx.ravel(), yy.ravel(), image.ravel(), uu.ravel(), vv.ravel())
+
+
+

pfb/utils/stokes2vis.py

@@ -60,7 +60,8 @@ def single_stokes(ds=None,
                              dc,
                              dc1,
                              dc2,
-                             chunks=data.chunks)
+                             chunks=data.chunks,
+                             meta=np.empty((0,0,0), dtype=data.dtype))
 
 
     nrow, nchan, ncorr = data.shape
@@ -90,7 +91,8 @@ def single_stokes(ds=None,
         # weight = 1.0/sigma**2
         weight = da.map_blocks(weight_from_sigma,
                                sigma,
-                               chunks=sigma.chunks)
+                               chunks=sigma.chunks,
+                               meta=np.empty((0,0,0), dtype=sigma.dtype))
     elif opts.weight_column is not None:
         weight = getattr(ds, opts.weight_column).data
         if opts.weight_column=='WEIGHT':

pfb/utils/weighting.py

@@ -177,14 +177,16 @@ def filter_extreme_counts(counts, nbox=16, nlevel=10):
                         counts, 'xy',
                         nbox, None,
                         nlevel, None,
-                        dtype=counts.dtype)
+                        dtype=counts.dtype,
+                        meta=np.empty((0,0), dtype=float))
 
 
 
-@njit(nogil=True, cache=True)
+# @njit(nogil=True, cache=True)
 def _filter_extreme_counts(counts, nbox=16, level=10.0):
     '''
-    Replaces extreme counts by local mean computed i
+    Replaces extremely small counts by median to prevent
+    upweighting nearly empty cells
     '''
     # nx, ny = counts.shape
     # I, J = np.where(counts>0)
@@ -203,10 +205,13 @@ def _filter_extreme_counts(counts, nbox=16, level=10.0):
     #     local_mean = np.mean(tmp[ix, iy])
     #     if counts[i,j] < local_mean/level:
     #         counts[i, j] = local_mean
-    # get the median ounts value
-    med = np.median(counts>0)
-    counts = np.where(counts > med/level, counts, med)
-
+    # get the median counts value
+    ix, iy = np.where(counts > 0)
+    cnts = counts[ix,iy]
+    med = np.median(cnts)
+    lowval = med/level
+    cnts = np.maximum(cnts, lowval)
+    counts[ix,iy] = cnts
     return counts