diff --git a/setup.py b/setup.py index 166b828..917ae42 100644 --- a/setup.py +++ b/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' ] diff --git a/tricolour/apps/tricolour/app.py b/tricolour/apps/tricolour/app.py index 836ee49..7515f8c 100644 --- a/tricolour/apps/tricolour/app.py +++ b/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"] diff --git a/tricolour/apps/tricolour/somthing.py b/tricolour/apps/tricolour/somthing.py new file mode 100644 index 0000000..332c06b --- /dev/null +++ b/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)