Derive the NGFF axis type from the axis name (0.5.0)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "polystore"
-version = "0.4.0"
+version = "0.5.0"
 description = "Framework-agnostic multi-backend storage abstraction for ML and scientific computing"
 readme = "README.md"
 requires-python = ">=3.11"

src/polystore/zarr_batch.py

@@ -6,6 +6,7 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from enum import Enum
+from types import MappingProxyType
 
 _ATTR_PREFIX = "polystore"
 ATTR_FILENAME_MAP = f"{_ATTR_PREFIX}_filename_map"
@@ -14,6 +15,15 @@
 ATTR_IMAGE_COORDINATE = f"{_ATTR_PREFIX}_image_coordinate"
 
 
+NGFF_AXIS_TYPES: Mapping[str, str] = MappingProxyType(
+    {"t": "time", "c": "channel", "z": "space", "y": "space", "x": "space"}
+)
+"""OME-NGFF axis names and the axis type each one declares."""
+
+PIXEL_AXIS_NAMES = ("y", "x")
+"""Storage-owned pixel axes that follow every declared batch axis."""
+
+
 class ZarrBatchAxisRole(Enum):
     """Storage role owned by one declared batch axis."""
 
@@ -31,17 +41,19 @@ class ZarrBatchAxis:
     """
 
     name: str
-    axis_type: str
     values: tuple[str, ...]
     role: ZarrBatchAxisRole = ZarrBatchAxisRole.ARRAY
 
     def __post_init__(self) -> None:
         if not self.name:
             raise ValueError("Zarr batch axis name cannot be empty")
-        if not self.axis_type:
-            raise ValueError("Zarr batch axis type cannot be empty")
-        if self.name in {"x", "y"}:
+        if self.name in PIXEL_AXIS_NAMES:
             raise ValueError("Zarr batch axes cannot redeclare pixel axes x or y")
+        if self.role is ZarrBatchAxisRole.ARRAY and self.name not in NGFF_AXIS_TYPES:
+            raise ValueError(
+                f"Zarr batch array axis {self.name!r} is not an OME-NGFF axis; "
+                f"array axes are named {sorted(set(NGFF_AXIS_TYPES) - set(PIXEL_AXIS_NAMES))}"
+            )
         if not self.values:
             raise ValueError(f"Zarr batch axis {self.name!r} has no values")
         if len(set(self.values)) != len(self.values):
@@ -58,7 +70,7 @@ def size(self) -> int:
     def ngff_declaration(self) -> dict[str, str]:
         """Return this axis in the representation consumed by ome-zarr."""
 
-        return {"name": self.name, "type": self.axis_type}
+        return {"name": self.name, "type": NGFF_AXIS_TYPES[self.name]}
 
 
 @dataclass(frozen=True, slots=True)
@@ -149,9 +161,8 @@ def array_shape(self) -> tuple[int, ...]:
     def ngff_axes(self) -> tuple[dict[str, str], ...]:
         """Return declared axes followed by the storage-owned pixel axes."""
 
-        return tuple(axis.ngff_declaration() for axis in self.array_axes) + (
-            {"name": "y", "type": "space"},
-            {"name": "x", "type": "space"},
+        return tuple(axis.ngff_declaration() for axis in self.array_axes) + tuple(
+            {"name": name, "type": NGFF_AXIS_TYPES[name]} for name in PIXEL_AXIS_NAMES
         )
 
     def dimensions_attribute(self) -> dict[str, Mapping[str, object]]:

tests/test_zarr_backend.py

@@ -152,7 +152,6 @@ def save_hcs_image_axis_batch(
         axes=(
             ZarrBatchAxis(
                 axis_name,
-                "field",
                 values,
                 ZarrBatchAxisRole.HCS_IMAGE,
             ),
@@ -306,15 +305,14 @@ def test_batch_save_and_load(self, temp_zarr_dir, strategy, expected_chunks):
         coordinates = ((1, 0, 1, 0), (0, 0, 0, 0), (0, 0, 1, 0), (1, 0, 0, 0))
         layout = ZarrBatchLayout(
             axes=(
-                ZarrBatchAxis("t", "time", ("1", "2")),
+                ZarrBatchAxis("t", ("1", "2")),
                 ZarrBatchAxis(
-                    "field",
                     "field",
                     ("1",),
                     ZarrBatchAxisRole.HCS_IMAGE,
                 ),
-                ZarrBatchAxis("c", "channel", ("1", "2")),
-                ZarrBatchAxis("z", "space", ("1",)),
+                ZarrBatchAxis("c", ("1", "2")),
+                ZarrBatchAxis("z", ("1",)),
             ),
             item_coordinates=coordinates,
         )
@@ -356,7 +354,7 @@ def test_batch_save_and_load(self, temp_zarr_dir, strategy, expected_chunks):
     def test_batch_operations_length_mismatch(self, zarr_backend, temp_zarr_dir):
         """Test that batch operations raise error on length mismatch."""
         layout = ZarrBatchLayout(
-            axes=(ZarrBatchAxis("c", "channel", ("1",)),),
+            axes=(ZarrBatchAxis("c", ("1",)),),
             item_coordinates=((0,),),
         )
         with pytest.raises(ValueError, match="equal lengths"):
@@ -384,7 +382,7 @@ def test_invalid_batch_preserves_existing_well(self, zarr_backend, tmp_path, bad
         )
         before = zarr_backend.load_batch(paths)
         layout = ZarrBatchLayout(
-            axes=(ZarrBatchAxis("c", "channel", ("1", "2")),),
+            axes=(ZarrBatchAxis("c", ("1", "2")),),
             item_coordinates=((0,), (1,)),
         )
         with pytest.raises(ValueError, match="image-plane shape|two-dimensional"):
@@ -411,15 +409,14 @@ def test_image_axis_creates_one_hcs_image_per_value(
         ]
         layout = ZarrBatchLayout(
             axes=(
-                ZarrBatchAxis("t", "time", ("1",)),
+                ZarrBatchAxis("t", ("1",)),
                 ZarrBatchAxis(
-                    "field",
                     "field",
                     ("3", "7"),
                     ZarrBatchAxisRole.HCS_IMAGE,
                 ),
-                ZarrBatchAxis("c", "channel", ("1",)),
-                ZarrBatchAxis("z", "space", ("1",)),
+                ZarrBatchAxis("c", ("1",)),
+                ZarrBatchAxis("z", ("1",)),
             ),
             item_coordinates=((0, 0, 0, 0), (0, 1, 0, 0)),
         )

tests/test_zarr_batch.py

@@ -15,18 +15,38 @@
 def test_layout_rejects_incomplete_dense_coordinates() -> None:
     with pytest.raises(ValueError, match="one item for every coordinate"):
         ZarrBatchLayout(
-            axes=(ZarrBatchAxis("c", "channel", ("1", "2")),),
+            axes=(ZarrBatchAxis("c", ("1", "2")),),
             item_coordinates=((0,),),
         )
 
 
+def test_array_axes_take_their_type_from_the_ngff_axis_name() -> None:
+    layout = ZarrBatchLayout(
+        axes=(
+            ZarrBatchAxis("t", ("1",)),
+            ZarrBatchAxis("c", ("1",)),
+            ZarrBatchAxis("z", ("1",)),
+        ),
+        item_coordinates=((0, 0, 0),),
+    )
+    assert layout.ngff_axes == (
+        {"name": "t", "type": "time"},
+        {"name": "c", "type": "channel"},
+        {"name": "z", "type": "space"},
+        {"name": "y", "type": "space"},
+        {"name": "x", "type": "space"},
+    )
+    with pytest.raises(ValueError, match="not an OME-NGFF axis"):
+        ZarrBatchAxis("band", ("1",))
+
+
 def test_layout_rejects_multiple_hcs_image_axes() -> None:
     image_role = ZarrBatchAxisRole.HCS_IMAGE
     with pytest.raises(ValueError, match="at most one HCS image axis"):
         ZarrBatchLayout(
             axes=(
-                ZarrBatchAxis("field", "field", ("1",), image_role),
-                ZarrBatchAxis("scene", "field", ("1",), image_role),
+                ZarrBatchAxis("field", ("1",), image_role),
+                ZarrBatchAxis("scene", ("1",), image_role),
             ),
             item_coordinates=((0, 0),),
         )