Expose chan_freq in BDA output

HISTORY.rst

@@ -4,6 +4,7 @@ History
 
 0.3.1 (2021-09-09)
 ------------------
+* Expose chan_freq in BDA output (:pr:`260`)
 * Restrict numba version to <= 0.54.0 (:pr:`259`)
 * Handle empty spectral indices in WSClean Model (:pr:`258`)
 * Support missing data during BDA (:pr:`252`)

africanus/averaging/bda_mapping.py

@@ -3,6 +3,7 @@
 from collections import namedtuple
 
 import numpy as np
+from africanus.rime.fast_beam_cubes import freq_grid_interp
 import numba
 from numba.experimental import jitclass
 import numba.types
@@ -335,7 +336,8 @@ def finalise_bin(self, auto_corr, uvw, time, interval,
 
 RowMapOutput = namedtuple("RowMapOutput",
                           ["map", "offsets", "decorr_chan_width",
-                           "time", "interval", "chan_width", "flag_row"])
+                           "time", "interval", "chan_freq", "chan_width",
+                           "flag_row"])
 
 
 @generated_jit(nopython=True, nogil=True, cache=True)
@@ -571,6 +573,9 @@ def update_lookups(finalised, bl):
         time_ret = np.full(out_row_chans, -1, dtype=time.dtype)
         int_ret = np.full(out_row_chans, -1, dtype=interval.dtype)
         chan_width_ret = np.full(out_row_chans, 0, dtype=chan_width.dtype)
+        chan_freq_ret = np.full(out_row_chans, 0, dtype=chan_freq.dtype)
+        freq_const = ((chan_freq[-1] - chan_freq[0]) / (nchan - 1)
+                      if nchan > 1 else 1)
 
         # Construct output flag row, if necessary
         out_flag_row = (None if flag_row is None else
@@ -618,6 +623,9 @@ def update_lookups(finalised, bl):
                 time_ret[out_offset] = bin_time
                 int_ret[out_offset] = bin_interval
 
+                # Compute channel frequency for each row
+                chan_freq_ret[out_offset] = chan_freq[0] + c*freq_const
+
                 # Add channel contribution for each row
                 chan_width_ret[out_offset] += chan_width[c]
 
@@ -627,6 +635,8 @@ def update_lookups(finalised, bl):
         return RowMapOutput(row_chan_map, offsets,
                             decorr_chan_width,
                             time_ret, int_ret,
-                            chan_width_ret, out_flag_row)
+                            chan_freq_ret,
+                            chan_width_ret,
+                            out_flag_row)
 
     return impl

africanus/averaging/dask.py

@@ -771,8 +771,9 @@ def bda(time, interval, antenna1, antenna2,
                                  meta, format=format)
     meta_interval = _bda_getitem_row(meta, 4, interval, ("row",),
                                      meta, format=format)
-    meta_chan_width = _bda_getitem_row(meta, 5, chan_width, ("row",), meta)
-    meta_flag_row = (_bda_getitem_row(meta, 6, flag_row, ("row",),
+    meta_chan_freq = _bda_getitem_row(meta, 5, chan_width, ("row",), meta)                            
+    meta_chan_width = _bda_getitem_row(meta, 6, chan_width, ("row",), meta)
+    meta_flag_row = (_bda_getitem_row(meta, 7, flag_row, ("row",),
                                       meta, format=format)
                      if flag_row is not None else None)
 
@@ -782,6 +783,7 @@ def bda(time, interval, antenna1, antenna2,
                             meta_decorr_cw,
                             meta_time,
                             meta_interval,
+                            meta_chan_freq,
                             meta_chan_width,
                             meta_flag_row,
                             row_data.antenna1,

africanus/averaging/tests/test_bda_averaging.py

@@ -185,7 +185,7 @@ def test_bda_avg(bda_test_map, inv_bda_test_map, flags):
 
     meta = RowMapOutput(bda_test_map, offsets,
                         chan_width, out_time, out_interval,
-                        None, out_flag_row)
+                        None, None, out_flag_row)
 
     ant1 = np.full(in_row, 0, dtype=np.int32)
     ant2 = np.full(in_row, 1, dtype=np.int32)
@@ -347,6 +347,8 @@ def test_dask_bda_avg(vis_format):
     dsk2_name = result2["visibilities"][0].name
     dsk3_name = result2["visibilities"][1].name
 
+    from pprint import pprint
+
     # For each task, compare the row dictionaries
     for k, v in dsk1.items():
         v2 = dsk2[(dsk2_name,) + k[1:]]

africanus/averaging/tests/test_bda_mapping.py

@@ -1,13 +1,13 @@
 # -*- coding: utf-8 -*-
 
 import numpy as np
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_equal, assert_array_almost_equal
 import pytest
 
 from africanus.averaging.bda_mapping import bda_mapper, Binner
 
 
-@pytest.fixture(scope="session", params=[4096])
+@pytest.fixture(scope="session", params=[16])
 def nchan(request):
     return request.param
 
@@ -205,11 +205,18 @@ def test_bda_mapper(time, synthesized_uvw, interval,
 
     # NUM_CHAN divides number of channels exactly
     num_chan = np.diff(row_meta.offsets)
-    _, remainder = np.divmod(chan_width.shape[0], num_chan)
-    assert np.all(remainder == 0)
+    div, mod = np.divmod(chan_width.shape[0], num_chan)
+    assert np.all(mod == 0)
     decorr_cw = chan_width.sum() / num_chan
     assert_array_equal(decorr_cw, row_meta.decorr_chan_width)
 
+    for s, e, d in zip(row_meta.offsets[:-1], row_meta.offsets[1:], div):
+        nchan = chan_freq.shape[0] // d
+        assert e - s == nchan
+        expected = np.linspace(chan_freq[0], chan_freq[-1], nchan)
+        # import pdb; pdb.set_trace()
+        assert_array_almost_equal(row_meta.chan_freq[s:e], expected)
+
 
 def test_bda_binner(time, interval, synthesized_uvw,
                     ref_freq, chan_freq, chan_width):