Merge pull request #74 from ska-sa/issue-62

FITStool cleanups: Issue 62/63/64

Owlcat/FitsTool.py

@@ -24,7 +24,7 @@
 # or write to the Free Software Foundation, Inc.,
 # 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 #
-
+from __future__ import print_function
 import sys
 import re
 import os.path
@@ -237,8 +237,12 @@ def main():
                       help="take mean of input images")
     parser.add_option("-d", "--diff", dest="diff", action="store_true",
                       help="take difference of 2 input images")
-    parser.add_option("--ratio", dest="ratio", action="store_true",
-                      help="take ratio of 2 input images")
+    parser.add_option("--div", dest="div", action="store_true",
+                      help="take ratio of 2 input images (e.g. as in primary beam correction)")
+    parser.add_option("--div-min", type="float",
+                      help="min value to use for second image, when taking a ratio")
+    parser.add_option("--div-min-mask", type="float", metavar="MIN", default=0,
+                      help="mask value to apply to output where second image < MIN")
     parser.add_option("--sum", dest="sum", action="store_true",
                       help="sum of input images")
     parser.add_option("--prod", dest="prod", action="store_true",
@@ -247,6 +251,13 @@ def main():
                       help="transfer data from image 2 into image 1, preserving the FITS header of image 1")
     parser.add_option("-z", "--zoom", dest="zoom", type="int", metavar="NPIX",
                       help="zoom into central region of NPIX x NPIX size")
+    parser.add_option("-p", "--paste", action="store_true",
+                      help="paste image(s) into first image, using the nearest WCS pixel positions. See also --empty-canvas")
+    parser.add_option("--empty-canvas", action="store_true",
+                      help="null first image, then paste image(s) into it")
+    parser.add_option("--null-region", type="str", metavar="NXxNY", help="null central region of specified size")
+    parser.add_option("-M", "--set-max", metavar="MAX", type="float",
+                      help="set values over MAX to MAX")
     parser.add_option("-R", "--rescale", dest="rescale", type="float",
                       help="rescale image values")
     parser.add_option("-E", "--edit-header", metavar="KEY=VALUE", type="string", action="append",
@@ -437,6 +448,26 @@ def main():
             print("Image %s: replaced %d points" % (name, wh.sum()))
         updated = True
 
+    def compute_in_out_slice(N, N0, R, R0):
+        """
+        given an input axis of size N, and an output axis of size N0, and reference pixels of I and I0
+        respectively, computes two slice objects such that
+                A0[slice_out] = A1[slice_in]
+        would do the correct assignment (with I mapping to I0, and the overlapping regions transferred)
+        """
+        i, j = 0, N     # input slice
+        i0 = R0 - R
+        j0 = i0 + N     # output slice
+        if i0 < 0:
+            i = -i0
+            i0 = 0
+        if j0 > N0:
+            j = N - (j0 - N0)
+            j0 = N0
+        if i >= j or i0 >= j0:
+            return None, None
+        return slice(i0, j0), slice(i, j)
+
     if options.transfer:
         if len(images) != 2:
             parser.error("The --transfer option requires exactly two input images.")
@@ -462,14 +493,18 @@ def main():
         for d in images[1:]:
             data += d[0].data
         updated = True
-    elif options.ratio:
+    elif options.div:
         if len(images) != 2:
-            parser.error("The --ratio option requires exactly two input images.")
+            parser.error("The --div option requires exactly two input images.")
         if autoname:
-            outname = "ratio_" + outname
+            outname = "div_" + outname
         print("Computing ratio")
-        data = images[0][0].data
-        data /= images[1][0].data
+        a, b = images[0][0].data, images[1][0].data
+        if options.div_min is not None:
+            b[b<=options.div_min] = options.div_min
+        a /= b
+        if options.div_min_mask is not None:
+            a[b<=options.div_min] = options.div_min_mask
         updated = True
     elif options.prod:
         if autoname:
@@ -489,6 +524,56 @@ def main():
         data /= len(images)
         images = [images[0]]
         updated = True
+    elif options.paste:
+        if autoname:
+            outname = "paste%d_" % len(images) + outname
+        print("Pasting images into canvas given by {}".format(imagenames[0]))
+        print("WARNING: the reference pixel of the pasted image will be aligned with the WCS of the canvas image,")
+        print("         but no WCS reprojection will be attempted. Results may be grossly incorrect if your WCSs")
+        print("         are sufficiently different. Use at own risk.")
+        # get original image size and WCS
+        data0 = images[0][0].data
+        ny0, nx0 = data0.shape[-2:]
+        hdr0 = images[0][0].header
+        wcs0 = WCS(hdr0, mode="pyfits")
+        if options.empty_canvas:
+            print("  canvas image will be cleared initially")
+            data0[:] = 0
+
+        for image, imagename in zip(images[1:], imagenames[1:]):
+            hdr = image[0].header
+            data = image[0].data
+            ny, nx = data.shape[-2:]
+            rx, ry = int(hdr["CRPIX1"]-1), int(hdr["CRPIX2"]-1)
+            # find coordinate of this image's reference pixel in the canvas WCS
+            rx0, ry0 = wcs0.wcs2pix(hdr["CRVAL1"], hdr["CRVAL2"])  # get WCS of center pixel
+            print("image {} rpix {},{} is at {},{} in canvas image".format(imagename, rx, ry, rx0, ry0))
+            rx0 = int(round(rx0))
+            ry0 = int(round(ry0))
+            xout, xin = compute_in_out_slice(nx, nx0, rx, rx0)
+            yout, yin = compute_in_out_slice(ny, ny0, ry, ry0)
+            if xout is not None and yout is not None:
+                print("  canvas X {} will be assigned from {}".format(xout, xin))
+                print("  canvas Y {} will be assigned from {}".format(yout, yin))
+                data0[..., yout, xout] += data[..., yin, xin]
+                updated = True
+            else:
+                print("  image {} footprint is outside of canvas".format(imagename))
+
+    if options.null_region:
+        if len(images) > 1:
+            print("Too many input images specified for this operation, at most 1 expected")
+            sys.exit(2)
+        data = images[0][0].data
+        nx = data.shape[-1]
+        ny = data.shape[-2]
+        xc, yc = nx//2, ny//2
+        rx, ry = map(int,options.null_region.split("x", 1))
+        rx, ry = min(rx, nx), min(ry, ny)
+        x0, y0 = xc - rx//2, yc - ry//2
+        data[..., y0:y0+ry, x0:x0+rx] = 0
+        print("Setting {}:{} {}:{} to zero".format(x0, x0+rx, y0, y0+ry))
+        updated = True
 
     if options.zoom:
         z = options.zoom
@@ -500,9 +585,18 @@ def main():
         data = images[0][0].data
         nx = data.shape[-1]
         ny = data.shape[-2]
-        zx0, zy0 = (nx - z)//2, (ny - z)//2
-        zcen = z//2
-        zdata = data[..., zy0:zy0+z, zx0:zx0+z]
+        # make output array of shape z x z
+        zdata_shape = list(data.shape)
+        zdata_shape[-1] = zdata_shape[-2] = z
+        zdata = numpy.zeros(zdata_shape, dtype=data.dtype)
+        # make input/output slices
+        rx, ry = nx//2, ny//2
+        rx0 = ry0 = z//2
+        xout, xin = compute_in_out_slice(nx, z, rx, rx0)
+        yout, yin = compute_in_out_slice(ny, z, ry, ry0)
+
+        zdata[..., yout, xout] = data[..., yin, xin]
+
         # update header
         hdr = images[0][0].header
         # hdr1 = hdr.copy()
@@ -513,19 +607,24 @@ def main():
         #         hdr1.remove(key + nax, ignore_missing=True)
         wcs = WCS(hdr) #, mode="pyfits")
         pixcoord = numpy.zeros((1, hdr["NAXIS"]), float)
-        pixcoord[0, 0] = zx0 + zcen
-        pixcoord[0, 1] = zy0 + zcen
+        pixcoord[0, 0] = nx//2
+        pixcoord[0, 1] = ny//2
         world = wcs.wcs_pix2world(pixcoord, 0)  # get WCS of center pixel
         hdr["CRVAL1"] = world[0,0]
         hdr["CRVAL2"] = world[0,1]
-        hdr["CRPIX1"] = zcen + 1
-        hdr["CRPIX2"] = zcen + 1
+        hdr["CRPIX1"] = rx0 + 1
+        hdr["CRPIX2"] = ry0 + 1
 
         print("Making zoomed image of shape", "x".join(map(str, zdata.shape)))
 
         images = [pyfits.PrimaryHDU(zdata, hdr)]
         updated = True
 
+    if options.set_max is not None:
+        print("Enforcing maximum value of {}".format(options.set_max))
+        images[0][0].data[images[0][0].data > options.set_max] = options.set_max
+        updated = True
+
     if options.rescale != 1:
         if autoname and not updated:
             outname = "rescale_" + outname
@@ -542,11 +641,11 @@ def main():
     if options.stats:
         for ff, filename in zip(images, imagenames):
             data = ff[0].data
-            min, max, dum1, dum2 = scipy.ndimage.measurements.extrema(data)
+            min1, max1, dum1, dum2 = scipy.ndimage.measurements.extrema(data)
             sum = data.sum()
             mean = sum / data.size
             std = math.sqrt(((data - mean) ** 2).mean())
-            print("%s: min %g, max %g, sum %g, np %d, mean %g, std %g" % (filename, min, max, sum, data.size, mean, std))
+            print("%s: min %g, max %g, sum %g, np %d, mean %g, std %g" % (filename, min1, max1, sum, data.size, mean, std))
         sys.exit(0)
 
     if updated:

Owlcat/bin/flag-ms.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
+# -*- coding: future_fstrings -*-
 
 #
 # % $Id$
@@ -30,6 +30,9 @@
 import sys
 import re
 import traceback
+import glob
+
+from casacore.tables.tableutil import tableexists
 import Owlcat.Tables
 
 flagger = parser = ms = msname = None
@@ -236,7 +239,7 @@ def parse_subset_options(options):
     group.add_option("--incr-stman", action="store_true",
                      help="force the use of the incremental storage manager for new BITFLAG columns. Default is to use same manager as DATA column.")
     group.add_option("-l", "--list", action="store_true",
-                     help="lists various info about the MS, including its flagsets.")
+                     help="lists various info about the MS, including its flagsets, and CASA flagversions, if available.")
     group.add_option("-s", "--stats", action="store_true",
                      help="prints per-flagset flagging stats.")
     group.add_option("-r", "--remove", metavar="FLAGSET(s)", type="string",
@@ -245,6 +248,8 @@ def parse_subset_options(options):
                      help="exports all flags to flag file. FILENAME may end with .gz to produce a gzip-compressed file. If any flagging actions are specified, these will be done before the export.")
     group.add_option("--import", type="string", dest="_import", metavar="FILENAME",
                      help="imports flags from flag file. If any flagging actions are specified, these will be done after the import.")
+    group.add_option("-R", "--restore", type="string", dest="restore", metavar="VERSION",
+                     help="imports flags from a CASA-style flagversion (use --list to list version names).")
     group.add_option("-v", "--verbose", metavar="LEVEL", type="int",
                      help="verbosity level for messages. Higher is more verbose, default is 0.")
     group.add_option("--timestamps", action="store_true",
@@ -282,9 +287,28 @@ def parse_subset_options(options):
             error("Error importing flags from %s, exiting" % options._import)
         print("Flags imported OK.")
 
+    # same for CASA flagversions
+    if options.restore:
+        flagvers = f"{msname}.flagversions/flags.{options.restore}"
+        if not Owlcat.tableexists(flagvers):
+            error(f"No such flagversion '{options.restore}'. Use --list to list available flag versions.")
+        try:
+            ms = get_ms(readonly=False)
+            ftab = Owlcat.table(flagvers)
+            flag_row = ftab.getcol("FLAG_ROW")
+            flag_col = ftab.getcol("FLAG")
+            ftab.close()
+            ms.putcol("FLAG_ROW", flag_row)
+            ms.putcol("FLAG", flag_col)
+            ms.close()
+        except:
+            traceback.print_exc()
+            error(f"Error restoring flags from {options.restore}, exiting")
+        print(f"Flag version '{options.restore}' restored.")
+
     # if no other actions supplied, enable stats (unless flags were imported, in which case just exit)
     if not (options.flag or options.unflag or options.copy or options.fill_legacy):
-        if options._import:
+        if options._import or options.restore:
             sys.exit(0)
         statonly = True
     else:
@@ -339,7 +363,7 @@ def parse_subset_options(options):
             fields = Owlcat.table(ms.getkeyword('FIELD'), ack=False).getcol('NAME')
 
             print("===> MS is %s" % msname)
-            print("  %d antennas: %s" % (len(ants), " ".join(ants)))
+            print("  %d antennas: %s" % (len(ants), " ".join(f"{num}:{name}" for num, name in enumerate(ants))))
             print("  %d DATA_DESC_ID(s): " % len(spwids))
             for i, (spw, pol) in enumerate(zip(spwids, polids)):
                 print("    %d: %.3f MHz, %d chans x %d correlations" % (
@@ -356,9 +380,18 @@ def parse_subset_options(options):
                         print("    '%s': %d (0x%02X)" % (name, mask, mask))
                 else:
                     print("  No flagsets.")
-            print("")
+            # now check for CASA flagversions
+            flagvers = [tab for tab in glob.glob(f"{msname}.flagversions/flags.*") if Owlcat.tableexists(tab)]
+            if flagvers:
+                print(f"  {len(flagvers)} CASA-style flagversions available for --restore")
+                for vers in flagvers:
+                    name = os.path.basename(vers).split('.', 1)[1]
+                    print(f"    '{name}'")
+            else:
+                print("  No CASA-style flagversions available.")
             if options.flag or options.unflag or options.copy or options.fill_legacy or options.remove:
                 print("-l/--list was in effect, so all other options were ignored.")
+            print("")
             sys.exit(0)
 
         if options.copy_legacy:

setup.py

@@ -4,20 +4,21 @@
 from setuptools import setup
 
 install_requires = [
-    'astropy',
-    'numpy',
-    'matplotlib',
-    'python-casacore',
-    'meqtrees_cattery',
-    'scipy',
-    'astro-kittens',
-    'six',
-    'future',
-    'bokeh'
+      'astropy',
+      'numpy',
+      'matplotlib',
+      'python-casacore',
+      'meqtrees_cattery',
+      'scipy',
+      'astro-kittens',
+      'future-fstrings',
+      'six',
+      'future',
+      'bokeh'
 ]
 
 setup(name='owlcat',
-      version='1.6.3',
+      version='1.6.4',
       description='miscellaneous utility scripts for manipulating radio interferometry data',
       author='Oleg Smirnov',
       author_email='Oleg Smirnov <[email protected]>',