Merge pull request #25 from SpheMakh/flagstats

Flagstats

.travis.yml

@@ -12,7 +12,7 @@ services:
 before_install:
 - true
 install:
-- docker build -t msutils .
+- docker build . -f .travis/${TARGET}.docker
 script:
 - true
 deploy:

.travis/mypy.docker

@@ -1,4 +1,4 @@
-FROM kernsuite/base:4
+FROM kernsuite/base:5
 RUN docker-apt-install python-pip git
 RUN pip install -U mypy pip setuptools
 ENV USER root

.travis/py2.docker

@@ -1,4 +1,4 @@
-FROM kernsuite/base:4
+FROM kernsuite/base:5
 RUN docker-apt-install python-pip git
 ADD . /code
 ENV USER root

.travis/py3.docker

@@ -1,4 +1,4 @@
-FROM kernsuite/base:4
+FROM kernsuite/base:5
 RUN docker-apt-install python3-pip git
 ADD . /code
 ENV USER root

Dockerfile

@@ -1,7 +1,4 @@
-FROM kernsuite/base:3
-RUN docker-apt-install python-casacore \ 
-    meqtrees
-RUN docker-apt-install python-pip
-RUN pip install -U pip pyyaml
+FROM kernsuite/base:5
+RUN docker-apt-install python3-casacore python3-pip
 ADD . /msutils-src
-RUN pip install /msutils-src
+RUN pip3 install /msutils-src

MSUtils/__init__.py

@@ -1 +1 @@
-import msutils
+from MSUtils import msutils

MSUtils/flag_stats.py

@@ -0,0 +1,109 @@
+#from contextlib import ExitStack
+#import bokeh
+from daskms import xds_from_ms, xds_to_table, xds_from_table
+import dask
+import dask.array as da
+import numpy
+import sys
+
+
+def _get_flags(names, antenna1, antenna2, flags):
+    names = names[0]
+    # chan and corr are assumed to have a single chunk
+    # so we contract twice to access this single ndarray
+    flags = flags[0][0]
+    nant = len(names)
+    fracs = numpy.zeros([nant,2], dtype=numpy.float64)
+    for i in range(nant):
+        flag_sum = flags[numpy.logical_or(antenna1==i, antenna2==i)]
+        fracs[i,0] += flag_sum.sum()
+        fracs[i,1] += numpy.product(flag_sum.shape)
+    return fracs
+
+def _chunk(x, keepdims, axis):
+    return x
+
+def _combine(x, keepdims, axis):
+    if isinstance(x, list):
+        return sum(x)
+    elif isinstance(x, numpy.ndarray):
+        return x
+    else:
+        raise TypeError("Invalid type %s" % type(x))
+
+def _aggregate(x, keepdims, axis):
+    return _combine(x, keepdims, axis)
+
+def antenna_flags_field(msname, fields=None, antennas=None):
+    ds_mss = xds_from_ms(msname, group_cols=["FIELD_ID", "DATA_DESC_ID"], chunks={'row': 100000})
+    ds_ant = xds_from_table(msname+"::ANTENNA")[0]
+    ds_field = xds_from_table(msname+"::FIELD")[0]
+
+    ant_names = ds_ant.NAME.data.compute()
+    field_names = ds_field.NAME.data.compute()
+
+    if fields:
+        if isinstance(fields[0], str):
+            field_ids = map(fields.indexof, fields)
+        else:
+            field_ids = fields
+    else:
+        field_ids = range(len(field_names))
+
+    if antennas:
+        if isinstance(antennas[0], str):
+            antennas_ids = map(antennas.indexof, antennas)
+        else:
+            ant_ids = antennas
+    else:
+        ant_ids = range(len(ant_names))
+
+    nant = len(ant_ids)
+    nfield = len(field_ids)
+
+    flag_sum_computes = []
+    for ds in ds_mss:
+        if ds.FIELD_ID not in field_ids:
+            continue
+        flag_sums = da.blockwise(_get_flags, ("row",),
+                                    ant_ids, ("ant",),
+                                    ds.ANTENNA1.data, ("row",),
+                                    ds.ANTENNA2.data, ("row",),
+                                    ds.FLAG.data, ("row","chan", "corr"),
+                                    adjust_chunks={"row": nant },
+                                    dtype=numpy.ndarray)
+
+        flags_redux = da.reduction(flag_sums,
+                                 chunk=_chunk,
+                                 combine=_combine,
+                                 aggregate=_aggregate,
+                                 concatenate=False,
+                                 dtype=numpy.float64)
+        flag_sum_computes.append(flags_redux)
+
+    #flag_sum_computes[0].visualize("graph.pdf")
+    sum_per_field_spw = dask.compute(flag_sum_computes)[0]
+    sum_all = sum(sum_per_field_spw)
+    fractions = sum_all[:,0]/sum_all[:,1]
+    stats = {}
+    for i,aid in enumerate(ant_ids):
+        ant_stats = {}
+        ant_stats["name"] = ant_names[aid]
+        ant_stats["frac"] = fractions[i]
+        ant_stats["sum"] = sum_all[i][0]
+        ant_stats["counts"] = sum_all[i][1]
+        stats[aid] = ant_stats
+
+    return stats
+
+#with ExitStack() as stack:
+#    from dask.diagnostics import Profiler, visualize
+#    prof = Profiler()
+#
+#    stack.enter_context(prof)
+#    result = dask.compute(writes)[0]
+#    print(sum(result))
+#
+#    import pdb; pdb.set_trace()
+#    
+#    visualize(prof)

MSUtils/msutils.py

@@ -130,7 +130,7 @@ def summary(msname, outfile=None, display=True):
     tab.close()
 
     if display:
-        print info
+        print(info)
 
     if outfile:
         with codecs.open(outfile, 'w', 'utf8') as stdw:

requirements.txt

@@ -1,4 +1,4 @@
 numpy
 scipy
 matplotlib
-#python-casacore
+dask-ms

setup.py

@@ -51,7 +51,7 @@ def src_pkg_dirs(pkg_name):
     return pkg_dirs
 
 setup(name='msutils',
-      version="0.9.7",
+      version="1.0.0",
       description='Tools for playing with Measurement sets.',
       long_description=readme(),
       url='https://github.com/SpheMakh/msutils',