From 632f2c7f4260de9546f8e0566a8457e9845949f0 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Mon, 16 Feb 2026 14:25:40 +0200 Subject: [PATCH 1/7] Field and source dataset --- src/xarray_kat/datatree_factory.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/xarray_kat/datatree_factory.py b/src/xarray_kat/datatree_factory.py index 3490b4f..7f7a6d7 100644 --- a/src/xarray_kat/datatree_factory.py +++ b/src/xarray_kat/datatree_factory.py @@ -186,6 +186,24 @@ def _build_antenna_dataset(self) -> Dataset: }, ) + def _build_field_and_source_dataset(self, target: Target) -> Dataset: + """Build a field and source dataset for a single scan""" + return Dataset( + data_vars={ + "FIELD_PHASE_CENTER_DIRECTION": Variable( + ("field_name", "sky_dir_label"), + [list(target.radec())], + {"type": "sky_coord", "units": "rad", "frame": "fk5"}, + ), + }, + coords={ + "field_name": Variable("field_name", [target.name]), + "sky_dir_label": Variable("sky_dir_label", ["ra", "dec"]), + "source_name": Variable("field_name", [target.name]), + }, + attrs={"type": "field_and_source"}, + ) + def _build_correlated_dataset( self, meta: ObservationMetadata, @@ -449,8 +467,11 @@ def WrappedArray(a): data_vars, ) + field_and_source_ds = self._build_field_and_source_dataset(target) + correlated_node_name = f"{self._data_products.instance.name}_{i:03d}" tree[correlated_node_name] = correlated_ds tree[f"{correlated_node_name}/antenna_xds"] = antenna_ds + tree[f"{correlated_node_name}/field_and_source_base_xds"] = field_and_source_ds return tree From 3466ad1df0a59e95fe9c2e41dec2eb431833510e Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Mon, 16 Feb 2026 15:08:18 +0200 Subject: [PATCH 2/7] Move repeated in-function default coordinate definitions into a global variable --- src/xarray_kat/datatree_factory.py | 32 ++++++++++++++++++++++++++---- tests/conftest.py | 30 ++++++++++------------------ 2 files changed, 39 insertions(+), 23 deletions(-) diff --git a/src/xarray_kat/datatree_factory.py b/src/xarray_kat/datatree_factory.py index 7f7a6d7..b356c76 100644 --- a/src/xarray_kat/datatree_factory.py +++ b/src/xarray_kat/datatree_factory.py @@ -279,6 +279,21 @@ def _build_correlated_dataset( }, ) + def _build_data_group(self, field_and_source_name: str): + """Build the correlated dataset base data group""" + return { + "data_groups": { + "base": { + "correlated_data": "VISIBILITY", + "description": "Data group associated with the VISIBILITY DataArray", + "field_and_source": field_and_source_name, + "date": datetime.now(timezone.utc).isoformat(), + "flag": "FLAG", + "weight": "WEIGHT", + } + } + } + def create(self) -> Dict[str, Dataset]: if self._chunks is not None: ArrayClass = DelayedBackendArray @@ -458,6 +473,8 @@ def WrappedArray(a): {"preferred_chunks": uvw_preferred_chunks}, ) + # Create the correlated dataset + base_path = f"{self._data_products.instance.name}_{i:03d}" correlated_ds = self._build_correlated_dataset( meta, scan_timestamps, @@ -467,11 +484,18 @@ def WrappedArray(a): data_vars, ) + # Create the field and source dataset field_and_source_ds = self._build_field_and_source_dataset(target) + field_and_source_node_path = f"{base_path}/field_and_source_base_xds" + + # Create the data on the correlated dataset + correlated_ds = correlated_ds.assign_attrs( + **self._build_data_group(field_and_source_node_path) + ) - correlated_node_name = f"{self._data_products.instance.name}_{i:03d}" - tree[correlated_node_name] = correlated_ds - tree[f"{correlated_node_name}/antenna_xds"] = antenna_ds - tree[f"{correlated_node_name}/field_and_source_base_xds"] = field_and_source_ds + # Assign to appropriate nodes in the tree + tree[base_path] = correlated_ds + tree[f"{base_path}/antenna_xds"] = antenna_ds + tree[field_and_source_node_path] = field_and_source_ds return tree diff --git a/tests/conftest.py b/tests/conftest.py index 3f1f38e..6907ef7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -30,6 +30,14 @@ logger = logging.getLogger(__name__) +# Format: "name, radec, RA, DEC" +DEFAULT_COORDS = { + "PKS1934": ("19:39:25.03", "-63:42:45.63"), + "3C286": ("13:31:08.29", "+30:30:33.0"), + "MockTarget": ("00:00:00.0", "+00:00:00.0"), +} + + @pytest.fixture(autouse=True) def clear_multitons(): with Multiton._INSTANCE_LOCK: @@ -270,13 +278,7 @@ def add_sensor_data( if target_name is None: target_str = "Nothing, special" else: - # Format: "name, radec, RA, DEC" - coords = { - "PKS1934": ("19:39:25.03", "-63:42:45.63"), - "3C286": ("13:31:08.29", "+30:30:33.0"), - "MockTarget": ("00:00:00.0", "+00:00:00.0"), - } - ra, dec = coords.get(target_name, ("00:00:00.0", "+00:00:00.0")) + ra, dec = DEFAULT_COORDS.get(target_name, ("00:00:00.0", "+00:00:00.0")) target_str = f"{target_name}, radec, {ra}, {dec}" telstate.add("obs_target", target_str, ts=scan_timestamp, immutable=False) @@ -306,12 +308,7 @@ def add_sensor_data( if target_name is None: target_str = "Nothing, special" else: - coords = { - "PKS1934": ("19:39:25.03", "-63:42:45.63"), - "3C286": ("13:31:08.29", "+30:30:33.0"), - "MockTarget": ("00:00:00.0", "+00:00:00.0"), - } - ra, dec = coords.get(target_name, ("00:00:00.0", "+00:00:00.0")) + ra, dec = DEFAULT_COORDS.get(target_name, ("00:00:00.0", "+00:00:00.0")) target_str = f"{target_name}, radec, {ra}, {dec}" # Add sensors for each antenna @@ -359,12 +356,7 @@ def add_sensor_data( if target_name is None: target_str = "Nothing, special" else: - coords = { - "PKS1934": ("19:39:25.03", "-63:42:45.63"), - "3C286": ("13:31:08.29", "+30:30:33.0"), - "MockTarget": ("00:00:00.0", "+00:00:00.0"), - } - ra, dec = coords.get(target_name, ("00:00:00.0", "+00:00:00.0")) + ra, dec = DEFAULT_COORDS.get(target_name, ("00:00:00.0", "+00:00:00.0")) target_str = f"{target_name}, radec, {ra}, {dec}" telstate.add( From bffa24350ab5a0ea242dac564538736d5275a69f Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Mon, 16 Feb 2026 15:19:09 +0200 Subject: [PATCH 3/7] pre-commit hooks --- tests/test_mock_http_server.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/test_mock_http_server.py b/tests/test_mock_http_server.py index 73a625f..84a7f23 100644 --- a/tests/test_mock_http_server.py +++ b/tests/test_mock_http_server.py @@ -12,6 +12,7 @@ import numpy as np import pytest import xarray +from katpoint import Target from pytest_httpserver import HTTPServer from tests.conftest import ( @@ -399,6 +400,31 @@ def test_scan_state_filtering(self, httpserver: HTTPServer, tmp_path): assert "VISIBILITY" in ds assert "time" in ds.dims + def test_field_and_source_xds_dataset(self, httpserver: HTTPServer, tmp_path): + from tests.conftest import DEFAULT_COORDS + + obs = SyntheticObservation("1234567890", ntime=16, nfreq=16, nants=4) + obs.add_scan(range(0, 8), "track", "PKS1934") + obs.add_scan(range(8, 16), "scan", "3C286") + obs.save_to_directory(tmp_path) + + token = setup_mock_archive_server( + httpserver, tmp_path, "1234567890", require_auth=True + ) + base_url = httpserver.url_for("/") + rdb_url = f"{base_url}1234567890/1234567890_sdp_l0.full.rdb?token={token}" + + child_fields = [sc["target_name"] for sc in obs.scan_configs] + + dt = xarray.open_datatree(rdb_url, engine="xarray-kat") + for i, child_name in enumerate(dt.children): + fns = dt[f"{child_name}/field_and_source_base_xds"] + assert all([child_fields[i] == fn for fn in fns.field_name]) + assert all([child_fields[i] == sn for sn in fns.source_name]) + coords = DEFAULT_COORDS[child_fields[i]] + ra, dec = Target(f"{child_fields[i]}, radec, {coords[0]}, {coords[1]}").radec() + np.testing.assert_allclose(fns.FIELD_PHASE_CENTER_DIRECTION.values, [[ra, dec]]) + def test_antenna_xds_dataset(self, httpserver: HTTPServer, tmp_path): """Test antenna_xds dataset structure, values, and scan invariance.""" from katpoint import Antenna as KatAntenna From 18bb6076cf5c267e02285e65b5fac7e137fa39f1 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Mon, 16 Feb 2026 18:15:26 +0200 Subject: [PATCH 4/7] Slight test case improvements --- tests/test_mock_http_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_mock_http_server.py b/tests/test_mock_http_server.py index 84a7f23..926766b 100644 --- a/tests/test_mock_http_server.py +++ b/tests/test_mock_http_server.py @@ -401,6 +401,7 @@ def test_scan_state_filtering(self, httpserver: HTTPServer, tmp_path): assert "time" in ds.dims def test_field_and_source_xds_dataset(self, httpserver: HTTPServer, tmp_path): + """Test field_and_source_xds structure and values""" from tests.conftest import DEFAULT_COORDS obs = SyntheticObservation("1234567890", ntime=16, nfreq=16, nants=4) @@ -419,11 +420,12 @@ def test_field_and_source_xds_dataset(self, httpserver: HTTPServer, tmp_path): dt = xarray.open_datatree(rdb_url, engine="xarray-kat") for i, child_name in enumerate(dt.children): fns = dt[f"{child_name}/field_and_source_base_xds"] + assert all(fns.sky_dir_label == ["ra", "dec"]) assert all([child_fields[i] == fn for fn in fns.field_name]) assert all([child_fields[i] == sn for sn in fns.source_name]) coords = DEFAULT_COORDS[child_fields[i]] ra, dec = Target(f"{child_fields[i]}, radec, {coords[0]}, {coords[1]}").radec() - np.testing.assert_allclose(fns.FIELD_PHASE_CENTER_DIRECTION.values, [[ra, dec]]) + np.testing.assert_allclose(fns.FIELD_PHASE_CENTER_DIRECTION, [[ra, dec]]) def test_antenna_xds_dataset(self, httpserver: HTTPServer, tmp_path): """Test antenna_xds dataset structure, values, and scan invariance.""" From de056790aefd67885b75869e65c7dee170814847 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Mon, 16 Feb 2026 18:16:48 +0200 Subject: [PATCH 5/7] [skip ci] Update changelog --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f720045..bc52683 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -3,6 +3,7 @@ Changelog X.Y.Z (DD-MM-YYY) ----------------- +* Add field and source dataset (:pr:`48`) * Add pre-computed UVW coordinates (:pr:`47`) * Test that xarray-kat and katdal produce the same data for the same data source (:pr:`45`) * Support ``antenna_xds`` dataset (:pr:`44`) From d3e98fa664c511f1161a5b04643c1c3232583048 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Tue, 17 Feb 2026 14:39:24 +0200 Subject: [PATCH 6/7] Support choosing the uvw sign convention --- src/xarray_kat/datatree_factory.py | 22 ++++++++++++++++------ src/xarray_kat/entrypoint.py | 7 ++++++- src/xarray_kat/xkat_types.py | 1 + tests/test_katdal.py | 15 +++++++++++---- 4 files changed, 34 insertions(+), 11 deletions(-) diff --git a/src/xarray_kat/datatree_factory.py b/src/xarray_kat/datatree_factory.py index b356c76..491ee18 100644 --- a/src/xarray_kat/datatree_factory.py +++ b/src/xarray_kat/datatree_factory.py @@ -5,7 +5,7 @@ import warnings from datetime import datetime, timezone from importlib.metadata import version as importlib_version -from typing import TYPE_CHECKING, Dict, Iterable, NamedTuple, Set +from typing import TYPE_CHECKING, Dict, Iterable, NamedTuple, Set, get_args import numpy as np import tensorstore as ts @@ -21,7 +21,7 @@ from xarray_kat.multiton import Multiton from xarray_kat.stores.vis_weight_flag_store_factory import VisWeightFlagFactory from xarray_kat.utils import corrprods_to_baseline_pols -from xarray_kat.xkat_types import VanVleckLiteralType +from xarray_kat.xkat_types import UvwSignConventionType, VanVleckLiteralType if TYPE_CHECKING: from katpoint import Target @@ -79,6 +79,7 @@ class DataTreeFactory: _scan_states: Set[str] _applycal: str | Iterable[str] _van_vleck: VanVleckLiteralType + _uvw_sign_convention: UvwSignConventionType _endpoint: str _token: str | None @@ -90,6 +91,7 @@ def __init__( data_products: Multiton[TelstateDataProducts], applycal: str | Iterable[str], scan_states: Iterable[str], + uvw_sign_convention: UvwSignConventionType, van_vleck: VanVleckLiteralType, endpoint: str, token: str | None = None, @@ -100,6 +102,7 @@ def __init__( self._data_products = data_products self._applycal = applycal self._scan_states = set(scan_states) + self._uvw_sign_convention = uvw_sign_convention self._van_vleck = van_vleck self._endpoint = endpoint self._token = token @@ -455,9 +458,16 @@ def WrappedArray(a): # Measurement Set `definition`_. # .. _CASA: https://casa.nrao.edu/Memos/CoordConvention.pdf # .. _definition: https://casa.nrao.edu/Memos/229.html#SECTION00064000000000000000 - uvw_coordinates = np.take(uvw_ant, ant1_index, axis=1) - np.take( - uvw_ant, ant2_index, axis=1 - ) + + if self._uvw_sign_convention == "fourier": + uvw = uvw_ant[:, ant2_index, :] - uvw_ant[:, ant1_index, :] + elif self._uvw_sign_convention == "casa": + uvw = uvw_ant[:, ant1_index, :] - uvw_ant[:, ant2_index, :] + else: + raise ValueError( + f"Invalid uvw sign convention {self._uvw_sign_convention} " + f"Should be one of {get_args(UvwSignConventionType)}" + ) flag_p_chunks = data_vars["FLAG"].encoding["preferred_chunks"] uvw_preferred_chunks = { @@ -468,7 +478,7 @@ def WrappedArray(a): data_vars["UVW"] = Variable( ("time", "baseline_id", "uvw_label"), - uvw_coordinates, + uvw, {"type": "uvw", "units": "m", "frame": "fk5"}, {"preferred_chunks": uvw_preferred_chunks}, ) diff --git a/src/xarray_kat/entrypoint.py b/src/xarray_kat/entrypoint.py index 1698cef..efb468e 100644 --- a/src/xarray_kat/entrypoint.py +++ b/src/xarray_kat/entrypoint.py @@ -18,7 +18,7 @@ from xarray_kat.datatree_factory import DataTreeFactory from xarray_kat.katdal_types import TelstateDataProducts, TelstateDataSource from xarray_kat.multiton import Multiton -from xarray_kat.xkat_types import VanVleckLiteralType +from xarray_kat.xkat_types import UvwSignConventionType, VanVleckLiteralType class KatStore(AbstractDataStore): @@ -33,6 +33,7 @@ class KatEntryPoint(BackendEntrypoint): "scan_states", "capture_block_id", "stream_name", + "uvw_sign_convention", "van_vleck", ] description = "Opens a MeerKAT data source" @@ -84,6 +85,7 @@ def open_datatree( scan_states: Iterable[str] = ("scan", "track"), capture_block_id: str | None = None, stream_name: str | None = None, + uvw_sign_convention: UvwSignConventionType = "casa", van_vleck: VanVleckLiteralType = "off", ): group_dicts = self.open_groups_as_dict( @@ -94,6 +96,7 @@ def open_datatree( scan_states=scan_states, capture_block_id=capture_block_id, stream_name=stream_name, + uvw_sign_convention=uvw_sign_convention, van_vleck=van_vleck, ) return DataTree.from_dict(group_dicts) @@ -108,6 +111,7 @@ def open_groups_as_dict( scan_states: Iterable[str] = ("scan", "track"), capture_block_id: str | None = None, stream_name: str | None = None, + uvw_sign_convention: UvwSignConventionType = "casa", van_vleck: VanVleckLiteralType = "off", ) -> Dict[str, Any]: url = str(filename_or_obj) @@ -143,6 +147,7 @@ def open_groups_as_dict( telstate_data_products, applycal, scan_states, + uvw_sign_convention, van_vleck, endpoint, token, diff --git a/src/xarray_kat/xkat_types.py b/src/xarray_kat/xkat_types.py index ce33110..bac2c8a 100644 --- a/src/xarray_kat/xkat_types.py +++ b/src/xarray_kat/xkat_types.py @@ -6,6 +6,7 @@ import numpy.typing as npt VanVleckLiteralType = Literal["off", "autocorr"] +UvwSignConventionType = Literal["fourier", "casa"] @dataclass(eq=True, unsafe_hash=True, slots=True, repr=True) diff --git a/tests/test_katdal.py b/tests/test_katdal.py index 45a6938..a27b595 100644 --- a/tests/test_katdal.py +++ b/tests/test_katdal.py @@ -1,5 +1,6 @@ import katdal import numpy as np +import pytest import xarray from pytest_httpserver import HTTPServer @@ -10,7 +11,10 @@ class TestKatdal: - def test_katdal_mock_server_basic(self, httpserver: HTTPServer, tmp_path): + @pytest.mark.parametrize("uvw_sign_convention", ["fourier", "casa"]) + def test_katdal_mock_server_basic( + self, httpserver: HTTPServer, uvw_sign_convention, tmp_path + ): """Tests that xarray-kat and katdal return the same data from the same datasource""" obs = SyntheticObservation("1234567890", ntime=8, nfreq=16, nants=4) obs.add_scan(range(0, 8), "track", "PKS1934") @@ -24,7 +28,9 @@ def test_katdal_mock_server_basic(self, httpserver: HTTPServer, tmp_path): rdb_url = f"{base_url}1234567890/1234567890_sdp_l0.full.rdb" ds = katdal.open(rdb_url) - dt = xarray.open_datatree(rdb_url, engine="xarray-kat") + dt = xarray.open_datatree( + rdb_url, engine="xarray-kat", uvw_sign_convention=uvw_sign_convention + ) def reorder_katdal_data(data): return ( @@ -48,6 +54,7 @@ def reorder_katdal_data(data): np.testing.assert_allclose(xarray_kat_flags, katdal_flags) xarray_kat_uvw = dt[children[0]].UVW.data - # Flip katdal sign convention to match CASA - katdal_uvw = np.stack([-ds.u, -ds.v, -ds.w], axis=2)[:, obs.corrprod_argsort] + katdal_uvw = np.stack([ds.u, ds.v, ds.w], axis=2)[:, obs.corrprod_argsort] + if uvw_sign_convention == "casa": + katdal_uvw = -katdal_uvw np.testing.assert_allclose(xarray_kat_uvw, katdal_uvw[:, :: obs.npol]) From a7791acbf0be116c71a4bae21afc3fbdfc03b216 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Tue, 17 Feb 2026 14:43:40 +0200 Subject: [PATCH 7/7] [skip ci] Update changelog --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bc52683..6fd9dc6 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -3,6 +3,7 @@ Changelog X.Y.Z (DD-MM-YYY) ----------------- +* Support choosing the UVW sign convention (:pr:`49`) * Add field and source dataset (:pr:`48`) * Add pre-computed UVW coordinates (:pr:`47`) * Test that xarray-kat and katdal produce the same data for the same data source (:pr:`45`)