Add astro/msutils.py: finished the "info" command

Based on https://github.com/SpheMakh/msutils (GPLv2)

1 files changed, 352 insertions, 0 deletions

diff --git a/astro/msutils.py b/astro/msutils.py
new file mode 100755
index 0000000..debbee7
--- /dev/null
+++ b/astro/msutils.py
@@ -0,0 +1,352 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2016-2017 Sphesihle Makhathini
+# Copyright (c) 2017 Aaron LI
+# GNU General Public License v2.0 (GPLv2)
+#
+
+"""
+MSUtils - A set of CASA MeasurementSet manipulation tools
+Based on: https://github.com/SpheMakh/msutils
+"""
+
+import argparse
+from collections import OrderedDict
+from pprint import pprint
+
+import numpy as np
+from casacore import tables
+from casacore.tables import table, maketabdesc, makearrcoldesc
+
+
+def getinfo(msname):
+    """
+    Summarize the basic information of a MS.
+
+    Parameters
+    ----------
+    msname : str
+        Name of the MS
+    """
+    tab = tables.table(msname, ack=False)
+
+    info = OrderedDict([
+        ("Ncol", tab.ncols()),
+        ("Nrow", tab.nrows()),
+        ("Ncor", tab.getcell("DATA", 0).shape[-1]),
+        ("Info", tab.info()),
+        ("Keywords", tab.getkeywords().keys()),
+        ("Columns", tab.colnames()),
+        ("ColKeywords", OrderedDict([
+            (cname, tab.getcolkeywords(cname)) for cname in tab.colnames()
+        ])),
+        ("Exposure", tab.getcell("EXPOSURE", 0)),
+        ("FIELD", OrderedDict()),
+        ("SPW", OrderedDict()),
+        ("SCAN", OrderedDict()),
+    ])
+    tabs = {
+        "FIELD": tables.table(msname+"/FIELD", ack=False),
+        "SPW":   tables.table(msname+"/SPECTRAL_WINDOW", ack=False),
+    }
+
+    field_ids = tabs["FIELD"].getcol("SOURCE_ID")
+    info["FIELD"]["STATE_ID"] = [None]*len(field_ids)
+    info["FIELD"]["PERIOD"] = [None]*len(field_ids)
+    for fid in field_ids:
+        ftab = tab.query("FIELD_ID=={0:d}".format(fid))
+        state_id = ftab.getcol("STATE_ID")[0]
+        info["FIELD"]["STATE_ID"][fid] = int(state_id)
+        scans = {}
+        total_length = 0
+        for scan in set(ftab.getcol("SCAN_NUMBER")):
+            stab = ftab.query("SCAN_NUMBER=={0:d}".format(scan))
+            length = (stab.getcol("TIME").max() - stab.getcol("TIME").min())
+            stab.close()
+            scans[str(scan)] = length
+            total_length += length
+
+        info["SCAN"][str(fid)] = scans
+        info["FIELD"]["PERIOD"][fid] = total_length
+        ftab.close()
+
+    for key, _tab in tabs.items():
+        if key == "SPW":
+            colnames = ["CHAN_FREQ", "MEAS_FREQ_REF",
+                        "REF_FREQUENCY", "TOTAL_BANDWIDTH",
+                        "NAME", "NUM_CHAN", "IF_CONV_CHAIN",
+                        "NET_SIDEBAND", "FREQ_GROUP_NAME"]
+        else:
+            colnames = _tab.colnames()
+        for name in colnames:
+            try:
+                info[key][name] = _tab.getcol(name).tolist()
+            except AttributeError:
+                info[key][name] = _tab.getcol(name)
+        _tab.close()
+
+    # Get the minimum and maximum baselines
+    uv = tab.getcol("UVW")[:, :2]
+    baselines = np.sqrt(np.sum(uv**2, axis=1))
+    info["Baseline"] = {"min": baselines.min(), "max": baselines.max()}
+
+    tab.close()
+    return info
+
+
+def addcol(msname, colname=None, shape=None,
+           data_desc_type="array",
+           valuetype=None,
+           init_with=None,
+           coldesc=None,
+           coldmi=None,
+           clone="DATA",
+           rowchunk=None):
+    """
+    Add a column to MS
+
+    Parameters
+    ----------
+    msanme : str
+        MS to which to add the column
+    colname : str
+        Name of the column to be added
+    shape : shape
+    valuetype : data type
+    data_desc_type :
+        * ``scalar`` - scalar elements
+        * ``array`` - array elements
+    init_with : value to initialize the column with
+    """
+    tab = table(msname, readonly=False)
+
+    if colname in tab.colnames():
+        print("Column already exists")
+        return "exists"
+
+    print("Attempting to add %s column to %s" % (colname, msname))
+
+    valuetype = valuetype or "complex"
+
+    if coldesc:
+        data_desc = coldesc
+        shape = coldesc["shape"]
+    elif shape:
+        data_desc = maketabdesc(makearrcoldesc(colname,
+                                               init_with,
+                                               shape=shape,
+                                               valuetype=valuetype))
+    elif valuetype == "scalar":
+        data_desc = maketabdesc(makearrcoldesc(colname,
+                                               init_with,
+                                               valuetype=valuetype))
+    elif clone:
+        element = tab.getcell(clone, 0)
+        try:
+            shape = element.shape
+            data_desc = maketabdesc(makearrcoldesc(colname,
+                                                   element.flatten()[0],
+                                                   shape=shape,
+                                                   valuetype=valuetype))
+        except AttributeError:
+            shape = []
+            data_desc = maketabdesc(makearrcoldesc(colname,
+                                                   element,
+                                                   valuetype=valuetype))
+
+    colinfo = [data_desc, coldmi] if coldmi else [data_desc]
+    tab.addcols(*colinfo)
+
+    print("Column added successfully.")
+
+    if init_with is None:
+        tab.close()
+        return "added"
+    else:
+        spwids = set(tab.getcol("DATA_DESC_ID"))
+        for spw in spwids:
+            print("Initializing column {0}. DDID is {1}".format(colname, spw))
+            tab_spw = tab.query("DATA_DESC_ID=={0:d}".format(spw))
+            nrows = tab_spw.nrows()
+
+            rowchunk = rowchunk or nrows/10
+            dshape = [0] + [a for a in shape]
+            for row0 in range(0, nrows, rowchunk):
+                nr = min(rowchunk, nrows-row0)
+                dshape[0] = nr
+                print("Wrtiting to column  %s (rows %d to %d)" %
+                      (colname, row0, row0+nr-1))
+                dtype = init_with.dtype
+                tab_spw.putcol(colname,
+                               np.ones(dshape, dtype=dtype) * init_with,
+                               row0, nr)
+            tab_spw.close()
+    tab.close()
+
+
+def sumcols(msname, col1=None, col2=None, outcol=None, cols=None,
+            subtract=False):
+    """
+    Add col1 to col2, or sum columns in "cols" list.
+
+    Parameters
+    ----------
+    subtract : bool
+        Subtract ``col2`` from ``col1``
+    """
+    tab = table(msname, readonly=False)
+    if outcol not in tab.colnames():
+        print("outcol {0:s} does not exist, will add it first.".format(outcol))
+        addcol(msname, outcol, clone=col1 or cols[0])
+
+    spws = set(tab.getcol("DATA_DESC_ID"))
+    for spw in spws:
+        tab_spw = tab.query("DATA_DESC_ID=={0:d}".format(spw))
+        nrows = tab_spw.nrows()
+        rowchunk = nrows//10 if nrows > 10000 else nrows
+        for row0 in range(0, nrows, rowchunk):
+            nr = min(rowchunk, nrows-row0)
+            print("Wrtiting to column  %s (rows %d to %d)" %
+                  (outcol, row0, row0+nr-1))
+            if subtract:
+                data = (tab_spw.getcol(col1, row0, nr) -
+                        tab_spw.getcol(col2, row0, nr))
+            else:
+                cols = cols or [col1, col2]
+                data = 0
+                for col in cols:
+                    data += tab.getcol(col, row0, nr)
+
+            tab_spw.putcol(outcol, data, row0, nr)
+        tab_spw.close()
+
+    tab.close()
+
+
+def copycol(msname, fromcol, tocol):
+    """
+        Copy data from one column to another
+    """
+
+    tab = table(msname, readonly=False)
+    if tocol not in tab.colnames():
+        addcol(msname, tocol, clone=fromcol)
+
+    spws = set(tab.getcol("DATA_DESC_ID"))
+    for spw in spws:
+        tab_spw = tab.query("DATA_DESC_ID=={0:d}".format(spw))
+        nrows = tab_spw.nrows()
+        rowchunk = nrows//10 if nrows > 5000 else nrows
+        for row0 in range(0, nrows, rowchunk):
+            nr = min(rowchunk, nrows-row0)
+            data = tab_spw.getcol(fromcol, row0, nr)
+            tab_spw.putcol(tocol, data, row0, nr)
+
+        tab_spw.close()
+    tab.close()
+
+
+def calc_vis_noise(msname, sefd, spw_id=0):
+    """
+    Calculate the nominal per-visibility noise
+    """
+    tab = table(msname)
+    spwtab = table(msname + "/SPECTRAL_WINDOW")
+
+    freq0 = spwtab.getcol("CHAN_FREQ")[spw_id, 0]
+    wavelength = 300e+6/freq0
+    bw = spwtab.getcol("CHAN_WIDTH")[spw_id, 0]
+    dt = tab.getcol("EXPOSURE", 0, 1)[0]
+    dtf = (tab.getcol("TIME", tab.nrows()-1, 1)-tab.getcol("TIME", 0, 1))[0]
+
+    tab.close()
+    spwtab.close()
+
+    print("%s: frequency %.2f MHz (lambda=%.2fm)" %
+          (msname, freq0/1e6, wavelength))
+    print("%s: bandwidth %.2g kHz, %.2fs integration, %.2fh synthesis" %
+          (bw*1e-3, dt, dtf/3600))
+    noise = sefd / np.sqrt(abs(2*bw*dt))
+    print("SEFD of %.2f Jy gives per-visibility noise of %.2f mJy" %
+          (sefd, noise*1000))
+
+    return noise
+
+
+def addnoise(msname, column="MODEL_DATA",
+             noise=0, sefd=551,
+             rowchunk=None,
+             addToCol=None,
+             spw_id=None):
+    """
+    Add Gaussian noise to MS, given a stdandard deviation (noise).
+    This noise can be also be calculated given SEFD value
+    """
+
+    tab = table(msname, readonly=False)
+
+    multi_chan_noise = False
+    if hasattr(noise, "__iter__"):
+        multi_chan_noise = True
+    elif hasattr(sefd, "__iter__"):
+        multi_chan_noise = True
+    else:
+        noise = noise or calc_vis_noise(msname, sefd=sefd,
+                                        spw_id=spw_id or 0)
+
+    spws = set(tab.getcol("DATA_DESC_ID"))
+    for spw in spws:
+        tab_spw = tab.query("DATA_DESC_ID=={0:d}".format(spw))
+        nrows = tab_spw.nrows()
+        nchan, ncor = tab_spw.getcell("DATA", 0).shape
+        rowchunk = rowchunk or nrows/10
+        for row0 in range(0, nrows, rowchunk):
+            nr = min(rowchunk, nrows-row0)
+            data = (np.random.randn(nr, nchan, ncor) +
+                    1j*np.random.randn(nr, nchan, ncor))
+            if multi_chan_noise:
+                noise = noise[np.newaxis, :, np.newaxis]
+            data *= noise
+
+            if addToCol:
+                data += tab_spw.getcol(addToCol, row0, nr)
+                print("%s + noise --> %s (rows %d to %d)" %
+                      (addToCol, column, row0, row0+nr-1))
+            else:
+                print("Adding noise to column %s (rows %d to %d)" %
+                      (column, row0, row0+nr-1))
+
+            tab_spw.putcol(column, data, row0, nr)
+        tab_spw.close()
+
+    tab.close()
+
+
+def cmd_info(args):
+    """
+    Sub-command: "info", show MS basic information
+    """
+    msname = args.msname
+    info = getinfo(msname)
+    pprint(info)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="CASA MeasurementSet (MS) manipulation utilities")
+
+    subparsers = parser.add_subparsers(dest="subparser_name",
+                                       title="sub-commands",
+                                       help="additional help")
+
+    # sub-command: "info"
+    parser_info = subparsers.add_parser("info", help="show MS basic info")
+    parser_info.add_argument("msname", help="MS name")
+    parser_info.set_defaults(func=cmd_info)
+
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()