From 5b71ac9307fdd15c0c31af908b7c6775de24a45b Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Tue, 19 Oct 2021 23:09:32 +0200 Subject: [PATCH 1/2] [skip ci] WIP: Return chan_freq per BDA output_row --- africanus/averaging/bda_mapping.py | 14 ++++++++++++-- africanus/averaging/dask.py | 6 ++++-- africanus/averaging/tests/test_bda_averaging.py | 4 +++- africanus/averaging/tests/test_bda_mapping.py | 15 +++++++++++---- 4 files changed, 30 insertions(+), 9 deletions(-) diff --git a/africanus/averaging/bda_mapping.py b/africanus/averaging/bda_mapping.py index 34d5de625..c40c98bc7 100644 --- a/africanus/averaging/bda_mapping.py +++ b/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 diff --git a/africanus/averaging/dask.py b/africanus/averaging/dask.py index 6f8937b0a..1153533e6 100644 --- a/africanus/averaging/dask.py +++ b/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, diff --git a/africanus/averaging/tests/test_bda_averaging.py b/africanus/averaging/tests/test_bda_averaging.py index a61e0eeda..d7863cd32 100644 --- a/africanus/averaging/tests/test_bda_averaging.py +++ b/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:]] diff --git a/africanus/averaging/tests/test_bda_mapping.py b/africanus/averaging/tests/test_bda_mapping.py index fc0508695..42358a467 100644 --- a/africanus/averaging/tests/test_bda_mapping.py +++ b/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): From 9ddfafdd45829837536ae4ab059f46db091a1e4f Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Wed, 20 Oct 2021 08:32:42 +0200 Subject: [PATCH 2/2] [skip ci] Update HISTORY.rst --- HISTORY.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/HISTORY.rst b/HISTORY.rst index c2e5db205..0f4a404e8 100644 --- a/HISTORY.rst +++ b/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`)