ddid fixes (#79)

* fix indexing the dataset ddids * computing meta fix - workaround for bug introduced upstream in dask 2020.12.0
ratt-ru · Feb 3, 2021 · 536ee52 · 536ee52
1 parent 4997adb
commit 536ee52
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 12 deletions.
diff --git a/tricolour/apps/tricolour/app.py b/tricolour/apps/tricolour/app.py
@@ -232,7 +232,7 @@ def support_tables(ms):
     support = {t: xds_from_table("::".join((ms, t)), group_cols="__row__")
                for t in ["FIELD", "POLARIZATION", "SPECTRAL_WINDOW"]}
     # These columns have fixed shapes
-    support.update({t: xds_from_table("::".join((ms, t)))
+    support.update({t: xds_from_table("::".join((ms, t)))[0]
                     for t in ["ANTENNA", "DATA_DESCRIPTION"]})
 
     # Reify all values upfront
@@ -305,11 +305,8 @@ def _main(args):
     spw_ds = st["SPECTRAL_WINDOW"]
     ant_ds = st["ANTENNA"]
 
-    assert len(ant_ds) == 1
-    assert len(ddid_ds) == 1
-
-    antspos = ant_ds[0].POSITION.data
-    antsnames = ant_ds[0].NAME.data
+    antspos = ant_ds.POSITION.data
+    antsnames = ant_ds.NAME.data
     fieldnames = [fds.NAME.data[0] for fds in field_ds]
 
     avail_scans = [ds.SCAN_NUMBER for ds in xds]
@@ -367,9 +364,8 @@ def _main(args):
                  "compute graph for processing"
                  .format(field_dict[ds.FIELD_ID], ds.SCAN_NUMBER))
 
-        ddid = ddid_ds[ds.attrs['DATA_DESC_ID']]
-        spw_info = spw_ds[ddid.SPECTRAL_WINDOW_ID.data[0]]
-        pol_info = pol_ds[ddid.POLARIZATION_ID.data[0]]
+        spw_info = spw_ds[ddid_ds.SPECTRAL_WINDOW_ID.data[ds.DATA_DESC_ID]]
+        pol_info = pol_ds[ddid_ds.POLARIZATION_ID.data[ds.DATA_DESC_ID]]
 
         nrow, nchan, ncorr = getattr(ds, data_column).data.shape
 

diff --git a/tricolour/window_statistics.py b/tricolour/window_statistics.py
@@ -123,21 +123,21 @@ def window_stats(flag_window, ubls, chan_freqs,
                          field_name, None,
                          ddid, None,
                          nchanbins, None,
-                         dtype=np.object)
+                         meta=np.empty((0,), dtype=np.object))
 
     # Create an empty stats object if the user hasn't supplied one
     if prev_stats is None:
         def _window_stat_creator():
             return WindowStatistics(nchanbins)
 
         prev_stats = da.blockwise(_window_stat_creator, (),
-                                  dtype=np.object)
+                                  meta=np.empty((), dtype=np.object))
 
     # Combine per-baseline stats into a single stats object
     return da.blockwise(_combine_baseline_window_stats, (),
                         stats, ("bl",),
                         prev_stats, (),
-                        dtype=np.object)
+                        meta=np.empty((), dtype=np.object))
 
 
 def _combine_window_stats(*args):