NGFF v06 transformations

pyproject.toml

@@ -35,7 +35,8 @@ dependencies = [
     "networkx",
     "numba>=0.55.0",
     "numpy",
-    "ome_zarr>=0.16.0",
+    "ome_zarr>=0.19.2",
+    "ome-zarr-models>=1.8.0",
     "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

@@ -6,6 +6,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
@@ -17,6 +20,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,
@@ -38,6 +42,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]]:
@@ -160,6 +166,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:

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 @@
+