Merge pull request #15 from mtazzari/export_uvtable

Add export_uvtable function

.gitignore

@@ -102,3 +102,6 @@ ENV/
 
 # JetBrains
 .idea/
+
+# MacOS
+*.DS_Store

README.rst

@@ -22,6 +22,7 @@ The current version implements the basic plotting functionality. More features a
 
 If you are interested, or have feature requests, or encounter issues, consider creating an Issue or writing me an `email  <[email protected]>`_. I am happy to have your feedback!
 
+What's new in **v0.2.0**: a new `export_uvtable` function to export visibilities from an MS to an ASCII table.
 
 Installation
 ------------
@@ -32,9 +33,18 @@ Installation
 
     pip install git+https://github.com/mtazzari/uvplot.git
 
+To make **uvplot** available in CASA, run from the shell:
+
+.. code-block :: bash
+
+    casa-pip install git+https://github.com/mtazzari/uvplot.git
+
+where `casa-pip` is a tool that can be downloaded `here <https://github.com/radio-astro-tools/casa-python>`_ .
+
+**uvplot** has been tested on CASA versions > 4.7.0.
+
 Example
 -------
-
 This is an example plot:
 
 .. image:: static/uvplot.png

setup.py

@@ -5,13 +5,13 @@
 
 setup(
     name="uvplot",
-    version="0.1.0",
+    version="0.2.0",
     packages=find_packages(),
     author="Marco Tazzari",
     author_email="[email protected]",
     description="Utilities for plotting interferometric visibilities.",
     long_description=open('README.rst').read(),
-    install_requires=["numpy>1.11", "matplotlib"],
+    install_requires=["numpy>=1.9", "matplotlib"],
     license="LGPLv3",
     url="tbd",
     classifiers=[

uvplot/__init__.py

@@ -15,5 +15,5 @@
 if matplotlib.get_backend().lower() == 'macosx':
     matplotlib.use('TkAgg')
 
-from .uvtable import UVTable
+from .uvtable import UVTable, export_uvtable
 from .constants import arcsec

uvplot/constants.py

@@ -4,5 +4,6 @@
 from __future__ import (division, print_function, absolute_import,
                         unicode_literals)
 
-arcsec = 4.84813681109536e-06       # radians
+arcsec = 4.84813681109536e-06   # radians
+clight = 2.99792458e+8          # [m/s] Speed of light
 

uvplot/uvtable.py

@@ -8,7 +8,9 @@
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec
 
-__all__ = ["UVTable"]
+from .constants import clight
+
+__all__ = ["UVTable", "export_uvtable"]
 
 
 class UVTable(object):
@@ -37,7 +39,7 @@ def __init__(self, uvtable=None, filename="", wle=1, **kwargs):
         if filename != "":
             self.filename = filename
 
-            uvdata = self.read_uvtable(self.filename, kwargs.get('format', 'uvtable'))
+            uvdata = self.read_uvtable(self.filename, kwargs.get('format', 'ascii'))
 
             u = uvdata[:, 0]
             v = uvdata[:, 1]
@@ -61,7 +63,7 @@ def __init__(self, uvtable=None, filename="", wle=1, **kwargs):
     @staticmethod
     def read_uvtable(filename, format):
         """ Read uvtable from file, given a specific format. """
-        if format == 'uvtable':
+        if format == 'ascii':
             uvdata = np.loadtxt(filename)
 
         else:
@@ -383,3 +385,167 @@ def plot(self, fig_filename=None, color='k', linestyle='.', label='',
             plt.savefig(fig_filename)
         else:
             return axes
+
+
+def export_uvtable(uvtable_filename, tb, vis="", split_args=None, split=None,
+                   dualpol=True, fmt='%10.6e', keepms=False, verbose=False):
+    """
+    Export visibilities from an MS Table to a uvtable. Requires execution inside CASA.
+    Currently the only uvtable format supported is ASCII.
+
+    Typicall call signature::
+
+    export_uvtable('uvtable_new.txt', tb, split=split,
+                   split_args={'vis': 'sample.ms', datacolumn': 'data', spw='0,1}, verbose=True)"
+
+    Parameters
+    ----------
+    uvtable_filename : str
+        Filename of the output uvtable, e.g. "uvtable.txt"
+    tb : CASA `tb` object
+        As tb parameter you **must** pass the tb object that is defined in the CASA shell.
+        Since tb cannot be accessed outside CASA, export_uvtable('uvtable.txt', tb, ...) 
+        can be executed only inside CASA.
+    vis : str, optional
+        MS Table filename, e.g. mstable.ms
+    split_args : dict, optional
+        Default is None. If provided, perform a split before exporting the uvtable.
+        The split_args dictionary is passed to the CASA::split task.
+        The CASA::split task must be provided in input as split.
+    split : optional
+        CASA split task
+    dualpol : bool, optional
+        If the MS Table contains dual polarisation data. Default is True.
+    fmt : str, optional
+        Format of the output ASCII uvtable.
+    keepms : bool, optional
+        If True, keeps the outputvis created by the split command.
+    verbose: bool, optional
+        If True, print informative messages.
+
+    Note
+    ----
+    By default, only the 1st spectral window (spw) is exported.
+    To export all the spws in an MS table provide split_args, e.g.::
+
+        split_args = {'vis': 'input.ms', 'outputvis':'input_tmp.ms', spw:'*'}
+
+    Example
+    -------
+    From within CASA, to extract all the visibilities from an MS table::
+
+        export_uvtable('uvtable.txt', tb, vis='sample.ms')
+
+    where `tb` is the CASA tb object (to inspect it type `tb` in the CASA shell).
+    For more information on `tb` see `<https://casa.nrao.edu/docs/CasaRef/table-Module.html>`_
+    
+    From within CASA, to extract the visibilities in spectral windows 0 and 2 use
+    the `split_args` parameter and the CASA `split` task::
+
+        export_uvtable('uvtable.txt', tb, split=split,
+         split_args={'vis': 'sample.ms' , 'datacolumn': 'data', 'spw':'0,2'})
+
+    To perform these operations without running CASA interactively::
+
+        casa --nologger --nogui -c "from uvplot import export_uvtable; export_uvtable(...)"
+
+    ensuring that the strings are inside the '..' characters and not the "..." one.
+
+    """
+    if vis != "":
+        MStb_name = vis
+
+    if split_args:
+        if split is None:
+            raise RuntimeError("Missing split parameter: provide the CASA split object in input. See typical call signature.")
+        if vis != "" and vis != split_args['vis']:
+            # raise RuntimeError("extract_uvtable: either provide `vis` or `split_args` as input parameters, not both.")
+            raise RuntimeError(
+                "extract_uvtable: vis={} input parameter doesn't match with split_args['vis']={}".format(
+                    vis, split_args['vis']))
+
+        if not 'outputvis' in split_args.keys():
+            split_args.update(outputvis='mstable_tmp.ms'.encode())
+
+        MStb_name = split_args['outputvis']
+
+        if verbose:
+            print("Applying split. Creating temporary MS table {} from original MS table {}".format(MStb_name, split_args['vis']))
+
+        split(**split_args)
+    else:
+        if vis == "":
+            raise RuntimeError("Missing vis parameter: provide a valid MS table filename.")
+
+    if verbose: print("Reading {}".format(MStb_name))
+
+    tb.open(MStb_name)
+
+    # get coordinates
+    uvw = tb.getcol("UVW".encode())
+    u, v, w = [uvw[i, :] for i in range(3)]
+
+    # get weights
+    weights_orig = tb.getcol("WEIGHT".encode())
+
+    # get visibilities
+    tb_columns = tb.colnames()
+    if "CORRECTED_DATA" in tb_columns:
+        if verbose: print("Reading CORRECTED_DATA column")
+        data = tb.getcol("CORRECTED_DATA".encode())
+    else:
+        if verbose: print("CORRECTED_DATA column not found. Reading DATA column")
+        data = tb.getcol("DATA".encode())
+
+    spw = tb.getcol("DATA_DESC_ID".encode())
+    nspw = len(np.unique(spw))
+    if nspw > 1:
+        if split_args is None or (split_args is not None and 'spw' not in split_args.keys()):
+            print(
+            "Warning: the MS table {} has {} spectral windows. By default all of them are exported."
+            " To choose which spws to export, provide split_args with the spw parameter.".format(
+                MStb_name, nspw))
+
+    ispw = 0
+
+    if dualpol:
+        # dual polarisation: extract the polarised visibilities and weights
+        V_XX = data[0, ispw, :]
+        V_YY = data[1, ispw, :]
+        weights_XX = weights_orig[0, :]
+        weights_YY = weights_orig[1, :]
+
+        # compute weighted average of the visibilities and weights
+        V = (V_XX * weights_XX + V_YY * weights_YY) / (weights_XX + weights_YY)
+        weights = 0.5 * (weights_XX + weights_YY)
+
+    else:
+        # single polarisation
+        V = data[0, ispw, :]
+        weigths = weights_orig
+
+    spw_path = tb.getkeyword('SPECTRAL_WINDOW'.encode()).split()[-1]
+    tb.close()
+
+    # get the mean observing frequency
+    tb.open(spw_path)
+    freqs = tb.getcol('CHAN_FREQ'.encode())  # [GHz]
+    tb.close()
+
+    wle = clight / freqs.mean()  # [m]
+
+    # export to file as ascii
+    if verbose: print(
+        "Exporting visibilities to {}".format(uvtable_filename))
+    np.savetxt(uvtable_filename,
+               np.column_stack([u, v, V.real, V.imag, weights]), fmt=fmt.encode(),
+               delimiter='\t',
+               header='Extracted from {}.\nwavelength[m] = {}\nColumns:\tu[m]\tv[m]\tRe(V)[Jy]\tIm(V)[Jy]\tweight'.format(
+                   MStb_name, wle))
+
+    if split_args:
+        if not keepms:
+            from subprocess import call
+            if verbose:
+                print("Removing temporary MS table {}".format(split_args['outputvis']))
+            call("rm -rf {}".format(split_args['outputvis']), shell=True)