diff --git a/src/xarray_kat/array.py b/src/xarray_kat/array.py index 757c3bc..4737967 100644 --- a/src/xarray_kat/array.py +++ b/src/xarray_kat/array.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from numbers import Integral -from typing import TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING, Any, Dict, Tuple import numpy as np import numpy.typing as npt @@ -17,9 +17,13 @@ explicit_indexing_adapter, ) +from xarray_kat.calibration import calc_correction_per_antenna + if TYPE_CHECKING: from rarg_python_patterns.multiton import Multiton + from xarray_kat.katdal_types import TelstateDataProducts + class AbstractMeerkatArchiveArray(ABC, BackendArray): """Require subclasses to implement ``dims`` and ``chunks`` properties. @@ -306,3 +310,87 @@ def _vindex_get(self, indexer): def __getitem__(self, key): return super().__getitem__(key).get_duck_array().read().result() + + +class CalibrationBackendArray(BackendArray): + __slots__ = ( + "_data_products", + "_timestamps", + "_antenna_map", + "_frequencies", + "shape", + "dtype", + ) + + _data_products: Multiton[TelstateDataProducts] + _timestamps: npt.NDArray + _antenna_map: Dict[str, int] + _frequencies: npt.NDArray + shape: Tuple[int, int, int, int] + dtype: npt.DTypeLike + + def __init__(self, data_products: Multiton[TelstateDataProducts]): + self._data_products = data_products + dp = data_products.instance + self._timestamps = dp.timestamps + self._antenna_map = {a.name: i for i, a in enumerate(dp.antennas)} + self._frequencies = dp.frequencies + self.dtype = np.dtype("complex64") + self.shape = ( + self._timestamps.shape[0], + len(self._antenna_map), + self._frequencies.shape[0], + 4, + ) + + def __reduce__(self): + return (CalibrationBackendArray, (self._data_products,)) + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, CalibrationBackendArray): + return NotImplemented + + return self._data_products == other._data_products + + def __getitem__(self, key): + return explicit_indexing_adapter( + key, + self.shape, + IndexingSupport.OUTER, + self.generate_calibration_solutions, + ) + + def generate_calibration_solutions(self, key): + key_arrays = [] + + for sel, dim_size in zip(key, self.shape, strict=True): + if isinstance(sel, slice): + key_arrays.append(np.arange(*sel.indices(dim_size))) + elif isinstance(sel, Integral): + key_arrays.append(np.array([sel])) + elif isinstance(sel, np.ndarray) and sel.ndim == 1: + key_arrays.append(sel) + else: + raise TypeError( + f"{tuple(map(type, key))} element types " + f"should be slice, integer or 1D ndarray" + ) + + key_shape = tuple(map(len, key_arrays)) + time_index, ant_index, chan_index, pol_index = key_arrays + + # Compute solutions for all channels + chan_slice = slice(0, len(self._frequencies)) + + solutions = np.ones(key_shape, self.dtype) + cal_params = self._data_products.instance.calibration_params + + # Expand to full calibration solutions for each timestamp + # the slice out the selection in the other dimensions + for t, dump in enumerate(time_index): + antenna_calsols = calc_correction_per_antenna( + dump, chan_slice, self._antenna_map, cal_params + ) + solutions[t, ...] = antenna_calsols[np.ix_(ant_index, chan_index, pol_index)] + + return solutions diff --git a/src/xarray_kat/calibration.py b/src/xarray_kat/calibration.py new file mode 100644 index 0000000..e2abbcf --- /dev/null +++ b/src/xarray_kat/calibration.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from bisect import insort +from collections import defaultdict +from itertools import product +from typing import TYPE_CHECKING, Dict + +import numpy as np + +from xarray_kat.utils import ANTENNA_RECEPTOR_REGEX + +if TYPE_CHECKING: + from xarray_kat.katdal_types import CorrectionParams + + +def calc_correction_per_antenna( + dump: int, channels: slice, antenna_map: Dict[str, int], params: CorrectionParams +): + """Prototype: gain correction per channel per *input* for a given dump. + + This is the per-antenna-feed counterpart to + :func:`calc_correction_per_corrprod`. It performs the same accumulation of + all requested calibration products into a gain per input (where an input is + an antenna-pol feed such as ``m000h``), but stops short of folding pairs of + inputs into correlation products. The returned gains can be combined into + per-corrprod gains as ``g[:, i] * conj(g[:, j])`` for a baseline of inputs + ``i`` and ``j``. + + Parameters + ---------- + dump : int + Dump index (applicable to full data set, i.e. absolute) + channels : slice + Channel indices (applicable to full data set, i.e. absolute) + antenna_map : dict + ``{antenna_name: index}`` dictionary mapping antenna names + into the antenna dimension. + params : :class:`CorrectionParams` + Corrections per input + + Returns + ------- + gains : array of complex64, shape (n_antenna, n_chans, n_pol) + Gain corrections per channel per input, columns ordered as + ``params.inputs`` + """ + + # Build a map of the form {"mxyz: [("h", i), ("v", j)]"} + # where mxyz is the antenna number and i and j + # are the indices in params.inputs of m000xh and m000xv respectively + antenna_receptor_map: defaultdict[str, list[tuple[str, int]]] = defaultdict(list) + + for i, antenna_receptor in enumerate(params.inputs): + if ( + (m := ANTENNA_RECEPTOR_REGEX.match(antenna_receptor)) is None + or (prefix := m.group("prefix")) is None + or (number := m.group("number")) is None + or (receptor := m.group("receptor")) is None + ): + raise ValueError(f"Invalid antenna receptor string '{antenna_receptor}'") + + insort(antenna_receptor_map[f"{prefix}{number}"], (receptor, i)) + + # This should always hold, but perform a sanity check + for ( + antenna, + receptor_map, + ) in antenna_receptor_map.items(): + if (ant_receptors := tuple(r for r, _ in receptor_map)) != ("h", "v"): + raise ValueError( + f"{antenna} does not have corrections for both " + f"h and v receptors: {ant_receptors}" + ) + + nant = len(antenna_map) + npol = 4 # Follows from the above invariant + nchan = channels.stop - channels.start + gains = np.ones((nant, nchan, npol), dtype="complex64") + + # Iterate over each gain type (cal_product) + for cal_product, product_corrections in params.corrections.items(): + channel_map = params.channel_maps[cal_product] + + # Construct each antenna's gains from it's feed receptors + for antenna_name, ((_, hi), (_, vi)) in antenna_receptor_map.items(): + a = antenna_map[antenna_name] + h_corrections = product_corrections[hi][dump] + v_corrections = product_corrections[vi][dump] + h_gains = channel_map(h_corrections, channels) + v_gains = channel_map(v_corrections, channels) + args = ((h_gains,), (v_gains,)) + + # Multiply feed receptor values into the net gain for each polarisation + for p, ((r1_gains,), (r2_gains,)) in enumerate(product(*(args, args))): + gains[a, :, p] *= r1_gains + gains[a, :, p] *= np.conj(r2_gains) + + return gains diff --git a/src/xarray_kat/datatree_factory.py b/src/xarray_kat/datatree_factory.py index 7bff72d..df2df67 100644 --- a/src/xarray_kat/datatree_factory.py +++ b/src/xarray_kat/datatree_factory.py @@ -342,13 +342,10 @@ def WrappedArray(a): return LazilyIndexedArray(a) telstate = self._data_products.instance.telstate - chunk_info = telstate["chunk_info"] # Time metadata - start_time = telstate["sync_time"] + telstate["first_timestamp"] - ntime = chunk_info["correlator_data"]["shape"][0] + timestamps = self._data_products.instance.timestamps integration_time = telstate["int_time"] - timestamps = start_time + np.arange(ntime) * integration_time # Observation information start_utc = calendar.timegm(time.gmtime(timestamps[0])) @@ -360,11 +357,9 @@ def WrappedArray(a): # Frequency metadata band = telstate["sub_band"] - nchan = telstate["n_chans"] - bandwidth = telstate["bandwidth"] center_freq = telstate["center_freq"] - channel_width = bandwidth / nchan - chan_freqs = (center_freq - (bandwidth / 2)) + np.arange(nchan) * channel_width + frequencies = self._data_products.instance.frequencies + channel_width = self._data_products.instance.channel_width # Antenna metadata antennas = self._data_products.instance.antennas @@ -403,7 +398,7 @@ def WrappedArray(a): end_iso=end_iso, observer=observer, experiment_id=experiment_id, - chan_freqs=chan_freqs, + chan_freqs=frequencies, band=band, center_freq=center_freq, channel_width=channel_width, diff --git a/src/xarray_kat/katdal_types.py b/src/xarray_kat/katdal_types.py index 7305d01..b7b3299 100644 --- a/src/xarray_kat/katdal_types.py +++ b/src/xarray_kat/katdal_types.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, List +import numpy as np from katpoint import Antenna from xarray_kat.third_party.vendored.katdal.datasources_minimal import ( @@ -16,6 +17,7 @@ from xarray_kat.third_party.vendored.katdal.visdatav4_minimal import VisibilityDataV4 if TYPE_CHECKING: + import numpy.typing as npt from katsdptelstate import TelescopeState from rarg_python_patterns.multiton import Multiton @@ -44,6 +46,31 @@ def telstate(self) -> TelescopeState: """Return the TelescopeState""" return self.datasource.telstate + @property + def timestamps(self) -> npt.NDarray: + """Return the timestamps for each dump in the observation""" + ts = self.telstate + chunk_info = ts["chunk_info"] + start_time = ts["sync_time"] + ts["first_timestamp"] + ntime = chunk_info["correlator_data"]["shape"][0] + integration_time = ts["int_time"] + return start_time + np.arange(ntime) * integration_time + + @property + def channel_width(self) -> float: + """Returns the channel width for the current Spectral Window""" + ts = self.telstate + return ts["bandwidth"] / ts["n_chans"] + + @property + def frequencies(self) -> npt.NDArray: + """Return the frequencies for the current Spectral Window""" + ts = self.telstate + nchan = ts["n_chans"] + bandwidth = ts["bandwidth"] + center_freq = ts["center_freq"] + return (center_freq - (bandwidth / 2)) + np.arange(nchan) * self.channel_width + @property def antennas(self) -> List[Antenna]: """Return a list of katpoint Antenna""" diff --git a/src/xarray_kat/utils/__init__.py b/src/xarray_kat/utils/__init__.py index 2c32792..f5812e9 100644 --- a/src/xarray_kat/utils/__init__.py +++ b/src/xarray_kat/utils/__init__.py @@ -104,7 +104,7 @@ def normalize_chunks( ANTENNA_RECEPTOR_REGEX = re.compile( - r"(?P[mMsSeE])(?P\d+)(?P[hHvV])" + r"(?P[mMsSeE])(?P\d+)(?P[hHvV])" ) @@ -118,14 +118,14 @@ def corrprods_to_baseline_pols(corrprods: npt.NDArray): (cp1_match := re.match(ANTENNA_RECEPTOR_REGEX, cp1)) is None or (cp1_prefix := cp1_match.group("prefix")) is None or (cp1_nr := cp1_match.group("number")) is None - or (cp1_pol := cp1_match.group("polarization")) is None + or (cp1_receptor := cp1_match.group("receptor")) is None ): raise ValueError(f"{cp1} is not a valid correlation product string {cp1_match}") if ( (cp2_match := re.match(ANTENNA_RECEPTOR_REGEX, cp2)) is None or (cp2_prefix := cp2_match.group("prefix")) is None or (cp2_nr := cp2_match.group("number")) is None - or (cp2_pol := cp2_match.group("polarization")) is None + or (cp2_receptor := cp2_match.group("receptor")) is None ): raise ValueError(f"{cp2} is not a valid correlation product string {cp2_match}") @@ -133,7 +133,7 @@ def corrprods_to_baseline_pols(corrprods: npt.NDArray): ( f"{cp1_prefix.lower()}{cp1_nr}", f"{cp2_prefix.lower()}{cp2_nr}", - f"{cp1_pol.lower()}{cp2_pol.lower()}", + f"{cp1_receptor.lower()}{cp2_receptor.lower()}", ) ) diff --git a/tests/conftest.py b/tests/conftest.py index f8644f1..6443929 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -413,7 +413,7 @@ def add_calibration_solutions(self) -> None: written into the telstate. This allows TelstateDataProducts to find a non-None calibration_params when constructed with applycal="all". - The solutions use unit-amplitude complex gains with random phases so + The solutions use complex gains with random amplitudes and phases so that the correction is well-defined but non-trivial. Example: @@ -432,7 +432,7 @@ def _add_cal_data(self, telstate: katsdptelstate.TelescopeState) -> None: spectral info, product_B_parts). - Per-capture-block mutable sensor keys for G, K, and B0 products. - All solutions use unit amplitude with random phases (seeded for + All solutions use random amplitudes and phases (seeded for reproducibility), so they are invertible but non-trivial. """ n_pols = 2 # 'h' and 'v' @@ -461,11 +461,14 @@ def _add_cal_data(self, telstate: katsdptelstate.TelescopeState) -> None: last_dump_ts = self.sync_time + self.ntime * self.int_time rng = np.random.default_rng(42) - # G (gain): shape (n_pols, n_ants), complex64, unit amplitude. + # G (gain): shape (n_pols, n_ants), complex64, amplitude away from unity. # Three solutions with different phases to exercise time interpolation. + # Non-unit amplitudes ensure autocorrelation pols (hh, vv) are not + # trivially 1.0, so gain magnitude handling is actually exercised. for ts in (first_dump_ts, mid_dump_ts, last_dump_ts): phases_G = rng.uniform(-np.pi, np.pi, (n_pols, n_ants)) - gains_G = np.exp(1j * phases_G).astype(np.complex64) + amps_G = rng.uniform(0.5, 2.0, (n_pols, n_ants)) + gains_G = (amps_G * np.exp(1j * phases_G)).astype(np.complex64) telstate.add( f"{self.capture_block_id}_cal_product_G", gains_G, @@ -484,11 +487,13 @@ def _add_cal_data(self, telstate: katsdptelstate.TelescopeState) -> None: immutable=False, ) - # B0 (bandpass): shape (n_chans, n_pols, n_ants), complex64, unit amplitude. + # B0 (bandpass): shape (n_chans, n_pols, n_ants), complex64, amplitude + # away from unity (see G above for rationale). # Two solutions (step function in time — no interpolation, zeroth-order hold). for ts in (first_dump_ts, mid_dump_ts): phases_B = rng.uniform(-np.pi, np.pi, (n_freqs, n_pols, n_ants)) - gains_B = np.exp(1j * phases_B).astype(np.complex64) + amps_B = rng.uniform(0.5, 2.0, (n_freqs, n_pols, n_ants)) + gains_B = (amps_B * np.exp(1j * phases_B)).astype(np.complex64) telstate.add( f"{self.capture_block_id}_cal_product_B0", gains_B, diff --git a/tests/test_cal_params.py b/tests/test_cal_params.py new file mode 100644 index 0000000..de376d3 --- /dev/null +++ b/tests/test_cal_params.py @@ -0,0 +1,206 @@ +"""Tests that opening a mock archive with applycal yields a valid CorrectionParams. + +A mock katdal archive is built with synthetic calibration solutions and opened +through the public ``xarray.open_datatree`` path. The ``TelstateDataProducts`` +instance created along the way is recovered from ``Multiton._INSTANCE_CACHE`` +and its ``calibration_params`` (a :class:`CorrectionParams`) is validated. +""" + +import pickle +from collections import defaultdict + +import numpy as np +import pytest +import xarray +from pytest_httpserver import HTTPServer +from rarg_python_patterns import Multiton +from xarray.core.indexing import LazilyIndexedArray + +from tests.conftest import ( + SyntheticObservation, + setup_mock_archive_server, +) +from xarray_kat.array import CalibrationBackendArray +from xarray_kat.calibration import calc_correction_per_antenna +from xarray_kat.katdal_types import TelstateDataProducts, TelstateDataSource +from xarray_kat.third_party.vendored.katdal.applycal_minimal import ( + CorrectionParams, + calc_correction_per_corrprod, +) +from xarray_kat.utils import ANTENNA_RECEPTOR_REGEX + +CBID = "1234567890" +NTIME = 25 +NFREQ = 16 +NANTS = 4 + + +def _build_archive(httpserver: HTTPServer, tmp_path, *, applycal, with_cal): + """Build a mock archive, and serve it""" + obs = SyntheticObservation(CBID, ntime=NTIME, nfreq=NFREQ, nants=NANTS) + obs.add_scan(range(0, 8), "track", "PKS1934") + obs.add_scan(range(8, 20), "scan", "3C286") + obs.add_scan(range(20, NTIME), "track", "PKS1934") + if with_cal: + obs.add_calibration_solutions() + obs.save_to_directory(tmp_path) + + setup_mock_archive_server(httpserver, tmp_path, CBID, require_auth=False) + return obs, f"{httpserver.url_for('/')}{CBID}/{CBID}_sdp_l0.full.rdb" + + +def _open_archive(httpserver: HTTPServer, tmp_path, *, applycal, with_cal): + """Build a mock archive, serve it, open it and return the TelstateDataProducts.""" + obs, rdb_url = _build_archive( + httpserver, tmp_path, applycal=applycal, with_cal=with_cal + ) + + # Open via the public entrypoint; this constructs (and caches) the + # TelstateDataProducts as a side effect. + xarray.open_datatree(rdb_url, engine="xarray-kat", applycal=applycal) + + # Recover the TelstateDataProducts instance from the multiton cache. The + # clear_multitons autouse fixture wipes the cache around every test, so the + # only instance present is the one created by the open_datatree call above. + products = [ + entry[0] + for entry in Multiton._INSTANCE_CACHE.values() + if isinstance(entry[0], TelstateDataProducts) + ] + assert len(products) == 1, f"expected one TelstateDataProducts, found {len(products)}" + return obs, products[0] + + +def test_cal_params(httpserver: HTTPServer, tmp_path): + """Opening a mock archive with applycal='all' builds a valid CorrectionParams.""" + obs, tdp = _open_archive(httpserver, tmp_path, applycal="all", with_cal=True) + params = tdp.calibration_params + + assert isinstance(params, CorrectionParams) + + # Cal products were discovered and recorded (l1.{G,K,B} in some order). + applied = tdp._dataset.applycal_products + assert applied, "no cal products were applied" + assert all(p.startswith("l1.") for p in applied) + + # inputs: one entry per antenna-pol (h/v) feed, sorted. + inputs = params.inputs + assert inputs == sorted(inputs) + assert len(inputs) == NANTS * 2 + assert all(inp[:-1] in obs.ant_names and inp[-1] in ("h", "v") for inp in inputs) + + # input1_index / input2_index: integer indices into `inputs`, one per + # correlation product, all within range. + n_corrprods = len(obs.bls_ordering) + for idx in (params.input1_index, params.input2_index): + idx = np.asarray(idx) + assert idx.shape == (n_corrprods,) + assert np.issubdtype(idx.dtype, np.integer) + assert idx.min() >= 0 and idx.max() < len(inputs) + + # corrections / channel_maps are keyed by the applied products; each + # corrections entry has one (dump-indexable) sequence per input, and each + # channel map is callable. + assert set(params.corrections) == set(applied) + assert set(params.channel_maps) == set(applied) + for product in applied: + per_input = params.corrections[product] + assert len(per_input) == len(inputs) + first = per_input[0] + assert first[0] is not None # indexable by dump + assert callable(params.channel_maps[product]) + + # Functional smoke test: a per-corrprod correction for the first dump over + # the full band has the expected shape and dtype. + gains = calc_correction_per_corrprod(0, slice(0, NFREQ), params) + assert gains.shape == (NFREQ, n_corrprods) + assert gains.dtype == np.complex64 + + +def test_calc_correction_per_antenna(httpserver: HTTPServer, tmp_path): + """The per-antenna prototype reproduces the per-corrprod autocorrelation gains. + + ``calc_correction_per_antenna`` only forms the four pol products of an + antenna's own h/v feeds, so it corresponds exactly to the autocorrelation + baselines of ``calc_correction_per_corrprod``; cross-antenna baselines are + not representable from a single antenna's feeds. + """ + _, tdp = _open_archive(httpserver, tmp_path, applycal="all", with_cal=True) + params = tdp.calibration_params + + # Group inputs by antenna into their (h, v) feed indices. Sorted antenna + # order matches the antenna axis of per_antenna (which iterates the sorted + # params.inputs). + antenna_receptors: defaultdict[str, dict[str, int]] = defaultdict(dict) + for idx, feed in enumerate(params.inputs): + if (m := ANTENNA_RECEPTOR_REGEX.match(feed)) is None: + raise ValueError(f"Invalid antenna receptor string {antenna_receptors}") + + antenna = f"{m.group('prefix')}{m.group('number')}" + receptor = m.group("receptor") + antenna_receptors[antenna][receptor] = idx + antenna_map = {k: i for i, k in enumerate(antenna_receptors.keys())} + assert len(antenna_map) == NANTS + + channels = slice(0, NFREQ) + per_antenna = calc_correction_per_antenna(0, channels, antenna_map, params) + per_corrprod = calc_correction_per_corrprod(0, channels, params) + + assert per_antenna.shape == (NANTS, NFREQ, 4) + assert per_antenna.dtype == np.complex64 + + # per_antenna[a, :, p] = g_in[ant_r1] * conj(g_in[ant_r2]) for the antenna's + # own feeds, with pols ordered (hh, hv, vh, vv). Locate the matching + # autocorrelation corrprod (same input pair) and compare. + input1 = np.asarray(params.input1_index) + input2 = np.asarray(params.input2_index) + for ant, a in antenna_map.items(): + hi, vi = antenna_receptors[ant]["h"], antenna_receptors[ant]["v"] + pol_inputs = [(hi, hi), (hi, vi), (vi, hi), (vi, vi)] # hh, hv, vh, vv + for p, (in1, in2) in enumerate(pol_inputs): + (matches,) = np.where((input1 == in1) & (input2 == in2)) + assert len(matches) == 1, f"expected one corrprod for {ant} pol {p}" + np.testing.assert_allclose( + per_antenna[a, :, p], per_corrprod[:, matches[0]], rtol=1e-6, atol=1e-6 + ) + + +@pytest.mark.parametrize("applycal", ["all"]) +def test_calibration_array_reduce(httpserver, tmp_path, applycal): + _, rdb_url = _build_archive(httpserver, tmp_path, applycal=applycal, with_cal=True) + + datasource = Multiton( + TelstateDataSource.from_url, + rdb_url, + chunk_store=None, + capture_block_id=CBID, + ) + + data_products = Multiton(TelstateDataProducts, datasource, applycal=applycal) + array = CalibrationBackendArray(data_products) + + assert array == pickle.loads(pickle.dumps(array)) + + +@pytest.mark.parametrize("applycal", ["all"]) +def test_calibration_array_backend(httpserver, tmp_path, applycal): + _, rdb_url = _build_archive(httpserver, tmp_path, applycal=applycal, with_cal=True) + + datasource = Multiton( + TelstateDataSource.from_url, + rdb_url, + chunk_store=None, + capture_block_id=CBID, + ) + + data_products = Multiton(TelstateDataProducts, datasource, applycal=applycal) + array = LazilyIndexedArray(CalibrationBackendArray(data_products)) + + dataset = xarray.Dataset( + {"DATA": (("time", "antenna_name", "frequency", "polarization"), array)} + ) + + subds = dataset.isel({"time": slice(0, 10), "frequency": [1, 4, 8]}) + subds.load() + + assert {"time": 10, "frequency": 3}.items() <= dict(subds.sizes).items()