support multi-column expression and depency on daskms master

setup.py

@@ -9,13 +9,16 @@
     readme = readme_file.read()
 
 requirements = [
+    'dask-ms'
+    '@git+https://github.com/ska-sa/dask-ms.git'
+    '@master',
     'dask[array] >= 2.2.0',
     'donfig >= 0.4.0',
     'numpy >= 1.14.0',
     'numba >= 0.43.0',
     'scipy >= 1.2.0',
     'threadpoolctl >= 1.0.0',
-    'dask-ms >= 0.2.3',
+    # 'dask-ms >= 0.2.3',
     'zarr >= 2.3.1'
 ]
 

tricolour/apps/tricolour/app.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 """ Main tricolour application """
-
+import warnings
 import re
 import argparse
 import contextlib
@@ -22,6 +22,7 @@
                               CacheProfiler, visualize)
 import numpy as np
 from daskms import xds_from_ms, xds_from_table, xds_to_table
+from daskms.expressions import data_column_expr
 from threadpoolctl import threadpool_limits
 
 from tricolour.apps.tricolour.strat_executor import StrategyExecutor
@@ -269,7 +270,14 @@ def _main(args):
         post_mortem_handler.disable_pdb_on_error()
 
     log.info("Flagging on the {0:s} column".format(args.data_column))
-    data_column = args.data_column
+    if "/" in args.data_column:
+        data_string = args.data_column.split("/")
+        data_column = data_string[0]
+        string_columns = data_string[1]
+    else:
+        data_column = args.data_column
+        string_columns = None
+
     masked_channels = [load_mask(fn, dilate=args.dilate_masks)
                        for fn in collect_masks()]
     GD = args.config
@@ -288,7 +296,18 @@ def _main(args):
                "ANTENNA1",
                "ANTENNA2"]
 
+    # Get the columns to read
+    if string_columns is not None and "(" in string_columns:
+        match = re.search(r'\(([A-Za-z0-9+-_]+)\)', string_columns)
+        multi_columns = re.split(r'\+|-', match.group(1))
+        multi_columns = list(filter(None, multi_columns))
+        columns.extend(multi_columns)
+    else:
+        if string_columns is not None:
+            columns.append(string_columns)
+
     if args.subtract_model_column is not None:
+        warnings.warn("Use -dc arg")
         columns.append(args.subtract_model_column)
 
     xds = list(xds_from_ms(args.ms,
@@ -297,6 +316,9 @@ def _main(args):
                            index_cols=index_cols,
                            chunks={"row": args.row_chunks}))
 
+    string = "EXPR = " + args.data_column
+    vis = data_column_expr(string, xds)
+
     # Get support tables
     st = support_tables(args.ms)
     ddid_ds = st["DATA_DESCRIPTION"]

tricolour/apps/tricolour/somthing.py

@@ -0,0 +1,82 @@
+import re
+import argparse
+import logging
+import logging.handlers
+
+
+def create_logger():
+    """ Create a console logger """
+    log = logging.getLogger("tricolour")
+    cfmt = logging.Formatter(u'%(name)s - %(asctime)s '
+                             '%(levelname)s - %(message)s')
+    log.setLevel(logging.DEBUG)
+    filehandler = logging.FileHandler("tricolour.log")
+    filehandler.setFormatter(cfmt)
+    log.addHandler(filehandler)
+    log.setLevel(logging.INFO)
+
+    console = logging.StreamHandler()
+    console.setLevel(logging.INFO)
+    console.setFormatter(cfmt)
+
+    log.addHandler(console)
+
+    return log
+
+
+log = create_logger()
+
+
+def create_parser():
+    # warnings.warn("Use -dc flag")
+    formatter = argparse.ArgumentDefaultsHelpFormatter
+    p = argparse.ArgumentParser(formatter_class=formatter)
+    p.add_argument("ms", help="Measurement Set")
+    p.add_argument("-dc", "--data-column", type=str, default="DATA",
+                   help="Name of visibility data column to flag. You can "
+                   "specify DATA(+-)MODEL for column arithmetic")
+    p.add_argument("-smc", "--subtract-model-column", default=None, type=str,
+                   help="Subtracts specified column from data column "
+                        "specified. "
+                        "Flagging will proceed on residual "
+                        "data.")
+
+    return p
+
+
+if __name__ == "__main__":
+    args = create_parser().parse_args()
+    if "/" in args.data_column:
+        data_string = args.data_column.split("/")
+        data_column = data_string[0]
+        string_columns = data_string[1]
+    else:
+        data_column = args.data_column
+        string_columns = None
+
+    print(data_column)
+    print(string_columns)
+
+    columns = [data_column,
+               "FLAG",
+               "TIME",
+               "ANTENNA1",
+               "ANTENNA2"]
+
+    if string_columns is not None and "(" in string_columns:
+        match = re.search(r'\(([A-Za-z0-9+-_]+)\)', string_columns)
+        multi_columns = re.split(r'\+|-', match.group(1))
+        multi_columns = list(filter(None, multi_columns))
+        print(str(multi_columns))
+        columns.extend(multi_columns)
+    else:
+        if string_columns is not None:
+            columns.append(string_columns)
+    print(columns)
+# Handles variety of cases
+# -dc DATA
+# -dc DATA/DATA_MODEL
+# -dc DATA/(DATA_MODEL)
+# -dc DATA/(DATA1 + DATA2 + DATA3)
+# -dc DATA/(-DATA1 + DATA2 + DATA3)
+# -dc DATA/(DATA1 + DATA2 - DATA3)