uda-xarray 0.2.2__tar.gz → 0.3.0__tar.gz

{uda_xarray-0.2.2/uda_xarray.egg-info → uda_xarray-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: uda-xarray
-Version: 0.2.2
+Version: 0.3.0
 Summary: UDA backend for xarray
 Author-email: Samuel Jackson <samuel.jackson@ukaea.uk>
 Maintainer-email: Samuel Jackson <samuel.jackson@ukaea.uk>
@@ -16,6 +16,7 @@ License-File: LICENSE
 Requires-Dist: uda>=2.9.2
 Requires-Dist: uda-mast
 Requires-Dist: xarray
+Requires-Dist: pydantic>=2.0
 Dynamic: license-file
 
 # uda-xarray

{uda_xarray-0.2.2 → uda_xarray-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uda-xarray"
-version = "0.2.2"
+version = "0.3.0"
 description = "UDA backend for xarray"
 readme = "README.md"
 requires-python = ">=3.9, <3.13"
@@ -23,11 +23,15 @@ dependencies = [
     "uda>=2.9.2",
     "uda-mast",
     "xarray",
+    "pydantic>=2.0",
 ]
 
 [tool.setuptools]
 packages = ["uda_xarray"]
 
+[tool.setuptools.package-data]
+uda_xarray = ["*.json"]
+
 [project.urls]
 repository = "https://github.com/samueljackson92/uda-xarray"
 
@@ -47,7 +51,7 @@ dev = [
 uda = "uda_xarray.main:UDABackendEntrypoint"
 
 [tool.bumpversion]
-current_version = "0.2.2"
+current_version = "0.3.0"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{uda_xarray-0.2.2 → uda_xarray-0.3.0}/tests/test_backend.py

@@ -1,11 +1,17 @@
-import pytest
-import xarray as xr
+"""Tests for the UDA xarray backend."""
+
+from unittest.mock import Mock
+
 import numpy as np
 import pyuda
-from unittest.mock import Mock
+import pytest
+import xarray as xr
+
+from uda_xarray.mappings import SignalMappings, SignalRange
 
 
 def test_open_uda_dataset(mocker):
+    """Test loading a 1D UDA signal as an xarray Dataset."""
     # Create mock signal object
     mock_signal = Mock()
     mock_signal.data = np.array([1.0, 2.0, 3.0])
@@ -45,6 +51,7 @@ def test_open_uda_dataset(mocker):
 
 
 def test_open_uda_dataset_2d(mocker):
+    """Test loading a 2D UDA signal as an xarray Dataset."""
     # Create mock 2D signal object
     mock_signal = Mock()
     mock_signal.data = np.array([[1.0, 2.0], [3.0, 4.0]])
@@ -81,6 +88,7 @@ def test_open_uda_dataset_2d(mocker):
 
 
 def test_open_uda_dataset_video(mocker):
+    """Test loading a UDA video signal as an xarray Dataset."""
     mock_signal = Mock()
     mock_signal.is_color = False
     mock_signal.frame_times = np.array([0.0, 0.033, 0.066])
@@ -115,6 +123,7 @@ def test_open_uda_dataset_video(mocker):
 
 
 def test_open_uda_dataset_invalid_signal(mocker):
+    """Test that an invalid UDA signal raises a RuntimeError."""
     # Mock the pyuda Client to raise an exception
     mock_client = Mock()
     mock_client.get.side_effect = pyuda.ServerException("Signal not found")
@@ -133,6 +142,7 @@ def test_open_uda_dataset_invalid_signal(mocker):
 
 
 def test_open_uda_dataset_invalid_format():
+    """Test that an invalid UDA URL format raises a ValueError."""
     try:
         xr.open_dataset("invalid_format", engine="uda")
     except ValueError as e:
@@ -152,6 +162,7 @@ def test_open_uda_dataset_invalid_format():
 
 
 def test_open_datatree(mocker):
+    """Test loading a 1D UDA signal as an xarray DataTree."""
     # Create mock signal object
     mock_signal = Mock()
     mock_signal.data = np.array([1.0, 2.0, 3.0])
@@ -195,6 +206,7 @@ def test_open_datatree(mocker):
 
 
 def test_open_datatree_2d(mocker):
+    """Test loading a 2D UDA signal as an xarray DataTree."""
     # Create mock 2D signal object
     mock_signal = Mock()
     mock_signal.data = np.array([[1.0, 2.0], [3.0, 4.0]])
@@ -236,6 +248,7 @@ def test_open_datatree_2d(mocker):
 
 
 def test_open_datatree_video(mocker):
+    """Test loading a UDA video signal as an xarray DataTree."""
     mock_signal = Mock()
     mock_signal.is_color = False
     mock_signal.frame_times = np.array([0.0, 0.033, 0.066])
@@ -276,7 +289,8 @@ def test_open_datatree_video(mocker):
 
 
 @pytest.mark.skip(reason="Requires access to real UDA server with video data")
-def test_open_uda_dataset_video_frame(mocker):
+def test_open_uda_dataset_video_frame():
+    """Test loading a single video frame from a real UDA server."""
     ds = xr.open_dataset(
         "uda://rba:30421", engine="uda", drop_variables=None, frame_number=1
     )
@@ -288,8 +302,208 @@ def test_open_uda_dataset_video_frame(mocker):
 
 @pytest.mark.skip(reason="Requires access to real UDA server with video data")
 def test_real_data_2d():
+    """Test loading real 2D data from a live UDA server."""
     ds = xr.open_dataset("uda://AYE_TE:30421", engine="uda")
     assert ds["data"].name == "data"
     assert ds["data"].dims == ("time", "radial_index")
     assert "time" in ds.coords
     assert "radial_index" in ds.coords
+
+
+def test_remapped_signal_uses_legacy_name(mocker):
+    """Requesting a canonical signal name for an old shot should transparently
+    remap it to the legacy stored name before calling the UDA client."""
+
+    # Build a minimal in-memory mapping:
+    #   canonical "AYC_TE" -> "AYC_TE" for shots 22832-30471
+    #                       -> "ATM_TE" for shots 1-22831
+    fake_mappings = SignalMappings(
+        mappings={
+            "AYC_TE": [
+                SignalRange(shot_min=22832, shot_max=30471, name="AYC_TE"),
+                SignalRange(shot_min=1, shot_max=22831, name="ATM_TE"),
+            ]
+        }
+    )
+    mocker.patch("uda_xarray.main._MAPPINGS", fake_mappings)
+
+    mock_signal = Mock()
+    mock_signal.data = np.array([1.0, 2.0, 3.0])
+    mock_signal.shape = (3,)
+    dim1 = Mock(label="time")
+    dim1.data = np.array([0.0, 1.0, 2.0])
+    mock_signal.dims = [dim1]
+    mock_signal.units = "eV"
+    mock_signal.time = Mock(label="time", data=np.array([0.0, 1.0, 2.0]))
+    mock_signal.errors = Mock(data=np.array([0.01, 0.01, 0.01]))
+
+    mock_client = Mock()
+    mock_client.get.return_value = mock_signal
+    mocker.patch("pyuda.Client", return_value=mock_client)
+    mocker.patch(
+        "uda_xarray.main.UDABackendEntrypoint._get_signal_type",
+        return_value="Signal",
+    )
+
+    # Shot 15000 is in the legacy range -> should call client with "ATM_TE"
+    ds = xr.open_dataset("uda://AYC_TE:15000", engine="uda")
+
+    mock_client.get.assert_called_once_with("ATM_TE", 15000)
+    assert ds["data"].attrs["uda_name"] == "ATM_TE"
+
+
+def test_remapped_signal_uses_modern_name(mocker):
+    """Requesting a canonical signal name for a recent shot should keep the
+    modern stored name."""
+
+    fake_mappings = SignalMappings(
+        mappings={
+            "AYC_TE": [
+                SignalRange(shot_min=22832, shot_max=30471, name="AYC_TE"),
+                SignalRange(shot_min=1, shot_max=22831, name="ATM_TE"),
+            ]
+        }
+    )
+    mocker.patch("uda_xarray.main._MAPPINGS", fake_mappings)
+
+    mock_signal = Mock()
+    mock_signal.data = np.array([1.0, 2.0, 3.0])
+    mock_signal.shape = (3,)
+    dim1 = Mock(label="time")
+    dim1.data = np.array([0.0, 1.0, 2.0])
+    mock_signal.dims = [dim1]
+    mock_signal.units = "eV"
+    mock_signal.time = Mock(label="time", data=np.array([0.0, 1.0, 2.0]))
+    mock_signal.errors = Mock(data=np.array([0.01, 0.01, 0.01]))
+
+    mock_client = Mock()
+    mock_client.get.return_value = mock_signal
+    mocker.patch("pyuda.Client", return_value=mock_client)
+    mocker.patch(
+        "uda_xarray.main.UDABackendEntrypoint._get_signal_type",
+        return_value="Signal",
+    )
+
+    # Shot 30000 is in the modern range -> should call client with "AYC_TE"
+    ds = xr.open_dataset("uda://AYC_TE:30000", engine="uda")
+
+    mock_client.get.assert_called_once_with("AYC_TE", 30000)
+    assert ds["data"].attrs["uda_name"] == "AYC_TE"
+
+
+def test_mastu_hcam_remapped_to_sanx_name(mocker):
+    """For a MASTU shot in the SAnx05 range, the canonical HCAM channel name
+    should be remapped to the SAnx05 digitiser path."""
+
+    # /XSX/HCAM/L/CH01/DATA is the canonical name; for shots 46353-49476 the
+    # data was stored as /xsx/SAnx05-01/ch14.
+    fake_mappings = SignalMappings(
+        mappings={
+            "/XSX/HCAM/L/CH01/DATA": [
+                SignalRange(
+                    shot_min=49904, shot_max=51056, name="/XSX/HCAM/L/CH01/DATA"
+                ),
+                SignalRange(shot_min=49476, shot_max=49904, name="/XSX/HCAM/L/CH01"),
+                SignalRange(shot_min=46353, shot_max=49476, name="/xsx/SAnx05-01/ch14"),
+                SignalRange(shot_min=44395, shot_max=46353, name="/XSX/HCAM/L/CH01"),
+            ]
+        }
+    )
+    mocker.patch("uda_xarray.main._MAPPINGS", fake_mappings)
+
+    mock_signal = Mock()
+    mock_signal.data = np.array([1.0, 2.0, 3.0])
+    mock_signal.shape = (3,)
+    dim1 = Mock(label="time")
+    dim1.data = np.array([0.0, 1.0, 2.0])
+    mock_signal.dims = [dim1]
+    mock_signal.units = "V"
+    mock_signal.time = Mock(label="time", data=np.array([0.0, 1.0, 2.0]))
+    mock_signal.errors = Mock(data=np.array([0.01, 0.01, 0.01]))
+
+    mock_client = Mock()
+    mock_client.get.return_value = mock_signal
+    mocker.patch("pyuda.Client", return_value=mock_client)
+    mocker.patch(
+        "uda_xarray.main.UDABackendEntrypoint._get_signal_type",
+        return_value="Signal",
+    )
+
+    # Shot 47000 falls in the SAnx05 range -> should call client with ch14 path
+    ds = xr.open_dataset("uda:///XSX/HCAM/L/CH01/DATA:47000", engine="uda")
+
+    mock_client.get.assert_called_once_with("/xsx/SAnx05-01/ch14", 47000)
+    assert ds["data"].attrs["uda_name"] == "/xsx/SAnx05-01/ch14"
+
+
+def test_mastu_hcam_modern_name_passthrough(mocker):
+    """For the most-recent MASTU shot range, the canonical HCAM channel name
+    should be used unchanged (it already matches the stored path)."""
+
+    fake_mappings = SignalMappings(
+        mappings={
+            "/XSX/HCAM/L/CH01/DATA": [
+                SignalRange(
+                    shot_min=49904, shot_max=51056, name="/XSX/HCAM/L/CH01/DATA"
+                ),
+                SignalRange(shot_min=49476, shot_max=49904, name="/XSX/HCAM/L/CH01"),
+                SignalRange(shot_min=46353, shot_max=49476, name="/xsx/SAnx05-01/ch14"),
+                SignalRange(shot_min=44395, shot_max=46353, name="/XSX/HCAM/L/CH01"),
+            ]
+        }
+    )
+    mocker.patch("uda_xarray.main._MAPPINGS", fake_mappings)
+
+    mock_signal = Mock()
+    mock_signal.data = np.array([1.0, 2.0, 3.0])
+    mock_signal.shape = (3,)
+    dim1 = Mock(label="time")
+    dim1.data = np.array([0.0, 1.0, 2.0])
+    mock_signal.dims = [dim1]
+    mock_signal.units = "V"
+    mock_signal.time = Mock(label="time", data=np.array([0.0, 1.0, 2.0]))
+    mock_signal.errors = Mock(data=np.array([0.01, 0.01, 0.01]))
+
+    mock_client = Mock()
+    mock_client.get.return_value = mock_signal
+    mocker.patch("pyuda.Client", return_value=mock_client)
+    mocker.patch(
+        "uda_xarray.main.UDABackendEntrypoint._get_signal_type",
+        return_value="Signal",
+    )
+
+    # Shot 50000 is in the most-recent range -> canonical name used as-is
+    ds = xr.open_dataset("uda:///XSX/HCAM/L/CH01/DATA:50000", engine="uda")
+
+    mock_client.get.assert_called_once_with("/XSX/HCAM/L/CH01/DATA", 50000)
+    assert ds["data"].attrs["uda_name"] == "/XSX/HCAM/L/CH01/DATA"
+
+
+def test_unmapped_signal_passes_through(mocker):
+    """A signal not present in the mappings should be forwarded as-is."""
+
+    fake_mappings = SignalMappings(mappings={})
+    mocker.patch("uda_xarray.main._MAPPINGS", fake_mappings)
+
+    mock_signal = Mock()
+    mock_signal.data = np.array([1.0, 2.0])
+    mock_signal.shape = (2,)
+    dim1 = Mock(label="time")
+    dim1.data = np.array([0.0, 1.0])
+    mock_signal.dims = [dim1]
+    mock_signal.units = "A"
+    mock_signal.time = Mock(label="time", data=np.array([0.0, 1.0]))
+    mock_signal.errors = Mock(data=np.array([0.1, 0.1]))
+
+    mock_client = Mock()
+    mock_client.get.return_value = mock_signal
+    mocker.patch("pyuda.Client", return_value=mock_client)
+    mocker.patch(
+        "uda_xarray.main.UDABackendEntrypoint._get_signal_type",
+        return_value="Signal",
+    )
+
+    ds = xr.open_dataset("uda://SOME_UNKNOWN_SIGNAL:30421", engine="uda")
+
+    mock_client.get.assert_called_once_with("SOME_UNKNOWN_SIGNAL", 30421)
+    assert ds["data"].attrs["uda_name"] == "SOME_UNKNOWN_SIGNAL"

uda_xarray-0.3.0/uda_xarray/__init__.py

@@ -0,0 +1 @@
+"""UDA backend for xarray."""

{uda_xarray-0.2.2 → uda_xarray-0.3.0}/uda_xarray/main.py

@@ -10,6 +10,10 @@ import xarray as xr
 from mast.mast_client import ListType
 from xarray.backends import BackendEntrypoint
 
+from uda_xarray.mappings import SignalMappings
+
+_MAPPINGS = SignalMappings.from_file()
+
 
 # Set up UDA environment variables with defaults if not already set. This is required for
 # the pyuda client to work correctly outside of Freia.
@@ -40,6 +44,7 @@ class UDABackendEntrypoint(BackendEntrypoint):
         *,
         drop_variables=None,
         frame_number: Optional[int] = None,  # noqa: F821
+        **kwargs,
     ) -> xr.Dataset:
         """Open a UDA dataset given a signal name and shot number.
 
@@ -65,6 +70,8 @@ class UDABackendEntrypoint(BackendEntrypoint):
         name = name.replace("uda://", "")
         shot = int(shot)
 
+        name = _MAPPINGS.resolve(name, shot) or name
+
         client = pyuda.Client()
         signal_type = self._get_signal_type(client, name, shot)
 
@@ -161,8 +168,10 @@ class UDABackendEntrypoint(BackendEntrypoint):
             return SignalType.IMAGE
         return SignalType.SIGNAL
 
-    def open_datatree(self, filename_or_obj, *, drop_variables=None):
-        dataset = self.open_dataset(filename_or_obj, drop_variables=drop_variables)
+    def open_datatree(self, filename_or_obj, *, drop_variables=None, **kwargs):
+        dataset = self.open_dataset(
+            filename_or_obj, drop_variables=drop_variables, **kwargs
+        )
         return xr.DataTree(dataset)
 
     def open_groups_as_dict(self, filename_or_obj, **kwargs):

uda_xarray-0.3.0/uda_xarray/mappings.py

@@ -0,0 +1,78 @@
+"""Signal name mappings for the UDA xarray backend."""
+
+from __future__ import annotations
+
+import importlib.resources
+import json
+from typing import Iterator
+
+from pydantic import BaseModel
+
+
+class SignalRange(BaseModel):
+    """A single shot-range entry mapping a canonical signal name to the name
+    used in the data store for that range of shots."""
+
+    shot_min: int
+    shot_max: int
+    name: str
+
+    def contains(self, shot: int) -> bool:
+        """Return True if *shot* falls within this range (inclusive)."""
+        return self.shot_min <= shot <= self.shot_max
+
+
+class SignalMappings(BaseModel):
+    """Complete signal-alias mapping table loaded from mast.json.
+
+    The mapping is a dict whose keys are the canonical (most recent) signal
+    names and whose values are lists of :class:`SignalRange` entries ordered
+    from newest to oldest.
+    """
+
+    mappings: dict[str, list[SignalRange]]
+
+    @classmethod
+    def from_file(cls, path: str | None = None) -> "SignalMappings":
+        """Load mappings from *path*.
+
+        When *path* is ``None`` all bundled ``*.json`` mapping files are
+        loaded and merged into a single :class:`SignalMappings` instance.
+        """
+        if path is None:
+            merged: dict[str, list] = {}
+            pkg = importlib.resources.files("uda_xarray")
+            for resource in pkg.iterdir():
+                if resource.name.endswith(".json"):
+                    with importlib.resources.as_file(resource) as p:
+                        for k, v in json.loads(p.read_text(encoding="utf-8")).items():
+                            if k in merged:
+                                merged[k].extend(v)
+                            else:
+                                merged[k] = list(v)
+            raw = merged
+        else:
+            with open(path, encoding="utf-8") as f:
+                raw = json.load(f)
+
+        return cls(mappings={k: [SignalRange(**r) for r in v] for k, v in raw.items()})
+
+    def resolve(self, signal: str, shot: int) -> str | None:
+        """Return the stored name for *signal* at the given *shot* number.
+
+        Returns *None* if *signal* is not in the mapping or no range covers
+        the shot.
+        """
+        ranges = self.mappings.get(signal)
+        if ranges is None:
+            return None
+        for r in ranges:
+            if r.contains(shot):
+                return r.name
+        return None
+
+    def __iter__(self) -> Iterator[tuple[str, list[SignalRange]]]:
+        return iter(self.mappings.items())
+
+    def __len__(self) -> int:
+        return len(self.mappings)