NGFF v06 transformations

.gitignore

@@ -69,3 +69,14 @@ plans/
 .coverage
 htmlcov/
 
+# local hatch config
+hatch.toml
+
+# Kilo and plans
+.kilo/
+plans/
+
+# test coverage
+.coverage
+htmlcov/
+

pyproject.toml

@@ -39,7 +39,8 @@ dependencies = [
   "networkx",
   "numba>=0.55",
   "numpy",
-  "ome-zarr>=0.16",
+  "ome-zarr>=0.19.2",
+  "ome-zarr-models>=1.8",
   "pandas",
   "pooch",
   "pyarrow",

src/spatialdata/_core/transformation_manager/exceptions.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from spatialdata._types import ArrayLike
+from spatialdata.transformations.graph.vert import Axis, CoordSystem
 from spatialdata.transformations.ngff.ngff_coordinate_system import NgffCoordinateSystem
 
 
@@ -227,3 +229,65 @@ class TransformationManagerWarning(UserWarning):
     """Base warning category for TransformationManager."""
 
     pass
+
+
+class IncompatibleCoordSystemsError(Exception):
+    def __init__(self, input: CoordSystem, output: CoordSystem, message: str | None = None) -> None:
+        self.input = input
+        self.output = output
+        super().__init__(message or "Output axes can't be mapped to input axes")
+
+
+class MissingAxisError(Exception):
+    def __init__(self, axis: Axis, cs: CoordSystem) -> None:
+        self.axis = axis
+        self.cs = cs
+        super().__init__(f"Axis {axis.name} is not in coordinate system {cs.name}")
+
+
+class UnexpectedShapeError(Exception):
+    def __init__(
+        self,
+        *,
+        array_shape: tuple[int, ...],
+        expected_shape: tuple[int, ...] | str | None = None,
+        array_name: str | None = None,
+    ) -> None:
+        self.array_shape = array_shape
+        self.expected_shape = expected_shape
+        message = "Unexpected array shape"
+        if array_name is not None:
+            message += f"for '{array_name}'"
+        message += f": {array_shape}"
+        if expected_shape is not None:
+            message += f" instead of {expected_shape}"
+        super().__init__(message)
+
+
+class DeterminantDifferentFromOne(Exception):
+    def __init__(self, matrix: ArrayLike) -> None:
+        self.matrix = matrix
+        super().__init__("Matrix does not have det(M) == 1")
+
+
+class NotOrthonormalError(Exception):
+    def __init__(self, matrix: ArrayLike) -> None:
+        self.matrix = matrix
+        super().__init__("Matrix is not orthonormal")
+
+
+class EmptyTransformSequenceError(Exception):
+    def __init__(self) -> None:
+        super().__init__("Empty sequence of transformations")
+
+
+class AxisRedefinitionError(Exception):
+    def __init__(self, axis: Axis) -> None:
+        super().__init__(f"Axis {axis.name} is defined multiple times")
+
+
+class UnmappedAxisError(Exception):
+    def __init__(self, axis: Axis, cs: CoordSystem) -> None:
+        self.axis = axis
+        self.cs = cs
+        super().__init__(f"Axis {axis.name} from coordinate system {cs.name} is not mapped to anything")

src/spatialdata/_io/io_raster.py

@@ -7,6 +7,9 @@
 
 import dask.array as da
 import numpy as np
+import ome_zarr as oz
+import ome_zarr_models.v06.coordinate_transforms as ozm06trans
+import xarray as xr
 import zarr
 from ome_zarr.format import Format
 from ome_zarr.io import ZarrLocation
@@ -18,6 +21,7 @@
 from ome_zarr.writer import write_multiscale as write_multiscale_ngff
 from ome_zarr.writer import write_multiscale_labels as write_multiscale_labels_ngff
 from xarray import DataArray, DataTree
+from xarray.indexes import RangeIndex
 
 from spatialdata._io._utils import (
     _get_transformations_from_ngff_dict,
@@ -30,7 +34,7 @@
     RasterFormatType,
     get_ome_zarr_format,
 )
-from spatialdata._types import ELEMENT_TYPE, ELEMENT_TYPE_RASTER, GROUP_NAME
+from spatialdata._types import ELEMENT_TYPE, ELEMENT_TYPE_RASTER
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models.models import ATTRS_KEY
 from spatialdata.models.pyramids_utils import dask_arrays_to_datatree
@@ -40,6 +44,8 @@
     _set_transformations,
     compute_coordinates,
 )
+from spatialdata.transformations.graph.edge import BaseTransformationEdge, parse_ngff_transf
+from spatialdata.transformations.graph.vert import Axis, CoordSystem
 
 
 def _is_flat_int_sequence(value: object) -> TypeGuard[Sequence[int]]:
@@ -162,6 +168,97 @@ def _prepare_storage_options(
     return prepared_options
 
 
+def try_read_ngff06_multiscale(store: Path) -> tuple[DataTree, Sequence[BaseTransformationEdge]]:
+    multiscale = oz.OMEZarrMultiscale.from_ome_zarr(str(store))
+    assert isinstance(multiscale, oz.OMEZarrMultiscale)  # disambiguate from OMEZarrLabel
+    return try_parse_ngff06_multiscale(multiscale)
+
+
+def try_parse_ngff06_multiscale(multiscale: oz.OMEZarrMultiscale) -> tuple[DataTree, Sequence[BaseTransformationEdge]]:
+    """Parse an OMEZarMultiscale into a DataTree and collects Multiscale-level transforms."""
+    name_to_cs: dict[str, CoordSystem] = {}
+    for cs in multiscale.metadata.coordinateSystems or ():
+        parsed_cs = CoordSystem.try_from_model(cs)
+        name_to_cs[cs.name] = parsed_cs
+
+    parsed_transfs: list[BaseTransformationEdge] = []
+    for transf in multiscale.metadata.coordinateTransformations or ():
+        in_cs_id = transf.input
+        out_cs_ref = transf.output
+        # these should not be None as per the spec
+        assert in_cs_id is not None
+        assert out_cs_ref is not None
+
+        # FIXME: not handling references into labels yet, which use name and path
+        in_cs_name = in_cs_id.name
+        out_cs_name = out_cs_ref.name
+        assert in_cs_name is not None
+        assert out_cs_name is not None
+
+        # assume CS references are valid via ome-zarr(-models)-py
+        input = name_to_cs[in_cs_name]
+        output = name_to_cs[out_cs_name]
+        parsed = parse_ngff_transf(input=input, output=output, model=transf)
+        parsed_transfs.append(parsed)
+
+    omero = multiscale.omero
+    channel_names = None if omero is None else [d.color for d in omero.channels]
+
+    data_tree = xr.DataTree()
+    for scale_idx, (ds_md, ds) in enumerate(zip(multiscale.metadata.datasets, multiscale.images, strict=True)):
+        transf = ds_md.coordinateTransformations[0]
+
+        intrinsic_cs = name_to_cs[multiscale.metadata.intrinsic_coordinate_system.name]
+        assert transf.input is not None
+        assert transf.input.path is not None
+
+        # This coord system doesn't exist explicitly in the NGFF file, nor will
+        # it exist in our graph of transformations; It is only created here
+        # for the sake of creating the transformations that will be expressed
+        # in levels of a xr.DataTree
+        pixel_cs = CoordSystem(
+            name=transf.input.path,
+            axes=tuple(Axis(name=ax.name, type=ax.type) for ax in intrinsic_cs.axes),
+            virtual=True,
+        )
+
+        pixel_cs_to_intrinsic_ngff = ozm06trans.Sequence(transformations=ds_md.coordinateTransformations)
+        seq = parse_ngff_transf(input=pixel_cs, output=intrinsic_cs, model=pixel_cs_to_intrinsic_ngff)
+        ds_shape = np.asarray(ds.data.shape)
+        transformed_start = seq.transform_points(np.zeros_like(ds_shape)[np.newaxis, :])[0]
+        transformed_stop = seq.transform_points((ds_shape - 1)[np.newaxis, :])[0]
+
+        coords: xr.Coordinates = xr.Coordinates()
+        for low, high, ax, extent in zip(transformed_start, transformed_stop, intrinsic_cs.axes, ds_shape, strict=True):
+            if ax.type == "channel" and channel_names is not None:
+                coords = coords.merge({ax.name: channel_names}).coords
+            else:
+                axis_index = xr.Coordinates.from_xindex(
+                    RangeIndex.linspace(
+                        start=low,
+                        stop=high,
+                        num=extent,
+                        endpoint=True,
+                        dim=ax.name,
+                    )
+                )
+                coords = coords.merge(axis_index).coords
+
+        # Note: the magic "image" and "scale<N> " strings mimic the current
+        # behavior from `dask_arrays_to_datatree`
+        data_tree[f"scale{scale_idx}"] = xr.Dataset(
+            {
+                "image": xr.DataArray(
+                    ds.data,
+                    name="image",
+                    dims=intrinsic_cs.axes_names,
+                    coords=coords,
+                )
+            },
+        )
+    return data_tree, parsed_transfs
+
+
 def _read_multiscale(
     store: str | Path, raster_type: ELEMENT_TYPE_RASTER, reader_format: Format
 ) -> DataArray | DataTree:
@@ -267,9 +364,7 @@ def _get_multiscale_nodes(image_nodes: list[Node], nodes: list[Node]) -> list[No
     return nodes
 
 
-def _get_raster_element_group(
-    raster_type: ELEMENT_TYPE_RASTER, group: zarr.Group, element_name: str
-) -> zarr.Group:
+def _get_raster_element_group(raster_type: ELEMENT_TYPE_RASTER, group: zarr.Group, element_name: str) -> zarr.Group:
     """Get the Zarr group holding a raster element that has just been written.
 
     Labels are nested one level deeper than images: ome-zarr writes them inside a "labels" group, so for them the
@@ -288,7 +383,7 @@ def _get_raster_element_group(
     -------
     The Zarr group of the raster element.
     """
-    if raster_type != "labels":
+    if raster_type != ELEMENT_TYPE.LABELS:
         return group
     labels_group = group["labels"]
     if not isinstance(labels_group, zarr.Group):

src/spatialdata/_io/io_zarr.py

@@ -302,7 +302,7 @@ def _get_groups_for_element(
     The Zarr groups for the root, element_type and element for a specific element.
     """
     if not isinstance(zarr_path, Path):
-        raise ValueError("zarr_path should be a Path object")
+        raise TypeError("zarr_path should be a Path object")
 
     if element_type not in [
         "images",

src/spatialdata/_types.py

@@ -7,15 +7,15 @@
 from xarray import DataArray, DataTree
 
 __all__ = [
+    "ELEMENT_TYPE",
+    "ELEMENT_TYPE_RASTER",
+    "ELEMENT_TYPE_VECTOR",
+    "GROUP_NAME",
     "ArrayLike",
     "ColorLike",
     "DTypeLike",
     "JSONValue",
     "Raster_T",
-    "ELEMENT_TYPE",
-    "ELEMENT_TYPE_RASTER",
-    "ELEMENT_TYPE_VECTOR",
-    "GROUP_NAME",
 ]
 
 from numpy.typing import DTypeLike, NDArray
@@ -32,6 +32,7 @@
 type Raster_T = DataArray | DataTree
 ColorLike = tuple[float, ...] | str
 
+
 # A value that survives a round-trip through JSON, which is the invariant that `SpatialData.attrs` must satisfy: the
 # attrs are persisted with `zarr.Group.attrs.put()`, which rejects anything that is not JSON-serializable (e.g. numpy
 # arrays, sets, DataFrames). Note that JSON has no tuples and only string keys, so a tuple is read back as a list and

src/spatialdata/_utils.py

@@ -17,7 +17,12 @@
 from xarray import DataArray, Dataset, DataTree
 
 from spatialdata._types import ArrayLike, ListOrNDArrayFloating
-from spatialdata.transformations import Sequence, Translation, get_transformation, set_transformation
+from spatialdata.transformations import (
+    Sequence,
+    Translation,
+    get_transformation,
+    set_transformation,
+)
 
 RT = TypeVar("RT")
 

src/spatialdata/transformations/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from spatialdata.transformations import graph
 from spatialdata.transformations.operations import (
     align_elements_using_landmarks,
     get_transformation,
@@ -20,6 +21,7 @@
 )
 
 __all__ = [
+    "graph",
     "BaseTransformation",
     "Identity",
     "MapAxis",

src/spatialdata/transformations/graph/__init__.py

@@ -0,0 +1 @@
+