From f57dc1dc360e75df263d4125708b060586b28b9c Mon Sep 17 00:00:00 2001 From: Tomaz Vieira Date: Fri, 21 Aug 2026 11:25:40 +0000 Subject: [PATCH 1/3] Creates the graph module based on the Ngff* classes Cleans up and modernizes the NGff classes so that they have stronger invariants and guarantee that they always have input and output. Reading and writing to zarr is done via ome-zarr-models-py. Uses graph module as part of io_raster.py::try_read_ngff06_multiscale to interpret Ngff trnasformations and produce an output that could be added to the new graph implementation. --- pyproject.toml | 3 +- src/spatialdata/_io/io_raster.py | 83 ++ src/spatialdata/transformations/__init__.py | 2 + .../transformations/graph/__init__.py | 1 + src/spatialdata/transformations/graph/edge.py | 948 ++++++++++++++++++ src/spatialdata/transformations/graph/vert.py | 201 ++++ 6 files changed, 1237 insertions(+), 1 deletion(-) create mode 100644 src/spatialdata/transformations/graph/__init__.py create mode 100644 src/spatialdata/transformations/graph/edge.py create mode 100644 src/spatialdata/transformations/graph/vert.py diff --git a/pyproject.toml b/pyproject.toml index f25a5fbb3..1f267868c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,8 @@ dependencies = [ "networkx", "numba>=0.55.0", "numpy", - "ome_zarr>=0.16.0", + # "ome_zarr>=0.16.0", + "ome_zarr @ git+https://github.com/ome/ome-zarr-py/@e859b55425b740335309876bb023fe52977a3045", "pandas", "pooch", "pyarrow", diff --git a/src/spatialdata/_io/io_raster.py b/src/spatialdata/_io/io_raster.py index 29973ad20..02570ef80 100644 --- a/src/spatialdata/_io/io_raster.py +++ b/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 BaseTransfEdge, parse_ngff_transf +from spatialdata.transformations.graph.vert import Axis, CoordSystem def _is_flat_int_sequence(value: object) -> TypeGuard[Sequence[int]]: @@ -160,6 +166,83 @@ def _prepare_storage_options( return prepared_options +def try_read_ngff06_multiscale(store: Path) -> tuple[DataTree, Sequence[BaseTransfEdge]]: + multiscale = oz.OMEZarrMultiscale.from_ome_zarr(str(store)) + assert isinstance(multiscale, oz.OMEZarrMultiscale) # disambiguate from OMEZarrLabel + + 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[BaseTransfEdge] = [] + 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 + + in_cs_name = in_cs_id.name + out_cs_name = out_cs_ref.name + # FIXME: not handling references into labels yet + 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] + + out_cs = name_to_cs[multiscale.metadata.intrinsic_coordinate_system.name] + assert transf.input is not None + assert transf.input.path is not None + in_cs = CoordSystem(name=str(transf.input.name), axes=[Axis(name=ax.name, type=ax.type) for ax in out_cs.axes]) + + ozm_seq = ozm06trans.Sequence(transformations=ds_md.coordinateTransformations) + seq = parse_ngff_transf(input=in_cs, output=out_cs, model=ozm_seq) + ds_shape = np.asarray(ds.data.shape) + transformed_start = seq.transform_points(np.zeros_like(ds.data.shape)[np.newaxis, :])[0] + transformed_stop = seq.transform_points((ds_shape - 1)[np.newaxis, :])[0] + + coords = xr.Coordinates() + for low, high, ax, extent in zip(transformed_start, transformed_stop, out_cs.axes, ds_shape, strict=True): + if ax.type == "channel" and channel_names is not None: + coords.merge({ax.name: channel_names}) + continue + coords = coords.merge( + xr.Coordinates.from_xindex( + RangeIndex.linspace( + start=low, + stop=high, + num=extent, + endpoint=True, + dim=ax.name, + ) + ) + ) + + data_tree[f"scale{scale_idx}"] = xr.Dataset( + { + "image": xr.DataArray( + ds.data, + name="image", + dims=out_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: diff --git a/src/spatialdata/transformations/__init__.py b/src/spatialdata/transformations/__init__.py index e95a92766..802cc47cd 100644 --- a/src/spatialdata/transformations/__init__.py +++ b/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", diff --git a/src/spatialdata/transformations/graph/__init__.py b/src/spatialdata/transformations/graph/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/src/spatialdata/transformations/graph/__init__.py @@ -0,0 +1 @@ + diff --git a/src/spatialdata/transformations/graph/edge.py b/src/spatialdata/transformations/graph/edge.py new file mode 100644 index 000000000..b01a11e97 --- /dev/null +++ b/src/spatialdata/transformations/graph/edge.py @@ -0,0 +1,948 @@ +# pyright: strict + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Sequence +from typing import Final + +import numpy as np +import ome_zarr_models.v06.coordinate_transforms as ozm06trans +import pydantic as pyd +import xarray as xr + +from spatialdata._types import ArrayLike +from spatialdata.transformations.graph.vert import Axis, CoordSystem + + +class GarbledInput(Exception): + def __init__(self, message: str, input: pyd.JsonValue) -> None: + import json + + super().__init__(message + "\n" + json.dumps(input, indent=4)) + self.input = input + + +class BaseTransfEdge(ABC): + """Base class for all the transformations defined by the NGFF specification.""" + + input: Final[CoordSystem] + output: Final[CoordSystem] + name: str | None + + def __init__( + self, + name: str | None, + *, + input: CoordSystem, + output: CoordSystem, + ) -> None: + self.input = input + self.output = output + self.name = name + super().__init__() + + def __repr__(self) -> str: + domain = ", ".join(self.input.axes_names) + codomain = ", ".join(self.output.axes_names) + return f"{type(self).__name__} ({domain} -> {codomain})" + + @abstractmethod + def inverse(self) -> BaseTransfEdge: + """Return the inverse of the transformation.""" + + @abstractmethod + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Notes + ------- + This function will check if the dimensionality of the input and output coordinate systems of the + transformation are compatible with the given points. + """ + + @abstractmethod + def to_affine(self) -> AffineEdge: + """Convert the transformation to an affine transformation, whenever the conversion can be made.""" + + def _validate_transform_points_shapes(self, points: xr.DataArray | xr.DataTree | ArrayLike) -> None: + """ + Validate if the shape of the points (coordinats to be transformed) are consistent with the input size of the + transformation. + """ + input_size = len(self.input.axes) + if len(points.shape) != 2 or points.shape[1] != input_size: + raise ValueError( + f"points must be a tensor of shape (n, d), where n is the number of points and d is the " + f"the number of spatial dimensions. Points shape: {points.shape}, input size: {input_size}" + ) + + # order of the composition: self is applied first, then the transformation passed as argument + def compose_with(self, transformation: BaseTransfEdge) -> BaseTransfEdge: + """ + Compose the transfomation object with another transformation + + Parameters + ---------- + transformation + The transformation to compose with. + + Returns + ------- + The compoesed transformation. + + Notes + ------- + Self is applied first, then the transformation passed as argument. + """ + return SequenceEdge([self, transformation], name=None) # FIXME: no name? + + @abstractmethod + def to_model(self) -> ozm06trans.AnyTransform: + pass + + +class AffineEdge(BaseTransfEdge): + """The Affine transformation from the NGFF specification.""" + + linear: Final[ArrayLike] + translation: Final[ArrayLike] + affine: Final[ArrayLike] + + def __init__( + self, + name: str | None, + *, + linear: ArrayLike, + translation: ArrayLike | None = None, + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Parameters + ---------- + name + A human readable name for this transformation + linear + The linear part of this transformation, i.e., the one that keeps + the origin in the same place. Shape must be (output.num_axes, input.num_axes) + translation y + The translation part of this transformation, of shape (output.num_axes,) + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + num_inputs = input.num_axes + num_outputs = output.num_axes + translation = np.zeros(num_outputs) if translation is None else translation + + expected_linear_shape = (num_outputs, num_inputs) + if linear.shape != expected_linear_shape: + raise ValueError(f"linear's shape is {linear.shape}. Expected f{(num_outputs, num_inputs)}") + expected_translation_shape = (num_outputs,) + if translation.shape != expected_translation_shape: + raise ValueError(f"translation's shape is {translation.shape}. Expected {expected_translation_shape}") + + self.linear = linear + self.translation = translation + + self.affine = np.zeros((num_outputs + 1, num_inputs + 1)) + self.affine[:-1, :-1] = self.linear + self.affine[:-1, -1] = self.translation + self.affine[-1, -1] = 1 + + super().__init__(input=input, output=output, name=name) + + def __repr__(self) -> str: + s = super().__repr__() + "\n" + s += "\n".join(str(row) for row in self.affine) + return s + + @classmethod + def from_affine_matrix( + cls, + *, + name: str | None, + affine_matrix: ArrayLike, + input: CoordSystem, + output: CoordSystem, + ) -> AffineEdge: + """Creates an AffineEdge from a raw affine matrix + + Parameters + ---------- + name + A human readable name for this transformation + affine_matrix + row-major, (output.num_axes + 1, input.num_axes + 1) matrix with: + - linear part at the top left + - translation as the rightmost column + - last row is [0, 0, ..., 0, 1] + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + return AffineEdge( + linear=affine_matrix[:-1, :-1], + translation=affine_matrix[-1, :-1], + input=input, + output=output, + name=name, + ) + + def inverse(self) -> BaseTransfEdge: + inv = np.linalg.inv(self.affine) + return AffineEdge( + linear=inv[:-1, :-1], + translation=inv[-1, :-1], + input=self.output, + output=self.input, + name=self.name and f"{self.name}__affine", + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + self._validate_transform_points_shapes(points) + p = np.vstack([points.T, np.ones(points.shape[0])]) + q = self.affine @ p + res = q[: self.output.num_axes, :].T + assert isinstance(res, np.ndarray) + return res + + def to_affine(self) -> AffineEdge: + return self + + def to_model(self) -> ozm06trans.Affine: + return ozm06trans.Affine( + name=self.name, + affine=tuple(tuple(row) for row in self.affine[:-1, :]), + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + ) + + +class IdentityEdge(BaseTransfEdge): + """The Identity transformation from the NGFF specification.""" + + def __init__( + self, + name: str | None, + *, + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Parameters + ---------- + name + A human readable name for this transformation + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + if input.num_axes != output.num_axes: + raise ValueError("Input and output must have the same number of dimensions") + super().__init__(input=input, output=output, name=name) + + def inverse(self) -> BaseTransfEdge: + return IdentityEdge(input=self.output, output=self.input, name=self.name and f"{self.name}__inverse") + + def transform_points(self, points: ArrayLike) -> ArrayLike: + self._validate_transform_points_shapes(points) + return points + + def to_affine(self) -> AffineEdge: + return AffineEdge( + linear=np.eye(self.input.num_axes), + input=self.input, + output=self.output, + name=self.name and f"{self.name}__affine", + ) + + def to_model(self) -> ozm06trans.Identity: + return ozm06trans.Identity( + name=self.name, + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + ) + + +class MapAxisEdge(BaseTransfEdge): + """The MapAxis transformation from the NGFF specification.""" + + def __init__( + self, + name: str | None, + *, + output_to_input: dict[str, str], + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Init the NgffMapAxis object. + Parameters + ---------- + name + A human readable name for this transformation + output_to_input + A dictionary mapping the output axes (keys) to the input axes (values). + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + for out_ax, inp_ax in output_to_input.items(): + if not input.has_axis(inp_ax): + raise ValueError(f"input has no axis named {inp_ax}") + if not output.has_axis(out_ax): + raise ValueError(f"output has no axis named {out_ax}") + if not (len(output_to_input) == output.num_axes == input.num_axes): + raise ValueError("input_to_output, input and output must have the same number of axes entries") + if len(set(output_to_input.values())) != len(output_to_input): + raise ValueError("input_to_output must map unique inputs to unique outputs") + + self.output_to_input = output_to_input + super().__init__(input=input, output=output, name=name) + + def __repr__(self) -> str: + s = super().__repr__() + "\n" + s += "\n".join(f" {out} <- {inp}\n" for out, inp in self.output_to_input.items()) + return s + + def inverse(self) -> BaseTransfEdge: + return MapAxisEdge( + output_to_input={v: k for k, v in self.output_to_input.items()}, + input=self.output, + output=self.input, + name=self.name and f"{self.name}__inverse", + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + input_axes = self.input.axes_names + output_axes = self.output.axes_names + self._validate_transform_points_shapes(points) + new_indices = [input_axes.index(self.output_to_input[ax]) for ax in output_axes] + mapped = points[:, new_indices] + assert isinstance(mapped, np.ndarray) + return mapped + + def to_affine(self) -> AffineEdge: + input_axes = self.input.axes_names + output_axes = self.output.axes_names + linear: ArrayLike = np.zeros((len(output_axes), len(input_axes)), dtype=float) + for i, des_axis in enumerate(output_axes): + for j, src_axis in enumerate(input_axes): + if src_axis == self.output_to_input[des_axis]: + linear[i, j] = 1 + affine = AffineEdge( + linear=linear, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + ) + return affine + + def to_model(self) -> ozm06trans.MapAxis: + mapAxis: list[int] = [] + for out_ax in self.output.axes_names: + in_ax = self.output_to_input[out_ax] + in_idx = self.input.axes_names.index(in_ax) + mapAxis.append(in_idx) + return ozm06trans.MapAxis( + name=self.name, + mapAxis=tuple(mapAxis), + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + ) + + +class TranslationEdge(BaseTransfEdge): + """The Translation transformation from the NGFF specification.""" + + def __init__( + self, + name: str | None, + *, + translation: ArrayLike, + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Init the NgffTranslation object. + Parameters + ---------- + name + A human readable name for this transformation + translation + A vector of shape (input.num_axes,) specifying the translation along each axis. + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + if input.num_axes != output.num_axes: + raise ValueError("Number of input and output axes must be the same") + self.translation = translation + super().__init__(input=input, output=output, name=name) + + def __repr__(self) -> str: + return super().__repr__() + str(self.translation) + + def inverse(self) -> BaseTransfEdge: + return TranslationEdge( + translation=-self.translation, + input=self.output, + output=self.input, + name=self.name and f"{self.name}__inverse", + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + self._validate_transform_points_shapes(points) + return points + self.translation + + def to_affine(self) -> AffineEdge: + return AffineEdge( + linear=np.identity(self.input.num_axes), + translation=self.translation, + input=self.input, + output=self.output, + name=self.name and f"{self.name}__affine", + ) + + def to_model(self) -> ozm06trans.Translation: + return ozm06trans.Translation( + name=self.name, + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + translation=tuple(self.translation), + ) + + +class ScaleEdge(BaseTransfEdge): + """The Scale transformation from the NGFF specification.""" + + def __init__( + self, + name: str | None, + *, + scale: ArrayLike, + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Init the NgffScale object. + Parameters + ---------- + scale + A list of numbers or a vector specifying the scale along each axis. + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + if scale.shape != (input.num_axes,): + raise ValueError(f"scale should be of shape f{(input.num_axes,)}") + if input.num_axes != output.num_axes: + raise ValueError("input and output must have same number of dimensions") + self.scale = scale + super().__init__(input=input, output=output, name=name) + + def __repr__(self) -> str: + return super().__repr__() + str(self.scale) + + def inverse(self) -> ScaleEdge: + if any(s == 0 for s in self.scale): + raise ValueError(f"Scaling {self} is not invertible") + new_scale = 1 / self.scale + return ScaleEdge( + scale=new_scale, input=self.output, output=self.input, name=self.name and f"{self.name}__inverse" + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + self._validate_transform_points_shapes(points) + return points * self.scale + + def to_affine(self) -> AffineEdge: + return AffineEdge( + linear=np.diag(self.scale), input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + ) + + def to_model(self) -> ozm06trans.Scale: + return ozm06trans.Scale( + name=self.name, + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + scale=tuple(self.scale), + ) + + +class RotationEdge(BaseTransfEdge): + """The Rotation transformation from the NGFF specification.""" + + rotation: Final[ArrayLike] + + def __init__( + self, + name: str | None, + *, + linear_matrix: ArrayLike, + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Init the NgffRotation object. + Parameters + ---------- + linear_matrix + an array of shape (output.num_axes, input.num_axes) representing the rotation + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + """ + expected_shape = (output.num_axes, input.num_axes) + if linear_matrix.shape != expected_shape: + raise ValueError(f"linear matrix should have shape {expected_shape}") + if input.num_axes != output.num_axes: + raise ValueError("input and output should have the same numbe rof axes") + if not np.isclose(np.linalg.det(linear_matrix), 1.0): + raise ValueError("det(linear_matrix) should be ~= 1") + linear_matrix.flags.writeable = False + self.rotation = linear_matrix + super().__init__(input=input, output=output, name=name) + + def __repr__(self) -> str: + s = super().__repr__() + "\n" + s += "\n".join(str(row) for row in self.rotation) + return s + + def inverse(self) -> BaseTransfEdge: + return RotationEdge( + linear_matrix=self.rotation.T, + input=self.output, + output=self.input, + name=self.name and f"{self.name}__inverse", + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + self._validate_transform_points_shapes(points) + res = (self.rotation @ points.T).T + assert isinstance(res, np.ndarray) + return res + + def to_affine(self) -> AffineEdge: + return AffineEdge( + linear=self.rotation, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + ) + + def to_model(self) -> ozm06trans.Rotation: + return ozm06trans.Rotation( + name=self.name, + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + rotation=tuple(tuple(row) for row in self.rotation), + ) + + +class SequenceEdge(BaseTransfEdge): + """The Sequence transformation from the NGFF specification.""" + + def __init__( + self, + transformations: Sequence[BaseTransfEdge], + name: str | None, + ) -> None: + """ + Init the NgffSequence object. + + Parameters + ---------- + transformations + The transformations which compose the sequence. + """ + if len(transformations) == 0: + raise ValueError("Empty transformation list") + previous_transf = transformations[0] + for current_transf in transformations[1:]: + if previous_transf.output != current_transf.input: + raise ValueError(f"Mismatched input/output from {previous_transf} to {current_transf}") + previous_transf = current_transf + self.transformations = transformations + super().__init__( + input=transformations[0].input, + output=transformations[-1].output, + name=name, + ) + + def __repr__(self) -> str: + from textwrap import indent + + out = super().__repr__() + " [\n" + for t in self.transformations: + out += indent(repr(t), prefix=" ") + "\n" + out += "]" + return out + + def inverse(self) -> SequenceEdge: + return SequenceEdge( + [t.inverse() for t in reversed(self.transformations)], name=self.name and f"{self.name}__inverse" + ) + + def to_affine(self) -> AffineEdge: + composed = self.transformations[0].to_affine().affine + for t in self.transformations[1:]: + a = t.to_affine() + composed = a.affine @ composed + return AffineEdge.from_affine_matrix( + affine_matrix=composed, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + return self.to_affine().transform_points(points) # FIXME + + def to_model(self) -> ozm06trans.Sequence: + return ozm06trans.Sequence( + name=self.name, + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + transformations=tuple(t.to_model() for t in self.transformations), + ) + + +class ByDimensionEdge(BaseTransfEdge): + """The ByDimension transformation from the NGFF specification.""" + + transformations: Final[Sequence[BaseTransfEdge]] + + def __init__( + self, + name: str | None, + *, + transformations: Sequence[BaseTransfEdge], + input: CoordSystem, + output: CoordSystem, + ) -> None: + """ + Init the ByDimension object. + + Parameters + ---------- + transformations + A list of transformations, whose set of output coordinate systems partition the output coordinate system of + the ByDimension transformation. + input + The input coordinate system of the transformation. + output + The output coordinate system of the transformation. + """ + # we check that: + # 1. each input from each transformation in self.transformation must appear in the set of input axes + # 2. each output from each transformation in self.transformation must appear at most once in the set of output + # axes + input_axes = input.axes_names + output_axes = output.axes_names + defined_output_axes: set[str] = set() + for t in transformations: + for ax in t.input.axes_names: + if ax not in input_axes: + raise ValueError(f"By dimension axis {ax} not in {input_axes}") + for ax in t.output.axes_names: + if ax not in output_axes: + raise ValueError(f"Axis {ax} not in output axes {output_axes}") + if ax in defined_output_axes: + raise ValueError(f"Output axis {ax} is defined more than once") + defined_output_axes.add(ax) + if len(output_axes) != len(defined_output_axes): + raise ValueError("Not all outputs are mapped") + + self.transformations = tuple(transformations) + super().__init__(input=input, output=output, name=name) + + def __repr__(self) -> str: + from textwrap import indent + + out = super().__repr__() + " [\n" + for t in self.transformations: + out += indent(repr(t), prefix=" ") + "\n" + out += "]" + return out + + def inverse(self) -> BaseTransfEdge: + inverse_transformations = [t.inverse() for t in self.transformations] + return ByDimensionEdge( + transformations=inverse_transformations, + input=self.output, + output=self.input, + name=self.name and f"{self.name}__inverse", + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + input_axes = self.input.axes_names + output_axes = self.output.axes_names + self._validate_transform_points_shapes(points) + output_columns: dict[str, ArrayLike] = {} + for t in self.transformations: + input_columns = [points[:, input_axes.index(ax)] for ax in t.input.axes_names] + input_columns_stacked: ArrayLike = np.stack(input_columns, axis=1) + output_columns_t = t.transform_points(input_columns_stacked) + for ax, col in zip(t.output.axes_names, output_columns_t.T, strict=True): + output_columns[ax] = col # type: ignore[assignment] + output: ArrayLike = np.stack([output_columns[ax] for ax in output_axes], axis=1) + return output + + def to_affine(self) -> AffineEdge: + input_axes = self.input.axes_names + output_axes = self.output.axes_names + m = np.zeros((len(output_axes) + 1, len(input_axes) + 1)) + m[-1, -1] = 1 + for t in self.transformations: + t_affine = t.to_affine() + target_output_indices = [output_axes.index(ax) for ax in t.output.axes_names if ax in output_axes] + source_output_indices = [t.output.axes_names.index(ax) for ax in t.output.axes_names] + target_input_indices = [input_axes.index(ax) for ax in t.input.axes_names] + [-1] + m[np.ix_(target_output_indices, target_input_indices)] = t_affine.affine[source_output_indices, :] + return AffineEdge.from_affine_matrix( + affine_matrix=m, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + ) + + def to_model(self) -> ozm06trans.ByDimension: + by_dim_transfs: list[ozm06trans.ByDimensionTransform] = [] + for t in self.transformations: + input_axes = tuple(self.input.axes_names.index(ax_name) for ax_name in t.input.axes_names) + output_axes = tuple(self.output.axes_names.index(ax_name) for ax_name in t.output.axes_names) + by_dim_transfs.append( + ozm06trans.ByDimensionTransform( + input_axes=input_axes, + output_axes=output_axes, + transformation=t.to_model(), + ) + ) + return ozm06trans.ByDimension( + name=self.name, + input=self.input.to_model_cs_ident(), + output=self.output.to_model_cs_ident(), + transformations=tuple(by_dim_transfs), + ) + + +class CsGen: + """A coordinate system generator + + Use it to create coordinate systems on the fly while avoiding + repeating names + """ + + def __init__(self, base_name: str): + self._base_name = base_name + self._cs_count: int = 0 + super().__init__() + + def generate(self, *, num_axes: int) -> CoordSystem: + out = CoordSystem( + name=f"{self._base_name}{self._cs_count}", + axes=[ + Axis( + name=f"axis_{ax_idx}", + type="space", # FIXME + ) + for ax_idx in range(num_axes) + ], + virtual=True, + ) + self._cs_count += 1 + return out + + def generate_like(self, other: CoordSystem) -> CoordSystem: + out = CoordSystem( + name=f"{self._base_name}{self._cs_count}", + axes=[ + Axis( + name=axis.name, + type=axis.type, + unit=axis.unit, + long_name=axis.long_name, + ) + for axis in other.axes + ], + virtual=True, + ) + self._cs_count += 1 + return out + + +def parse_identity( + model: ozm06trans.Identity, + *, + input: CoordSystem, + out: CoordSystem | CsGen, +) -> IdentityEdge: + output = out.generate_like(input) if isinstance(out, CsGen) else out + return IdentityEdge(name=model.name, input=input, output=output) + + +def parse_translation( + model: ozm06trans.Translation, + *, + input: CoordSystem, + out: CoordSystem | CsGen, +) -> TranslationEdge: + output = out.generate_like(input) if isinstance(out, CsGen) else out + return TranslationEdge( + translation=np.asarray(model.translation, dtype=float), + input=input, + output=output, + name=input.name, + ) + + +def parse_scale( + model: ozm06trans.Scale, + *, + input: CoordSystem, + out: CoordSystem | CsGen, +) -> ScaleEdge: + output = out.generate_like(input) if isinstance(out, CsGen) else out + return ScaleEdge( + scale=np.asarray(model.scale, dtype=float), + input=input, + output=output, + name=model.name, + ) + + +def parse_map_axis( + model: ozm06trans.MapAxis, + *, + input: CoordSystem, + out: CoordSystem | CsGen, +) -> MapAxisEdge: + output = out.generate(num_axes=len(model.mapAxis)) if isinstance(out, CsGen) else out + return MapAxisEdge( + input=input, + output=output, + name=model.name, + output_to_input={ # FIXME: double check this. Feels like we depend a lot on order + output.axes[output_axis].name: input.axes[input_axis].name + for output_axis, input_axis in enumerate(model.mapAxis) + }, + ) + + +def parse_affine( + model: ozm06trans.Affine, + *, + input: CoordSystem, + output: CoordSystem | CsGen, +) -> AffineEdge: + num_output_axes = len(model.affine_matrix) # spec doesn't save last row + output = output.generate(num_axes=num_output_axes) if isinstance(output, CsGen) else output + affine_array = np.asarray(model.affine_matrix, dtype=float) + return AffineEdge( + name=model.name, linear=affine_array[:, :-1], translation=affine_array[:, -1], input=input, output=output + ) + + +def parse_rotation( + model: ozm06trans.Rotation, + *, + input: CoordSystem, + out: CoordSystem | CsGen, +) -> RotationEdge: + num_output_axes = len(model.rotation_matrix) + output = out.generate(num_axes=num_output_axes) if isinstance(out, CsGen) else out + return RotationEdge( + name=model.name, + linear_matrix=np.asarray(model.rotation_matrix, dtype=float), + input=input, + output=output, + ) + + +def parse_sequence( + model: ozm06trans.Sequence, + *, + input: CoordSystem, + output: CoordSystem | CsGen, +) -> SequenceEdge: + parsed_inners: list[BaseTransfEdge] = [] + + base_name = "intermediate" + ("" if not model.name else f"_for_{model.name}") + cs_gen: CsGen = output if isinstance(output, CsGen) else CsGen(base_name=base_name) + parsed = parse_ngff_transf( + input=input, + model=model.transformations[0], + output=cs_gen if len(model.transformations) > 1 else output, + ) + parsed_inners.append(parsed) + + for t in model.transformations[1:-1]: + parsed = parse_ngff_transf(input=parsed.output, output=cs_gen, model=t) + parsed_inners.append(parsed) + + if len(model.transformations) > 1: + parsed = parse_ngff_transf(input=parsed.output, output=output, model=model.transformations[-1]) + parsed_inners.append(parsed) + + return SequenceEdge(name=model.name, transformations=parsed_inners) + + +def parse_by_dimension( + model: ozm06trans.ByDimension, + *, + input: CoordSystem, + output: CoordSystem | CsGen, +) -> ByDimensionEdge: + if not isinstance(output, CoordSystem): + max_out_idx = max(ax_idx for t in model.transformations for ax_idx in t.output_axes) + output = output.generate(num_axes=max_out_idx + 1) + + piecewise_transforms: list[BaseTransfEdge] = [] + for t in model.transformations: + inp_axes = [input.axes[i] for i in t.input_axes] + partial_input = CoordSystem( + axes=inp_axes, + name=f"{input.name}_{','.join(ax.name for ax in inp_axes)}", + virtual=True, + ) + + out_axes = [output.axes[i] for i in t.output_axes] + partial_out = CoordSystem( + axes=out_axes, + name=f"{output.name}_{','.join(ax.name for ax in inp_axes)}", + virtual=True, + ) + + parsed_t = parse_ngff_transf(model=t.transformation, input=partial_input, output=partial_out) + piecewise_transforms.append(parsed_t) + + return ByDimensionEdge( + input=input, + output=output, + name=model.name, + transformations=piecewise_transforms, + ) + + +def parse_ngff_transf( + input: CoordSystem, + model: ozm06trans.AnyTransform, + output: CoordSystem | CsGen, +) -> BaseTransfEdge: + if isinstance(model, ozm06trans.Identity): + return parse_identity(model, input=input, out=output) + elif isinstance(model, ozm06trans.Translation): + return parse_translation(model, input=input, out=output) + elif isinstance(model, ozm06trans.Scale): + return parse_scale(model, input=input, out=output) + elif isinstance(model, ozm06trans.MapAxis): + return parse_map_axis(model, input=input, out=output) + elif isinstance(model, ozm06trans.Affine): + return parse_affine(model, input=input, output=output) + elif isinstance(model, ozm06trans.Rotation): + return parse_rotation(model, input=input, out=output) + elif isinstance(model, ozm06trans.Sequence): + return parse_sequence(model, input=input, output=output) + elif isinstance(model, ozm06trans.ByDimension): + return parse_by_dimension(model, input=input, output=output) + else: + raise NotImplementedError(f"Unsupported transformation: {model.type}") diff --git a/src/spatialdata/transformations/graph/vert.py b/src/spatialdata/transformations/graph/vert.py new file mode 100644 index 000000000..8819d04e9 --- /dev/null +++ b/src/spatialdata/transformations/graph/vert.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Final, Literal + +import ome_zarr.classes.image as ozi +import pydantic as pyd + + +class AxisParsingException(Exception): + pass + + +class Axis: + """ + Representation of an axis, following the NGFF specification. + + Attributes + ---------- + name + name of the axis. + type + type of the axis. Should be in ["channel", "space"]. + unit + unit of the axis. For a set of valid options see https://ngff.openmicroscopy.org/ + long_name: + a longer, human-friendly name for this axis + """ + + name: Final[str] + type: Final[Literal["space", "channel"]] + unit: Final[str | None] + long_name: Final[str | None] + + class LegacyModel(pyd.BaseModel): + name: Literal["x", "y", "z", "c"] + type: Literal["space", "channel"] + + def __init__( + self, *, name: str, type: Literal["space", "channel"], unit: str | None = None, long_name: str | None = None + ): + self.name = name + self.type = type + self.unit = unit + self.long_name = long_name + + def cloned_with(self, *, unit: str | None) -> Axis: + return Axis(name=self.name, type=self.type, unit=unit or self.unit, long_name=self.long_name) + + def __hash__(self) -> int: + return hash((self.name, self.type, self.unit, self.long_name)) + + def __repr__(self) -> str: + return f"NgffAxis(name={self.name}, type={self.type})" + + def __eq__(self, value: object, /) -> bool: + if not isinstance(value, Axis): + return False + return ( + self.name == value.name + and self.type == value.type + and self.unit == value.unit + and self.long_name == value.long_name + ) + + @classmethod + def try_from_model(cls, model: ozi.Axis) -> Axis: + name = model.name + if name is None: + raise AxisParsingException("Axis doesn't have a name") + if model.type != "channel" and model.type != "space": + raise AxisParsingException(f"Can't handle axis of type {model.type}") + if not isinstance(model.unit, str): + raise AxisParsingException("Can't handle axis unit") + return Axis( + name=name, + type=model.type, + unit=model.unit, + long_name=model.longName, + ) + + @classmethod + def try_from_dict(cls, d: pyd.JsonValue) -> Axis: + model = ozi.Axis.model_validate(d) + return Axis.try_from_model(model) + + def to_model(self) -> ozi.Axis: + return ozi.Axis( + discrete=False, + longName=self.long_name, + name=self.name, + type=self.type, + unit=self.unit, + ) + + +class CoordSystemParsingException(Exception): + pass + + +class CoordSystem: + """ + Representation of a coordinate system, following the NGFF specification. + + Parameters + ---------- + name + name of the coordinate system + axes + names of the axes of the coordinate system + """ + + name: Final[str] + axes: Final[tuple[Axis, ...]] + + virtual: Final[bool] + """A virtual coordinate system exists as an intermediate step between + non-virtual coordinate systems and is usually ignored during serialization""" + + class LegacyAxes: + pass + + def __init__(self, name: str, axes: Sequence[Axis], virtual: bool = False): + self.name = name + self.axes = tuple(axes) + self.virtual = virtual + if len(self.axes) != len({axis.name for axis in self.axes}): + raise ValueError("Axes names must be unique") + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.name!r}, {self.axes})" + + def __hash__(self) -> int: + return hash((self.name, self.axes, self.virtual)) + + @classmethod + def try_from_model(cls, model: ozi.CoordinateSystem) -> CoordSystem: + axes: list[Axis] = [] + for axis in model.axes: + if isinstance(parsed := Axis.try_from_model(axis), Exception): + raise CoordSystemParsingException(parsed) # FIXME + axes.append(parsed) + return CoordSystem( + name=model.name, + axes=axes, + ) + + @classmethod + def try_from_model_or_default[T](cls, model: ozi.CoordinateSystem | None, *, default: T) -> CoordSystem | T: + if model is not None: + return CoordSystem.try_from_model(model) + return default + + def to_model(self) -> ozi.CoordinateSystem | None: + if self.virtual: + return None + return ozi.CoordinateSystem( + name=self.name, + axes=tuple(ax.to_model() for ax in self.axes), + ) + + def to_model_cs_ident(self) -> ozi.CoordinateSystemIdentifier | None: + input = self.to_model() + if input is None: + return None + return ozi.CoordinateSystemIdentifier(name=input.name) + + @property + def num_axes(self) -> int: + return len(self.axes) + + @property + def axes_names(self) -> tuple[str, ...]: + """Get axes' names""" + return tuple([ax.name for ax in self.axes]) + + @property + def axes_types(self) -> tuple[str, ...]: + """Get axes' types""" + return tuple([ax.type for ax in self.axes]) + + def has_axis(self, name: str) -> bool: + """ + Check the coordinate system has an axis of the given name. + + Parameters + ---------- + name + name of the axis. + """ + return any(axis.name == name for axis in self.axes) + + def get_axis(self, name: str) -> Axis | None: + """Get the axis by name""" + for axis in self.axes: + if axis.name == name: + return axis + return None + + def get_spatial_axes(self) -> Sequence[Axis]: + return [axis for axis in self.axes if axis.type == "space"] From b351aa0b51c21ae296acf302b7e6ae8bf24a288e Mon Sep 17 00:00:00 2001 From: Tomaz Vieira Date: Wed, 2 Sep 2026 16:11:38 +0000 Subject: [PATCH 2/3] addresses some PR comments, simplifies MapAxis --- src/spatialdata/transformations/graph/edge.py | 162 ++++++++++-------- src/spatialdata/transformations/graph/vert.py | 19 +- 2 files changed, 94 insertions(+), 87 deletions(-) diff --git a/src/spatialdata/transformations/graph/edge.py b/src/spatialdata/transformations/graph/edge.py index b01a11e97..42396175f 100644 --- a/src/spatialdata/transformations/graph/edge.py +++ b/src/spatialdata/transformations/graph/edge.py @@ -9,7 +9,6 @@ import numpy as np import ome_zarr_models.v06.coordinate_transforms as ozm06trans import pydantic as pyd -import xarray as xr from spatialdata._types import ArrayLike from spatialdata.transformations.graph.vert import Axis, CoordSystem @@ -48,7 +47,7 @@ def __repr__(self) -> str: return f"{type(self).__name__} ({domain} -> {codomain})" @abstractmethod - def inverse(self) -> BaseTransfEdge: + def inverse(self, name: str | None = None) -> BaseTransfEdge: """Return the inverse of the transformation.""" @abstractmethod @@ -63,23 +62,23 @@ def transform_points(self, points: ArrayLike) -> ArrayLike: """ @abstractmethod - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: """Convert the transformation to an affine transformation, whenever the conversion can be made.""" - def _validate_transform_points_shapes(self, points: xr.DataArray | xr.DataTree | ArrayLike) -> None: + def _validate_transform_points_shapes(self, points: ArrayLike) -> None: """ - Validate if the shape of the points (coordinats to be transformed) are consistent with the input size of the + Validate if the shape of the points (coordinates to be transformed) are consistent with the input size of the transformation. """ input_size = len(self.input.axes) if len(points.shape) != 2 or points.shape[1] != input_size: raise ValueError( f"points must be a tensor of shape (n, d), where n is the number of points and d is the " - f"the number of spatial dimensions. Points shape: {points.shape}, input size: {input_size}" + f"the number of dimensions. Points shape: {points.shape}, input size: {input_size}" ) # order of the composition: self is applied first, then the transformation passed as argument - def compose_with(self, transformation: BaseTransfEdge) -> BaseTransfEdge: + def compose_with(self, transformation: BaseTransfEdge, name: str | None) -> BaseTransfEdge: """ Compose the transfomation object with another transformation @@ -96,7 +95,7 @@ def compose_with(self, transformation: BaseTransfEdge) -> BaseTransfEdge: ------- Self is applied first, then the transformation passed as argument. """ - return SequenceEdge([self, transformation], name=None) # FIXME: no name? + return SequenceEdge([self, transformation], name=name) @abstractmethod def to_model(self) -> ozm06trans.AnyTransform: @@ -112,8 +111,8 @@ class AffineEdge(BaseTransfEdge): def __init__( self, - name: str | None, *, + name: str | None = None, linear: ArrayLike, translation: ArrayLike | None = None, input: CoordSystem, @@ -169,7 +168,7 @@ def from_affine_matrix( input: CoordSystem, output: CoordSystem, ) -> AffineEdge: - """Creates an AffineEdge from a raw affine matrix + """Creates an AffineEdge from a raw affine matrix in homogenous coordinates Parameters ---------- @@ -193,14 +192,14 @@ def from_affine_matrix( name=name, ) - def inverse(self) -> BaseTransfEdge: + def inverse(self, name: str | None = None) -> BaseTransfEdge: inv = np.linalg.inv(self.affine) return AffineEdge( linear=inv[:-1, :-1], translation=inv[-1, :-1], input=self.output, output=self.input, - name=self.name and f"{self.name}__affine", + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: @@ -211,8 +210,10 @@ def transform_points(self, points: ArrayLike) -> ArrayLike: assert isinstance(res, np.ndarray) return res - def to_affine(self) -> AffineEdge: - return self + def to_affine(self, name: str | None = None) -> AffineEdge: + return AffineEdge( + input=self.input, output=self.output, linear=self.linear, translation=self.translation, name=name + ) def to_model(self) -> ozm06trans.Affine: return ozm06trans.Affine( @@ -247,19 +248,19 @@ def __init__( raise ValueError("Input and output must have the same number of dimensions") super().__init__(input=input, output=output, name=name) - def inverse(self) -> BaseTransfEdge: - return IdentityEdge(input=self.output, output=self.input, name=self.name and f"{self.name}__inverse") + def inverse(self, name: str | None = None) -> BaseTransfEdge: + return IdentityEdge(input=self.output, output=self.input, name=name) def transform_points(self, points: ArrayLike) -> ArrayLike: self._validate_transform_points_shapes(points) return points - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: return AffineEdge( linear=np.eye(self.input.num_axes), input=self.input, output=self.output, - name=self.name and f"{self.name}__affine", + name=name, ) def to_model(self) -> ozm06trans.Identity: @@ -270,6 +271,13 @@ def to_model(self) -> ozm06trans.Identity: ) +class UnmappableCoordSystemsError(Exception): + def __init__(self, input: CoordSystem, output: CoordSystem) -> None: + self.input = input + self.output = output + super().__init__("Output axes can't be mapped to input axes") + + class MapAxisEdge(BaseTransfEdge): """The MapAxis transformation from the NGFF specification.""" @@ -277,7 +285,6 @@ def __init__( self, name: str | None, *, - output_to_input: dict[str, str], input: CoordSystem, output: CoordSystem, ) -> None: @@ -287,67 +294,56 @@ def __init__( ---------- name A human readable name for this transformation - output_to_input - A dictionary mapping the output axes (keys) to the input axes (values). input Input coordinate system of the transformation. output - Output coordinate system of the transformation. + Output coordinate system of the transformation, whose axes + must be a shuffling of `input` """ - for out_ax, inp_ax in output_to_input.items(): - if not input.has_axis(inp_ax): - raise ValueError(f"input has no axis named {inp_ax}") - if not output.has_axis(out_ax): - raise ValueError(f"output has no axis named {out_ax}") - if not (len(output_to_input) == output.num_axes == input.num_axes): - raise ValueError("input_to_output, input and output must have the same number of axes entries") - if len(set(output_to_input.values())) != len(output_to_input): - raise ValueError("input_to_output must map unique inputs to unique outputs") - - self.output_to_input = output_to_input + + if set(input.axes) != set(output.axes): + raise UnmappableCoordSystemsError(input=input, output=output) super().__init__(input=input, output=output, name=name) def __repr__(self) -> str: s = super().__repr__() + "\n" - s += "\n".join(f" {out} <- {inp}\n" for out, inp in self.output_to_input.items()) + s += "\n".join( + f" {out.name} <- {inp.name}\n" for out, inp in zip(self.output.axes, self.input.axes, strict=True) + ) return s - def inverse(self) -> BaseTransfEdge: + def inverse(self, name: str | None = None) -> BaseTransfEdge: return MapAxisEdge( - output_to_input={v: k for k, v in self.output_to_input.items()}, input=self.output, output=self.input, - name=self.name and f"{self.name}__inverse", + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: - input_axes = self.input.axes_names - output_axes = self.output.axes_names self._validate_transform_points_shapes(points) - new_indices = [input_axes.index(self.output_to_input[ax]) for ax in output_axes] + new_indices = [self.input.axes.index(out_ax.name) for out_ax in self.output.axes] mapped = points[:, new_indices] assert isinstance(mapped, np.ndarray) return mapped - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: input_axes = self.input.axes_names output_axes = self.output.axes_names linear: ArrayLike = np.zeros((len(output_axes), len(input_axes)), dtype=float) for i, des_axis in enumerate(output_axes): for j, src_axis in enumerate(input_axes): - if src_axis == self.output_to_input[des_axis]: + if src_axis == des_axis: linear[i, j] = 1 affine = AffineEdge( - linear=linear, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + linear=linear, + input=self.input, + output=self.output, + name=name, ) return affine def to_model(self) -> ozm06trans.MapAxis: - mapAxis: list[int] = [] - for out_ax in self.output.axes_names: - in_ax = self.output_to_input[out_ax] - in_idx = self.input.axes_names.index(in_ax) - mapAxis.append(in_idx) + mapAxis: list[int] = [self.input.axes.index(out_ax) for out_ax in self.output.axes] return ozm06trans.MapAxis( name=self.name, mapAxis=tuple(mapAxis), @@ -388,25 +384,25 @@ def __init__( def __repr__(self) -> str: return super().__repr__() + str(self.translation) - def inverse(self) -> BaseTransfEdge: + def inverse(self, name: str | None = None) -> BaseTransfEdge: return TranslationEdge( translation=-self.translation, input=self.output, output=self.input, - name=self.name and f"{self.name}__inverse", + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: self._validate_transform_points_shapes(points) return points + self.translation - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: return AffineEdge( linear=np.identity(self.input.num_axes), translation=self.translation, input=self.input, output=self.output, - name=self.name and f"{self.name}__affine", + name=name, ) def to_model(self) -> ozm06trans.Translation: @@ -450,21 +446,27 @@ def __init__( def __repr__(self) -> str: return super().__repr__() + str(self.scale) - def inverse(self) -> ScaleEdge: + def inverse(self, name: str | None = None) -> ScaleEdge: if any(s == 0 for s in self.scale): raise ValueError(f"Scaling {self} is not invertible") new_scale = 1 / self.scale return ScaleEdge( - scale=new_scale, input=self.output, output=self.input, name=self.name and f"{self.name}__inverse" + scale=new_scale, + input=self.output, + output=self.input, + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: self._validate_transform_points_shapes(points) return points * self.scale - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: return AffineEdge( - linear=np.diag(self.scale), input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + linear=np.diag(self.scale), + input=self.input, + output=self.output, + name=name, ) def to_model(self) -> ozm06trans.Scale: @@ -516,12 +518,12 @@ def __repr__(self) -> str: s += "\n".join(str(row) for row in self.rotation) return s - def inverse(self) -> BaseTransfEdge: + def inverse(self, name: str | None = None) -> BaseTransfEdge: return RotationEdge( linear_matrix=self.rotation.T, input=self.output, output=self.input, - name=self.name and f"{self.name}__inverse", + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: @@ -530,9 +532,12 @@ def transform_points(self, points: ArrayLike) -> ArrayLike: assert isinstance(res, np.ndarray) return res - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: return AffineEdge( - linear=self.rotation, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + linear=self.rotation, + input=self.input, + output=self.output, + name=name, ) def to_model(self) -> ozm06trans.Rotation: @@ -583,18 +588,22 @@ def __repr__(self) -> str: out += "]" return out - def inverse(self) -> SequenceEdge: + def inverse(self, name: str | None = None) -> SequenceEdge: return SequenceEdge( - [t.inverse() for t in reversed(self.transformations)], name=self.name and f"{self.name}__inverse" + [t.inverse() for t in reversed(self.transformations)], + name=name, ) - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: composed = self.transformations[0].to_affine().affine for t in self.transformations[1:]: a = t.to_affine() composed = a.affine @ composed return AffineEdge.from_affine_matrix( - affine_matrix=composed, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + affine_matrix=composed, + input=self.input, + output=self.output, + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: @@ -667,13 +676,13 @@ def __repr__(self) -> str: out += "]" return out - def inverse(self) -> BaseTransfEdge: + def inverse(self, name: str | None = None) -> BaseTransfEdge: inverse_transformations = [t.inverse() for t in self.transformations] return ByDimensionEdge( transformations=inverse_transformations, input=self.output, output=self.input, - name=self.name and f"{self.name}__inverse", + name=name, ) def transform_points(self, points: ArrayLike) -> ArrayLike: @@ -690,7 +699,7 @@ def transform_points(self, points: ArrayLike) -> ArrayLike: output: ArrayLike = np.stack([output_columns[ax] for ax in output_axes], axis=1) return output - def to_affine(self) -> AffineEdge: + def to_affine(self, name: str | None = None) -> AffineEdge: input_axes = self.input.axes_names output_axes = self.output.axes_names m = np.zeros((len(output_axes) + 1, len(input_axes) + 1)) @@ -702,7 +711,10 @@ def to_affine(self) -> AffineEdge: target_input_indices = [input_axes.index(ax) for ax in t.input.axes_names] + [-1] m[np.ix_(target_output_indices, target_input_indices)] = t_affine.affine[source_output_indices, :] return AffineEdge.from_affine_matrix( - affine_matrix=m, input=self.input, output=self.output, name=self.name and f"{self.name}__affine" + affine_matrix=m, + input=self.input, + output=self.output, + name=name, ) def to_model(self) -> ozm06trans.ByDimension: @@ -816,15 +828,19 @@ def parse_map_axis( input: CoordSystem, out: CoordSystem | CsGen, ) -> MapAxisEdge: - output = out.generate(num_axes=len(model.mapAxis)) if isinstance(out, CsGen) else out + if isinstance(out, CoordSystem): + output = out + else: + dummy_cs = out.generate(num_axes=len(model.mapAxis)) + output = CoordSystem( + name=dummy_cs.name, + axes=[input.axes[i] for i in model.mapAxis], + virtual=True, + ) return MapAxisEdge( input=input, output=output, name=model.name, - output_to_input={ # FIXME: double check this. Feels like we depend a lot on order - output.axes[output_axis].name: input.axes[input_axis].name - for output_axis, input_axis in enumerate(model.mapAxis) - }, ) diff --git a/src/spatialdata/transformations/graph/vert.py b/src/spatialdata/transformations/graph/vert.py index 8819d04e9..128c50706 100644 --- a/src/spatialdata/transformations/graph/vert.py +++ b/src/spatialdata/transformations/graph/vert.py @@ -4,7 +4,7 @@ from typing import Final, Literal import ome_zarr.classes.image as ozi -import pydantic as pyd +import ome_zarr_models.v06.coordinate_transforms as ozm06ct class AxisParsingException(Exception): @@ -32,10 +32,6 @@ class Axis: unit: Final[str | None] long_name: Final[str | None] - class LegacyModel(pyd.BaseModel): - name: Literal["x", "y", "z", "c"] - type: Literal["space", "channel"] - def __init__( self, *, name: str, type: Literal["space", "channel"], unit: str | None = None, long_name: str | None = None ): @@ -64,13 +60,13 @@ def __eq__(self, value: object, /) -> bool: ) @classmethod - def try_from_model(cls, model: ozi.Axis) -> Axis: + def try_from_model(cls, model: ozm06ct.Axis) -> Axis: name = model.name if name is None: raise AxisParsingException("Axis doesn't have a name") if model.type != "channel" and model.type != "space": raise AxisParsingException(f"Can't handle axis of type {model.type}") - if not isinstance(model.unit, str): + if not isinstance(model.unit, (str, type(None))): raise AxisParsingException("Can't handle axis unit") return Axis( name=name, @@ -79,13 +75,8 @@ def try_from_model(cls, model: ozi.Axis) -> Axis: long_name=model.longName, ) - @classmethod - def try_from_dict(cls, d: pyd.JsonValue) -> Axis: - model = ozi.Axis.model_validate(d) - return Axis.try_from_model(model) - - def to_model(self) -> ozi.Axis: - return ozi.Axis( + def to_model(self) -> ozm06ct.Axis: + return ozm06ct.Axis( discrete=False, longName=self.long_name, name=self.name, From 1d4d7cb8362fd86ab490b19d866442aac624360d Mon Sep 17 00:00:00 2001 From: Tomaz Vieira Date: Thu, 3 Sep 2026 16:32:06 +0000 Subject: [PATCH 3/3] Adds ProjectAxisEdge, makes inverse return optional --- src/spatialdata/transformations/graph/edge.py | 160 +++++++++++++++--- 1 file changed, 132 insertions(+), 28 deletions(-) diff --git a/src/spatialdata/transformations/graph/edge.py b/src/spatialdata/transformations/graph/edge.py index 42396175f..74664cff7 100644 --- a/src/spatialdata/transformations/graph/edge.py +++ b/src/spatialdata/transformations/graph/edge.py @@ -31,8 +31,8 @@ class BaseTransfEdge(ABC): def __init__( self, - name: str | None, *, + name: str | None = None, input: CoordSystem, output: CoordSystem, ) -> None: @@ -47,8 +47,8 @@ def __repr__(self) -> str: return f"{type(self).__name__} ({domain} -> {codomain})" @abstractmethod - def inverse(self, name: str | None = None) -> BaseTransfEdge: - """Return the inverse of the transformation.""" + def inverse(self, name: str | None = None) -> BaseTransfEdge | None: + """Return the inverse of the transformation if it exists""" @abstractmethod def transform_points(self, points: ArrayLike) -> ArrayLike: @@ -192,8 +192,26 @@ def from_affine_matrix( name=name, ) - def inverse(self, name: str | None = None) -> BaseTransfEdge: - inv = np.linalg.inv(self.affine) + @classmethod + def mapping(cls, input: CoordSystem, output: CoordSystem, name: str | None = None) -> AffineEdge: + linear: ArrayLike = np.zeros((output.num_axes, input.num_axes), dtype=float) + for i, des_axis in enumerate(output.axes): + for j, src_axis in enumerate(input.axes): + if src_axis.name == des_axis.name: # FIXME: compare the entire axis? + linear[i, j] = 1 + return AffineEdge( + linear=linear, + input=input, + output=output, + name=name, + ) + + def inverse(self, name: str | None = None) -> BaseTransfEdge | None: + try: + # FIXME: I think there are more efficient/precise ways to invert a matrix + inv = np.linalg.inv(self.affine) + except np.linalg.LinAlgError: + return None return AffineEdge( linear=inv[:-1, :-1], translation=inv[-1, :-1], @@ -327,20 +345,7 @@ def transform_points(self, points: ArrayLike) -> ArrayLike: return mapped def to_affine(self, name: str | None = None) -> AffineEdge: - input_axes = self.input.axes_names - output_axes = self.output.axes_names - linear: ArrayLike = np.zeros((len(output_axes), len(input_axes)), dtype=float) - for i, des_axis in enumerate(output_axes): - for j, src_axis in enumerate(input_axes): - if src_axis == des_axis: - linear[i, j] = 1 - affine = AffineEdge( - linear=linear, - input=self.input, - output=self.output, - name=name, - ) - return affine + return AffineEdge.mapping(input=self.input, output=self.output, name=name) def to_model(self) -> ozm06trans.MapAxis: mapAxis: list[int] = [self.input.axes.index(out_ax) for out_ax in self.output.axes] @@ -352,6 +357,97 @@ def to_model(self) -> ozm06trans.MapAxis: ) +class AxisNotInCoordSystemError(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 ProjectAxisEdge(BaseTransfEdge): + dropped_inputs: Final[set[Axis]] + created_outputs: Final[set[Axis]] + + def __init__( + self, + *, + name: str | None = None, + input: CoordSystem, + output: CoordSystem, + dropped_inputs: set[Axis], + created_outputs: set[Axis], + ) -> None: + for axis in dropped_inputs: + if axis not in input.axes: + raise AxisNotInCoordSystemError(axis=axis, cs=input) + for axis in dropped_inputs: + if axis not in output.axes: + raise AxisNotInCoordSystemError(axis=axis, cs=output) + self.dropped_inputs = set(dropped_inputs) + self.created_outputs = set(created_outputs) + super().__init__(name=name, input=input, output=output) + + def to_affine(self, name: str | None = None) -> AffineEdge: + linear = np.zeros((self.output.num_axes, self.input.num_axes), dtype=float) + + for out_idx, out_ax in enumerate(self.output.axes): + if out_ax in self.created_outputs: + continue + for in_idx, in_ax in enumerate(self.input.axes): + if in_ax not in self.dropped_inputs: + linear[out_idx, in_idx] = 1 + + return AffineEdge(name=name, input=self.input, output=self.output, linear=linear) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + return self.to_affine().transform_points(points) + + def inverse(self, name: str | None = None) -> BaseTransfEdge | None: + # FIXME: there may be other cases where this is invertible + if self.input.num_axes != self.output.num_axes: + return None + if len(self.dropped_inputs) > 0: + return None + if len(self.created_outputs) > 0: + return None + return ProjectAxisEdge( + input=self.output, + output=self.input, + dropped_inputs=set(), + created_outputs=set(), + name=name, + ) + + def to_model(self) -> ozm06trans.ProjectAxis: + return ozm06trans.ProjectAxis( + createdOutputs=tuple(self.output.axes.index(co) for co in self.created_outputs) or None, + droppedInputs=tuple(self.input.axes.index(di) for di in self.dropped_inputs) or None, + ) + + +def parse_project_axis( + model: ozm06trans.ProjectAxis, + *, + input: CoordSystem, + out: CoordSystem | CsGen, +) -> ProjectAxisEdge: + if isinstance(out, CoordSystem): + output = out + else: + num_dropped_inputs = len(model.droppedInputs or ()) + num_created_outputs = len(model.droppedInputs or ()) + num_output_axes = input.num_axes - num_dropped_inputs + num_created_outputs + output = out.generate(num_axes=num_output_axes) + + return ProjectAxisEdge( + created_outputs={output.axes[i] for i in model.createdOutputs or ()}, + dropped_inputs={input.axes[i] for i in model.droppedInputs or ()}, + input=input, + output=output, + name=model.name, + ) + + class TranslationEdge(BaseTransfEdge): """The Translation transformation from the NGFF specification.""" @@ -446,9 +542,9 @@ def __init__( def __repr__(self) -> str: return super().__repr__() + str(self.scale) - def inverse(self, name: str | None = None) -> ScaleEdge: + def inverse(self, name: str | None = None) -> ScaleEdge | None: if any(s == 0 for s in self.scale): - raise ValueError(f"Scaling {self} is not invertible") + return None new_scale = 1 / self.scale return ScaleEdge( scale=new_scale, @@ -588,11 +684,14 @@ def __repr__(self) -> str: out += "]" return out - def inverse(self, name: str | None = None) -> SequenceEdge: - return SequenceEdge( - [t.inverse() for t in reversed(self.transformations)], - name=name, - ) + def inverse(self, name: str | None = None) -> SequenceEdge | None: + inverted: list[BaseTransfEdge] = [] + for t in self.transformations: + inv = t.inverse() + if inv is None: + return None + inverted.append(inv) + return SequenceEdge(inverted, name=name) def to_affine(self, name: str | None = None) -> AffineEdge: composed = self.transformations[0].to_affine().affine @@ -676,8 +775,13 @@ def __repr__(self) -> str: out += "]" return out - def inverse(self, name: str | None = None) -> BaseTransfEdge: - inverse_transformations = [t.inverse() for t in self.transformations] + def inverse(self, name: str | None = None) -> BaseTransfEdge | None: + inverse_transformations: list[BaseTransfEdge] = [] + for t in self.transformations: + inv = t.inverse() + if inv is None: + return None + inverse_transformations.append(inv) return ByDimensionEdge( transformations=inverse_transformations, input=self.output,