Export calibration solutions

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

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

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,

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"""

src/xarray_kat/utils/__init__.py

@@ -104,7 +104,7 @@ def normalize_chunks(
 
 
 ANTENNA_RECEPTOR_REGEX = re.compile(
-  r"(?P<prefix>[mMsSeE])(?P<number>\d+)(?P<polarization>[hHvV])"
+  r"(?P<prefix>[mMsSeE])(?P<number>\d+)(?P<receptor>[hHvV])"
 )
 
 
@@ -118,22 +118,22 @@ 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}")
 
     result.append(
       (
         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()}",
       )
     )
 

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,