diff --git a/.gitignore b/.gitignore index d157ddaae..f44cfbecf 100644 --- a/.gitignore +++ b/.gitignore @@ -54,3 +54,29 @@ _version.py # benchmarking and profiling .asv/ profile.speedscope.json + +# other +node_modules/ + +.mypy_cache +.ruff_cache + +# Kilo and plans +.kilo/ +plans/ + +# test coverage +.coverage +htmlcov/ + +# local hatch config +hatch.toml + +# Kilo and plans +.kilo/ +plans/ + +# test coverage +.coverage +htmlcov/ + diff --git a/pyproject.toml b/pyproject.toml index 4cc02bcb4..9ccebb29c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,8 @@ dependencies = [ "networkx", "numba>=0.55", "numpy", - "ome-zarr>=0.16", + "ome-zarr>=0.21.0", + "ome-zarr-models>=1.8", "pandas", "pooch", "pyarrow", @@ -51,7 +52,7 @@ dependencies = [ "spatial-image>=1.2.3", "typing-extensions>=4.8", "universal-pathlib>=0.2.6", - "xarray>=2024.10", + "xarray>=2025.4", "xarray-spatial>=0.3.5", "zarr>=3", ] diff --git a/src/spatialdata/__init__.py b/src/spatialdata/__init__.py index 96e7032dd..1c04924a4 100644 --- a/src/spatialdata/__init__.py +++ b/src/spatialdata/__init__.py @@ -59,7 +59,6 @@ "polygon_query": "spatialdata._core.query.spatial_query", # _core.spatialdata "SpatialData": "spatialdata._core.spatialdata", - # _io._utils "get_dask_backing_files": "spatialdata._io._utils", # _io.format "SpatialDataFormatType": "spatialdata._io.format", diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py index 9f5c91c7a..34e25aa64 100644 --- a/src/spatialdata/_core/spatialdata.py +++ b/src/spatialdata/_core/spatialdata.py @@ -22,6 +22,7 @@ from zarr.errors import GroupNotFoundError from spatialdata._core._elements import Images, Labels, Points, Shapes, Tables +from spatialdata._core.transformation_manager import TransformationManager from spatialdata._core.validation import ( check_all_keys_case_insensitively_unique, check_target_region_column_symmetry, @@ -132,6 +133,7 @@ def __init__( self._shapes: Shapes = Shapes(shared_keys=self._shared_keys) self._tables: Tables = Tables(shared_keys=self._shared_keys) self.attrs = attrs if attrs else {} # type: ignore[assignment] + self.transformation_manager = TransformationManager() # Initialize the TransformationManager element_names = list(chain.from_iterable([e.keys() for e in [images, labels, points, shapes] if e is not None])) diff --git a/src/spatialdata/_core/transformation_manager/__init__.py b/src/spatialdata/_core/transformation_manager/__init__.py new file mode 100644 index 000000000..079fe5cce --- /dev/null +++ b/src/spatialdata/_core/transformation_manager/__init__.py @@ -0,0 +1,635 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import networkx as nx + +from spatialdata._core.transformation_manager.exceptions import ( + CannotRemoveCoordinateSystemError, + CoordinateSystemAlreadyExistsError, + CoordinateSystemHasElementsError, + CoordinateSystemHasTransformationsError, + CoordinateSystemNotFoundError, + ElementAlreadyExistsError, + ElementNotRegisteredToAnyCoordinateSystemError, + InvalidPathError, + TransformationNotFoundError, + TransformationPathAmbiguousError, + TransformationPathAmbiguousMultipleEdgeExpectedError, + TransformationPathAmbiguousNoEdgeExpectedError, + TransformationPathNotFoundError, + TransformationPathNotSimple, +) +from spatialdata.transformations import transformations as sd_transforms +from spatialdata.transformations.ngff.ngff_coordinate_system import NgffCoordinateSystem + +TRANSFORM_KEY = "transformation" +type EDGE_DEF = tuple[NgffCoordinateSystem, NgffCoordinateSystem, sd_transforms.BaseTransformation] + + +class TransformationManager: + def __init__(self) -> None: + """Initialize a TransformationManager with empty graph and mappings.""" + self._graph: nx.MultiDiGraph[NgffCoordinateSystem] = nx.MultiDiGraph() + # MultiDiGraph with NgffCoordinateSystem objects as nodes and transforms as edge attributes + self._element_to_cs_mapping: dict[str, NgffCoordinateSystem] = {} + # mapping element_name to the coordinate system to which the element belongs + + def assert_element_exists(self, element_name: str) -> None: + """ + Assert that an element exists in the transformation manager. + + Parameters + ---------- + element_name + The name of the element to check. + + Raises + ------ + ElementNotFoundError + If the element does not exist. + """ + if element_name not in self._element_to_cs_mapping: + raise ElementNotRegisteredToAnyCoordinateSystemError(element_name) + + def assert_element_does_not_exist(self, element_name: str) -> None: + """ + Assert than an element doesn't exist in the transformation manager. + + Parameters + ---------- + element_name + The name of the element to check. + + Raises + ------ + ElementAlreadyExistsError + If the element already exists. + """ + if element_name in self._element_to_cs_mapping: + raise ElementAlreadyExistsError(element_name) + + def assert_coordinate_system_exists(self, cs: NgffCoordinateSystem) -> None: + """ + Assert that coordinate system exists in the graph. + + Parameters + ---------- + cs + The coordinate system to check. + + Raises + ------ + CoordinateSystemNotFoundError + If the coordinate system does not exist. + """ + if cs not in self._graph: + raise CoordinateSystemNotFoundError(cs.name) + + def assert_coordinate_system_does_not_exist(self, cs: NgffCoordinateSystem) -> None: + """ + Assert that coordinate system doesn't exist in the graph. + + Parameters + ---------- + cs + The coordinate system to check. + + Raises + ------ + CoordinateSystemAlreadyExistsError + If the coordinate system already exists. + """ + if cs in self._graph: + raise CoordinateSystemAlreadyExistsError(cs.name) + + def assert_an_edge_exists_between_coordinate_systems( + self, source_cs: NgffCoordinateSystem, target_cs: NgffCoordinateSystem, edge_key: str | None = None + ) -> None: + """ + Assert that an edge exists between coordinate systems. + + Parameters + ---------- + source_cs + The input coordinate system. + target_cs + The output coordinate system. + edge_key + If specified, asserts that the specific edge exists between coordinate systems + + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + TransformationNotFoundError + If the edge does not exist. + """ + self.assert_coordinate_system_exists(source_cs) + self.assert_coordinate_system_exists(target_cs) + if not self._graph.has_edge(source_cs, target_cs): + raise TransformationNotFoundError(source_cs.name, target_cs.name, edge_key) + + def assert_coordinate_system_has_no_transformations(self, cs: NgffCoordinateSystem) -> None: + """ + Assert that coordinate system has no associated transformations. + + Parameters + ---------- + cs + The coordinate system to check. + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + CoordinateSystemHasTransformationsError + If the coordinate system has associated transformations. + + """ + self.assert_coordinate_system_exists(cs) + has_successors = next(self._graph.successors(cs), None) is not None + has_predecessors = next(self._graph.predecessors(cs), None) is not None + + if has_successors or has_predecessors: + raise CoordinateSystemHasTransformationsError(cs.name) + + def assert_coordinate_system_has_no_elements(self, cs: NgffCoordinateSystem) -> None: + """ + Assert that a coordinate system doesn't have elements belonging to it. + + Parameters + ---------- + cs + The coordinate system to check + + Raises + ------ + CoordinateSystemHasElementsError + if the coordinate system has elements belonging to it. + """ + elements = self._get_elements_belonging_to_cs(cs) + # also checks if cs exists + if elements: + raise CoordinateSystemHasElementsError(cs.name, elements) + + def add_coordinate_system(self, cs: NgffCoordinateSystem) -> None: + """ + Add a coordinate system to the transformation manager. + + Parameters + ---------- + cs + The coordinate system to add. + + Raises + ------ + CoordinateSystemAlreadyExistsError + If the coordinate system already exists. + """ + self.assert_coordinate_system_does_not_exist(cs) + self._graph.add_node(cs) + + def remove_coordinate_system(self, cs: NgffCoordinateSystem) -> None: + """ + Remove a coordinate system from the transformation manager. + + Parameters + ---------- + cs + The coordinate system to remove. + + Raises + ------ + CoordinateSystemNotFoundError + If the coordinate system is not found + CannotRemoveCoordinateSystemError + If the coordinate system cannot be removed. + Caused by + CoordinateSystemHasTransformationsError or + CoordinateSystemHasElementsError when the system has dependencies. + """ + try: + self.assert_coordinate_system_has_no_transformations(cs) + # also asserts that cs exists + self.assert_coordinate_system_has_no_elements(cs) + except (CoordinateSystemHasTransformationsError, CoordinateSystemHasElementsError) as err: + raise CannotRemoveCoordinateSystemError(cs.name) from err + + self._graph.remove_node(cs) + + def list_coordinate_systems(self) -> list[NgffCoordinateSystem]: + """ + List all registered coordinate systems. + + Returns + ------- + A list of coordinate system objects. + """ + return list(self._graph.nodes()) + + def add_element(self, element_name: str, coordinate_system: NgffCoordinateSystem) -> None: + """ + Register an element and associate it with a coordinate system. + + Parameters + ---------- + element_name + The name of the element. + coordinate_system + The coordinate system to which the element belongs. + + Raises + ------ + CoordinateSystemNotFoundError + If the coordinate system is not found. + """ + self.assert_element_does_not_exist(element_name) + self.assert_coordinate_system_exists(coordinate_system) + self._element_to_cs_mapping[element_name] = coordinate_system + + def get_element_coordinate_system(self, element_name: str) -> NgffCoordinateSystem: + """ + Get the coordinate system to which an element belongs. + + Parameters + ---------- + element_name + The name of the element. + + Returns + ------- + The coordinate system + + Raises + ------ + ElementNotFoundError + If the element does not exist. + """ + self.assert_element_exists(element_name) + return self._element_to_cs_mapping[element_name] + + def _unset_element(self, element_name: str) -> None: + """ + Unregister an element from the coordinate system to which it belongs. + + Parameters + ---------- + element_name + The name of the element. + + Raises + ------ + ElementNotFoundError + If the element has not been registered to any coordinate system. + """ + self.assert_element_exists(element_name) + del self._element_to_cs_mapping[element_name] + + @staticmethod + def _get_edge_key(edge_def: EDGE_DEF) -> str: + """ + Encode the definition of an edge into a string to be used as graph-wide edge identifier. + + Parameters + ---------- + edge_def + tuple, (source_cs, target_cs, transformation) + + Returns + ------- + edge_key + str, the encoded edge identifier + """ + source_cs, target_cs, transform = edge_def + return f"{source_cs.name}_{target_cs.name}_{id(transform)}" + + def add_transformation( + self, + source_cs: NgffCoordinateSystem, + target_cs: NgffCoordinateSystem, + transformation: sd_transforms.BaseTransformation, + ) -> None: + """ + Add a transformation between coordinate systems. + + Parameters + ---------- + source_cs + The input coordinate system. + target_cs + The output coordinate system. + transformation + The transformation to add. + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + """ + self.assert_coordinate_system_exists(source_cs) + self.assert_coordinate_system_exists(target_cs) + + edge_key = self._get_edge_key((source_cs, target_cs, transformation)) + edge_attributes = {TRANSFORM_KEY: transformation} + self._graph.add_edge(source_cs, target_cs, key=edge_key, **edge_attributes) + + def get_direct_transformations( + self, source_cs: NgffCoordinateSystem, target_cs: NgffCoordinateSystem + ) -> list[sd_transforms.BaseTransformation]: + """ + Retrieve transformations directly defined between coordinate systems. + + Parameters + ---------- + source_cs + The input coordinate system. + target_cs + The output coordinate system. + + Returns + ------- + List of transformations (empty if there are no direct transformations). + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + """ + self.assert_coordinate_system_exists(source_cs) + self.assert_coordinate_system_exists(target_cs) + + transforms = [] + if self._graph.has_edge(source_cs, target_cs): + for edge_data in self._graph[source_cs][target_cs].values(): + transform: sd_transforms.BaseTransformation = edge_data[TRANSFORM_KEY] + transforms.append(transform) + return transforms + + def remove_specific_transformation( + self, + source_cs: NgffCoordinateSystem, + target_cs: NgffCoordinateSystem, + transformation: sd_transforms.BaseTransformation, + ) -> None: + """ + Remove a specific transformation between coordinate systems. + + Parameters + ---------- + source_cs + The input coordinate system. + target_cs + The output coordinate system. + transformation + The transformation to remove. + (mainly useful for cases with multiple transformations between the same coordinate systems). + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + TransformationNotFoundError + If the transformation does not exist. + """ + expected_edge_key = self._get_edge_key((source_cs, target_cs, transformation)) + self.assert_an_edge_exists_between_coordinate_systems(source_cs, target_cs, expected_edge_key) + # also checks if source_cs and target_cs exist + self._graph.remove_edge(source_cs, target_cs, key=expected_edge_key) + + def remove_all_transformations_between_coordinate_systems( + self, + source_cs: NgffCoordinateSystem, + target_cs: NgffCoordinateSystem, + ) -> None: + """ + Remove all transformation between coordinate systems. + + Parameters + ---------- + source_cs + The input coordinate system. + target_cs + The output coordinate system. + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + TransformationNotFoundError + If no transformation exists between the coordiante systems + """ + self.assert_coordinate_system_exists(source_cs) + self.assert_coordinate_system_exists(target_cs) + if self._graph.has_edge(source_cs, target_cs): + for edge_key in list(self._graph[source_cs][target_cs].keys()): + # need to covert keys() to list to freeze it, else it will change during the following removal + self._graph.remove_edge(source_cs, target_cs, key=edge_key) + + def _get_transformation_sequences_from_simple_paths_after_disambiguation( + self, + paths: list[list[NgffCoordinateSystem]], + expected_intermediate_edges: Sequence[EDGE_DEF], + ) -> list[sd_transforms.Sequence]: + """ + Traverses paths to form sequence of Transformations. + + In case of ambiguity looks into `expected_intermediate_transformations` to disambiguate. + + Parameters + ---------- + paths: + sequence of list of nodes + expected_intermediate_edges: + list of edge definitions, for use when multiple edges are found between two coordinate systems in a path + An edge is defined as a tuple, (source_cs, target_cs, transform). + + Returns + ------- + list of transformations, each an instance of the transformation class Sequence + + Raises + ------ + TransformationPathNotSimple + if any path in `paths` is not simple, i.e., has recurring coordinate systems + """ + for path in paths: + if len(set(path)) != len(path): + raise TransformationPathNotSimple(path) + + expected_intermediate_transformation_edge_keys = set() + if expected_intermediate_edges is not None: + for edge_def in expected_intermediate_edges: + edge_key = self._get_edge_key(edge_def) + expected_intermediate_transformation_edge_keys |= {edge_key} + + all_sequences = [] + deduplicated_paths = list({repr(x): x for x in paths}.values()) + for path in deduplicated_paths: + if len(path) <= 1: + raise InvalidPathError(path) + + transformation_list = [] + for i in range(len(path) - 1): + source_cs_here, target_cs_here = path[i], path[i + 1] + edge_data = self._graph[source_cs_here][target_cs_here] + if len(edge_data) > 1: + # when there are multiple edges between a pair of coordinate systems in the path + # find the expected edge based on key + expected_intermediate_transformation_key_here = ( + expected_intermediate_transformation_edge_keys & set(edge_data.keys()) + ) + if len(expected_intermediate_transformation_key_here) == 0: + # transformation was not specified in `expected_intermediate_transformations` for disambiguation + raise TransformationPathAmbiguousNoEdgeExpectedError(source_cs_here.name, target_cs_here.name) + if len(expected_intermediate_transformation_key_here) > 1: + # multiple transformations were specified in `expected_intermediate_transformations` + # for disambiguation + raise TransformationPathAmbiguousMultipleEdgeExpectedError( + source_cs_here.name, target_cs_here.name, len(expected_intermediate_transformation_key_here) + ) + edge_key_to_use = list(expected_intermediate_transformation_key_here)[0] + transformation_list.append(edge_data[edge_key_to_use][TRANSFORM_KEY]) + else: + # Only one edge, no ambiguity + edge_key = next(iter(edge_data.keys())) + transformation_list.append(edge_data[edge_key][TRANSFORM_KEY]) + all_sequences.append(sd_transforms.Sequence(transformation_list)) + return all_sequences + + def get_all_shortest_transformation_sequences( + self, + source_cs: NgffCoordinateSystem, + target_cs: NgffCoordinateSystem, + expected_intermediate_edges: Sequence[EDGE_DEF] = (), + ) -> list[sd_transforms.Sequence]: + """ + Get all shortest sequences of transformations between two coordinate systems. + + Parameters + ---------- + source_cs + The source coordinate system. + target_cs + The target coordinate system. + expected_intermediate_edges + list of intermediate edges expected. + An edge is defined as a tuple, (source_cs, target_cs, transform). + Used to choose an edge when multiple edges are found between the same coordinate systems. + + Returns + ------- + list[Sequence] + All shortest sequences of transformations from source_cs to target_cs. + When multiple transformations are defined between the same coordinate systems, only those containing + a transformation among intermediate transformations are included. + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + TransformationPathNotFoundError + If no path exists between the source and target coordinate systems. + TransformationPathAmbiguousError + When multiple transformations are defined between the same coordinate systems and transformations + are not specified in `expeceted_intermediate_transformations` for disambiguation. + """ + try: + paths = list(nx.all_shortest_paths(self._graph, source=source_cs, target=target_cs)) + except nx.NetworkXNoPath as nxe: + raise TransformationPathNotFoundError(source_cs.name, target_cs.name) from nxe + + try: + return self._get_transformation_sequences_from_simple_paths_after_disambiguation( + paths, expected_intermediate_edges + ) + except TransformationPathAmbiguousError as tpae: + raise TransformationPathAmbiguousError(source_cs.name, target_cs.name) from tpae + + def get_all_transformation_sequences( + self, + source_cs: NgffCoordinateSystem, + target_cs: NgffCoordinateSystem, + expected_intermediate_edges: Sequence[EDGE_DEF] = (), + ) -> list[sd_transforms.Sequence]: + """ + Get all existing sequences of transformations between two coordinate systems. + + Parameters + ---------- + source_cs + The source coordinate system. + target_cs + The target coordinate system. + expected_intermediate_edges + list of intermediate edges expected. + An edge is defined as a tuple, (source_cs, target_cs, transform). + Used to choose an edge when multiple edges are found between the same coordinate systems. + + Returns + ------- + list[Sequence] + All existing sequences of transformations from source_cs to target_cs. + When multiple transformations are defined between the same coordinate systems, only those containing + a transformation among intermediate transformations are included. + + Raises + ------ + CoordinateSystemNotFoundError + If either coordinate system does not exist. + TransformationPathNotFoundError + If no path exists between the source and target coordinate systems. + TransformationPathAmbiguousError + When multiple transformations are defined between the same coordinate systems and transformations + are not specified in `expected_intermediate_transformations` for disambiguation. + """ + try: + paths = list(nx.all_simple_paths(self._graph, source=source_cs, target=target_cs)) + except nx.NetworkXNoPath as nxe: + raise TransformationPathNotFoundError(source_cs.name, target_cs.name) from nxe + + try: + return self._get_transformation_sequences_from_simple_paths_after_disambiguation( + paths, expected_intermediate_edges + ) + except TransformationPathAmbiguousError as tpae: + raise tpae from TransformationPathAmbiguousError(source_cs.name, target_cs.name) + + def _get_elements_belonging_to_cs(self, cs: NgffCoordinateSystem) -> list[str]: + """ + Get all elements belonging to a coordinate system. + + Parameters + ---------- + cs + The coordinate system to check. + + Returns + ------- + List of element names belonging to the coordinate system. + + Raises + ------ + CoordinateSystemNotFoundError + If the coordinate system does not exist. + + """ + self.assert_coordinate_system_exists(cs) + + elements = [] + for element_name, element_cs in self._element_to_cs_mapping.items(): + if element_cs == cs: + elements.append(element_name) + + return elements + + def __repr__(self) -> str: + """Return a string representation of the TransformationManager.""" + return ( + f"TransformationManager(" + f" coordinate_systems={list(self._graph.nodes())}, " + f" coordinate_transforms={[x[TRANSFORM_KEY] for *_, x in self._graph.edges(data=True)]}, " + f" elements={list(self._element_to_cs_mapping.keys())}" + f")" + ) diff --git a/src/spatialdata/_core/transformation_manager/exceptions.py b/src/spatialdata/_core/transformation_manager/exceptions.py new file mode 100644 index 000000000..9e0caff42 --- /dev/null +++ b/src/spatialdata/_core/transformation_manager/exceptions.py @@ -0,0 +1,295 @@ +# pyright: reportImplicitStringConcatenation=false + +from __future__ import annotations + +from spatialdata._types import ArrayLike +from spatialdata.transformations.graph import vert +from spatialdata.transformations.ngff.ngff_coordinate_system import NgffCoordinateSystem + + +class CoordinateSystemNotFoundError(ValueError): + """ + Exception raised when a coordinate system is not found in the transformation manager. + + Attributes + ---------- + name : str + The name of the coordinate system that was not found. + """ + + def __init__(self, name: str) -> None: + self.name = name + super().__init__(f"Coordinate system '{name}' not found in the transformation manager.") + + +class CoordinateSystemAlreadyExistsError(ValueError): + """ + Exception raised when a coordinate system already exists in the transformation manager. + + Attributes + ---------- + name : str + The name of the coordinate system to check. + """ + + def __init__(self, name: str) -> None: + self.name = name + super().__init__(f"Coordinate system '{name}' already exists in the transformation manager.") + + +class ElementNotRegisteredToAnyCoordinateSystemError(KeyError): + """ + Exception raised when an element is not found in the transformation manager. + + Attributes + ---------- + element_name : str + The name of the element that was not found. + """ + + def __init__(self, element_name: str) -> None: + self.element_name = element_name + super().__init__(f"Element '{element_name}' has not been registered to any coordinate system.") + + +class TransformationNotFoundError(KeyError): + """ + Exception raised when a transformation is not found between coordinate systems. + + Attributes + ---------- + input_cs_name : str + The name of the input coordinate system. + output_cs_name : str + The name of the output coordinate system. + edge_key: str or None + key used when adding transformation + """ + + def __init__(self, source_cs_name: str, target_cs_name: str, edge_key: str | None = None) -> None: + self.input_cs_name = source_cs_name + self.output_cs_name = target_cs_name + self.edge_key = edge_key + msg = f"Transformation from '{source_cs_name}' to '{target_cs_name}' not found" + if edge_key is not None: + msg += f" with key '{edge_key}'" + super().__init__(msg) + + +class ElementAlreadyExistsError(ValueError): + """ + Exception raised when trying to add an Element that already exists. + + Attributes + ---------- + name : str + The name of the Element that already exists. + """ + + def __init__(self, name: str) -> None: + self.name = name + super().__init__(f"Element '{name}' already exists in the transformation manager") + + +class InvalidPathError(ValueError): + """ + Exception raised when a path is defined with less than 2 nodes. + + Attributes + ---------- + invalid_path : list[NgffCoordinateSystem] + """ + + def __init__(self, invalid_path: list[NgffCoordinateSystem]) -> None: + super().__init__(f"Found an invalid path with less than 2 nodes: {invalid_path}") + + +class TransformationPathNotFoundError(ValueError): + """ + Exception raised when no transformation path exists between coordinate systems. + + Attributes + ---------- + source_cs_name : str + The name of the source coordinate system. + target_cs_name : str + The name of the target coordinate system. + """ + + def __init__(self, source_cs_name: str, target_cs_name: str) -> None: + self.source_cs_name = source_cs_name + self.target_cs_name = target_cs_name + super().__init__(f"No transformation path found from {source_cs_name} to {target_cs_name}") + + +class TransformationPathAmbiguousError(ValueError): + """ + Exception raised when multiple transformation path exists between coordinate systems. + + Attributes + ---------- + source_cs_name : str + The name of the source coordinate system. + target_cs_name : str + The name of the target coordinate system. + """ + + def __init__(self, source_cs_name: str, target_cs_name: str) -> None: + self.source_cs_name = source_cs_name + self.target_cs_name = target_cs_name + base_msg = f"Transformation Path ambiguous from {source_cs_name} to {target_cs_name}." + cause_of_confusion = self.cause_of_confusion() + msg = f"{base_msg} {cause_of_confusion}" if cause_of_confusion else base_msg + super().__init__(msg) + + def cause_of_confusion(self) -> str: + return "" + + +class TransformationPathAmbiguousNoEdgeExpectedError(TransformationPathAmbiguousError): + def __init__(self, source_cs_name: str, target_cs_name: str) -> None: + super().__init__(source_cs_name, target_cs_name) + + def cause_of_confusion(self) -> str: + + return "Multiple edges found. None of them were specified to be expected" + + +class TransformationPathAmbiguousMultipleEdgeExpectedError(TransformationPathAmbiguousError): + def __init__(self, source_cs_name: str, target_cs_name: str, number_of_edges_expected: int) -> None: + self.number_of_edges_expected = number_of_edges_expected + super().__init__(source_cs_name, target_cs_name) + + def cause_of_confusion(self) -> str: + return f"Multiple ({self.number_of_edges_expected}) edges were specified to be expected" + + +class TransformationPathNotSimple(ValueError): + """ + Exception raised when a path represented by a list of coordinate systems is not simple. + + A simple path is one in which each coordinate system appears only once + """ + + def __init__(self, path: list[NgffCoordinateSystem]) -> None: + self.path = path + css_formatted_one_per_line = "\n".join(repr(cs) for cs in path) + super().__init__( + f"Transformation Path not simple, i.e., some coordinate systems appear multiple times:\n" + f"{css_formatted_one_per_line}" + ) + + +class CannotRemoveCoordinateSystemError(ValueError): + """Exception raised when trying to remove a coordinate system. + + Attributes + ---------- + name: str + name of the coordinate system + """ + + def __init__(self, name: str) -> None: + self.name = name + super().__init__(f"Cannot remove coordinate system with name {name}.") + + +class CoordinateSystemHasTransformationsError(ValueError): + """ + Exception raised when trying to remove a coordinate system that has associated transformations. + + Attributes + ---------- + name : str + The name of the coordinate system that has transformations. + """ + + def __init__(self, name: str) -> None: + self.name = name + super().__init__(f"Coordinate System ('{name}') has transformations.") + + +class CoordinateSystemHasElementsError(ValueError): + """ + Exception raised when trying to remove a coordinate system that has associated elements. + + Attributes + ---------- + name : str + The name of the coordinate system that has elements. + associated_elements : list[str] + List of element names associated with the coordinate system. + """ + + def __init__(self, name: str, associated_elements: list[str]) -> None: + self.name = name + self.associated_elements = associated_elements + super().__init__(f"Coordinate system '{name}' has elements belonging to it: {associated_elements}") + + +class TransformationManagerWarning(UserWarning): + """Base warning category for TransformationManager.""" + + pass + + +class IncompatibleCoordSystemsError(Exception): + def __init__(self, input: vert.CoordSystem, output: vert.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: vert.Axis, cs: vert.CoordSystem) -> None: + self.axis = axis + self.cs = cs + super().__init__(f"Axis {axis.name} is not in coordinate system {cs}") + + +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__(f"Matrix does not have det(M) == 1:\n{matrix}") + + +class NotOrthonormalError(Exception): + def __init__(self, matrix: ArrayLike) -> None: + self.matrix = matrix + super().__init__(f"Matrix is not orthonormal:\n{matrix}") + + +class EmptyTransformSequenceError(Exception): + def __init__(self) -> None: + super().__init__("Empty sequence of transformations") + + +class AxisRedefinitionError(Exception): + def __init__(self, axis: vert.Axis) -> None: + super().__init__(f"Axis {axis.name} is defined multiple times") + + +class UnmappedAxisError(Exception): + def __init__(self, axis: vert.Axis, cs: vert.CoordSystem) -> None: + self.axis = axis + self.cs = cs + super().__init__(f"Axis {axis.name} from coordinate system {cs} is not mapped to anything") diff --git a/src/spatialdata/_io/io_raster.py b/src/spatialdata/_io/io_raster.py index 3806fabf7..61371829b 100644 --- a/src/spatialdata/_io/io_raster.py +++ b/src/spatialdata/_io/io_raster.py @@ -1,12 +1,19 @@ +# pyright: standard + from __future__ import annotations +from dataclasses import dataclass import warnings from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any, Literal, TypeGuard +from pathlib import Path, PurePosixPath +from typing import Any, Final, Iterator, Literal, TypeGuard, assert_type, get_args +import itertools import dask.array as da import numpy as np +import ome_zarr as oz +import ome_zarr_models.v06.coordinate_transforms as ozm06ct +import xarray as xr import zarr from ome_zarr.format import Format from ome_zarr.io import ZarrLocation @@ -18,6 +25,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,6 +38,7 @@ RasterFormatType, get_ome_zarr_format, ) +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 @@ -39,6 +48,8 @@ _set_transformations, compute_coordinates, ) +from spatialdata.transformations.graph.edge import BaseTransformationEdge, IdentityEdge, ScaleEdge, TranslationEdge, parse_ngff_transf +from spatialdata.transformations.graph.vert import Axis, CoordSystem, CsVert, Edge, Edges, NgffCsPath, Provenance, Spaces def _is_flat_int_sequence(value: object) -> TypeGuard[Sequence[int]]: @@ -161,11 +172,544 @@ def _prepare_storage_options( return prepared_options +@dataclass(frozen=True, eq=True) +class ElementKey: + namespace: Final[PurePosixPath] + name: Final[str] + + def __post_init__(self): + assert not self.namespace.is_absolute() + + @classmethod + def from_path(cls, path: PurePosixPath) -> ElementKey: + return ElementKey(namespace=path.parent, name=path.name) + + def to_path(self) -> PurePosixPath: + return self.namespace / self.name + + def is_sibling_of(self, other: ElementKey) -> bool: + return self.namespace == other.namespace + + def is_relative_to(self, other: ElementKey) -> bool: + return self.namespace.is_relative_to(other.namespace) + + +@dataclass(frozen=True, eq=True) +class MultiscalesKeys: + raster: ElementKey + labels: Mapping[PurePosixPath, ElementKey] + + +class FutureSpatialData: + _graph: Spaces # FIXME: use TransformationManager + _images: dict[ElementKey, xr.DataTree] + _labels: dict[ElementKey, xr.DataTree] + _intrinsics: dict[ElementKey, CsVert] + + def __init__(self, graph: Spaces | None = None) -> None: + self._graph = graph or Spaces() + self._images = {} + self._labels = {} + self._intrinsics = {} + super().__init__() + + def incorporate(self, other: FutureSpatialData): + # FIXME: there might be more conflicts that need to be checked before merging + # so we avoid a half-merged SpatialData + for key in other._images.keys(): + if key in self._images: + raise KeyError(f"Can't merge key {key}") # FIXME: custm error + + for key in other._labels.keys(): + if key in self._images: + raise KeyError(f"Can't merge key {key}") # FIXME: custm error + + self._graph.incorporate(other._graph) + self._images.update(other._images) + self._labels.update(other._labels) + self._intrinsics.update(other._intrinsics) + + def incorporate_graph(self, graph: Spaces): + self._graph.incorporate(graph) + + def incorporate_ngff_graph( + self, + group_path: PurePosixPath, + coord_systems: Sequence[ozm06ct.CoordinateSystem], + transforms: Sequence[ozm06ct.AnyTransform], + ) -> None: + self._graph.incorporate_ngff_graph(group_path, coord_systems, transforms) + + def add_image(self, key: ElementKey, image: xr.DataTree, into: CsVert): + if into not in self._graph.verts(): + raise KeyError(f"{CsVert} is not in the transformation graph") + if key in self._labels: + raise KeyError(f"Already a label") + self._images[key] = image + self._intrinsics[key] = into + + def get_image(self, key: ElementKey) -> xr.DataTree | None: + return self._images[key] + + def add_label(self, key: ElementKey, label: xr.DataTree, into: CsVert): + if into not in self._graph.verts(): + raise KeyError(f"{CsVert} is not in the transformation graph") + if key in self._images: + raise KeyError(f"Already an image") + self._labels[key] = label + self._intrinsics[key] = into + + def get_label(self, key: ElementKey) -> xr.DataTree | None: + return self._labels[key] + + def get_intrinsic(self, key: ElementKey) -> CsVert: + return self._intrinsics[key] + + def add_cs(self, cs: CoordSystem, provenance: Provenance): + self._graph.add_cs(cs=cs, provenance=provenance) + + def find_node(self, *, ngff_path: NgffCsPath) -> CsVert: + return self._graph.find_node(ngff_path=ngff_path) + + def add_transformation(self, *, input: CsVert, output: CsVert, transf: BaseTransformationEdge, name: str | None): + self._graph.add_transformation(input=input, output=output, transf=transf, name=name) + + def edges(self) -> Edges: + return self._graph.edges() + + def ngff_label_dumper(self, *, label_key: ElementKey) -> MultiscalesDumper: + return self._ngff_raster_dumper(MultiscalesKeys(raster=label_key, labels={})) + + def ngff_multiscales_dumper(self, keys: MultiscalesKeys) -> MultiscalesDumper: + return self._ngff_raster_dumper(keys) + + def _ngff_raster_dumper(self, keys: MultiscalesKeys) -> MultiscalesDumper: + assert keys.raster in self._images or (keys.raster in self._labels and len(keys.labels) == 0) + + intrinsic = self.get_intrinsic(key=keys.raster) + multiscales_transforms: set[Edge] = self.edges().containing(cs=intrinsic).ngff_siblings().collect(set) + + label_dumpers: dict[PurePosixPath, MultiscalesDumper] = {} + for label_rel_path, label_key in keys.labels.items(): + label_intrinsic = self.get_intrinsic(label_key) + ngff_transforms_to_label = ( + self.edges() + .directed(fro=intrinsic, to=label_intrinsic) #mutlsicale-local transfs can point down to labels + .transf_is((IdentityEdge, ScaleEdge, TranslationEdge)) + .collect(set) + ) + multiscales_transforms |= ngff_transforms_to_label + label_dumpers[label_rel_path] = self.ngff_label_dumper(label_key=label_key) + + return MultiscalesDumper( + intrinsic=intrinsic, + key=keys.raster, + labels=label_dumpers, + sdata=self, + transform_edges=list(multiscales_transforms), + ) + + def candidate_ngff_labels( + self, + *, + image_key: ElementKey, + gather_by_ngff_provenance: bool, + gather_by_unique_connection: bool, + ) -> Mapping[PurePosixPath, ElementKey]: + image_intrinsic = self.get_intrinsic(image_key) + out: dict[PurePosixPath, ElementKey] = {} + + for label_key in self._labels.keys(): + label_intrinsic = self.get_intrinsic(label_key) + edges = ( + self.edges() + .directed(fro=image_intrinsic, to=label_intrinsic) + .transf_is((IdentityEdge, ScaleEdge, TranslationEdge)) + .collect(set) + ) + + proceed = False + if not proceed: + proceed = gather_by_unique_connection and len(edges) == 1 + if not proceed: + proceed = gather_by_ngff_provenance and any(e.points_down_ngff_hierarchy() for e in edges) + if not proceed: + continue + + if label_key.is_relative_to(image_key): + label_rel_path = label_key.namespace.relative_to(image_key.namespace) + else: + label_rel_path = PurePosixPath(label_key.name) + suffix = 2 + while label_rel_path in out: + label_rel_path = PurePosixPath(f"{label_key.name}__{suffix}") + suffix += 1 + out[label_rel_path] = label_key + + return out + + def scene_dumper( + self, + children: Mapping[PurePosixPath, MultiscalesKeys], + ) -> SceneDumper: + scene_coord_systems: set[CsVert] = set(self._graph.verts()) + + ms_dumpers = { + ms_path: self.ngff_multiscales_dumper(mskey) + for ms_path, mskey in children.items() + } + ms_coord_systems = set(itertools.chain.from_iterable( + ms.all_coord_systems() for ms in ms_dumpers.values() + )) + scene_coord_systems: set[CsVert] = set() + scene_edges: set[Edge] = set() + for edge in self.edges(): + if edge.origin in ms_coord_systems and edge.target not in ms_coord_systems: + scene_cs = edge.target + elif edge.origin not in ms_coord_systems and edge.target in ms_coord_systems: + scene_cs = edge.origin + else: + continue + scene_coord_systems.add(scene_cs) + scene_edges.add(edge) + + return SceneDumper( + multiscales=ms_dumpers, + coord_systems=tuple(scene_coord_systems), + transform_edges=tuple(scene_edges), + ) + + def default_scene_dumper(self) -> SceneDumper: + multiscale_keys: dict[PurePosixPath, MultiscalesKeys] = {} + for img_key, img in self._images.items(): + candidate_labels = self.candidate_ngff_labels( + image_key=img_key, gather_by_ngff_provenance=True, gather_by_unique_connection=False, + ) + # FIXME: raising for duplicate names doens't seem appropriate + img_path = img_key.namespace + if img_path in multiscale_keys: + raise ValueError(f"Multiple images with conflicting paths: {img_path}") # FIXME: custom exception + multiscale_keys[img_path] = MultiscalesKeys( + raster=img_key, + labels=candidate_labels + ) + return self.scene_dumper(multiscale_keys) + +def try_parse_ngff06_scene( + group_path: PurePosixPath, + scene: oz.OMEZarrScene, +) -> FutureSpatialData: + sdata = FutureSpatialData() + + for img_path, img in scene.images.items(): + ms_sdata = try_parse_ngff06_multiscale( + group_path=group_path / img_path.lstrip("/"), + multiscale=img, + ) + sdata.incorporate(ms_sdata) + + sdata.incorporate_ngff_graph( + group_path=group_path, + coord_systems=scene.coordinate_systems, + transforms=scene.coordinate_transformations, + ) + + return sdata + +def try_read_ngff06_multiscale( + store: Path | zarr.Group, +) -> FutureSpatialData: + multiscale = oz.OMEZarrMultiscale.from_ome_zarr(str(store) if isinstance(store, Path) else store) + assert isinstance(multiscale, oz.OMEZarrMultiscale) # disambiguate from OMEZarrLabel + # FIXME: we should probably use something liek a zarr-interpretable URL isntead of a path + return try_parse_ngff06_multiscale(multiscale=multiscale, group_path=PurePosixPath(".")) + +def try_parse_ngff06_multiscale( + group_path: PurePosixPath, + multiscale: oz.OMEZarrMultiscale | oz.OMEZarrLabels, +) -> FutureSpatialData: + sdata = FutureSpatialData() + + if isinstance(multiscale, oz.OMEZarrMultiscale): + for label_path, label in (multiscale.labels or {}).items(): + # spec says labels are under the "labels/" directory + label_namespace = group_path / "labels" / label_path # FIXME: is label_path safe? + label_sdata = try_parse_ngff06_multiscale( + group_path=label_namespace, + multiscale=label, + ) + sdata.incorporate(label_sdata) + + # mutliscales transforms can reference CSs in `labels/` , so we process those last + sdata.incorporate_ngff_graph( + group_path=group_path, + coord_systems = multiscale.metadata.coordinateSystems, + transforms = multiscale.metadata.coordinateTransformations or (), + ) + + channel_names: Sequence[str] | None = None + if isinstance(multiscale, oz.OMEZarrMultiscale): + omero = multiscale.omero + channel_names = None if omero is None else [d.color for d in omero.channels] + elif isinstance(multiscale, oz.OMEZarrLabels): + # FIXME: ome zarr has special annotations for labels + pass + else: + assert_never(multiscale) + + intrinsic_cs_path = NgffCsPath(group_path=group_path, name=multiscale.metadata.intrinsic_coordinate_system.name) + intrinsic = sdata.find_node(ngff_path=intrinsic_cs_path) + intrinsic_cs = intrinsic.coord_system + data_tree = xr.DataTree() + for scale_idx, (ds_md, ds) in enumerate(zip(multiscale.metadata.datasets, multiscale.images, strict=True)): + # 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.create_dummy(axes=intrinsic_cs.num_axes) + pixel_cs_to_intrinsic_ngff = ozm06ct.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 " 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, + ) + }, + ) + + sdata_key = ElementKey(namespace=group_path, name=multiscale.name) + if isinstance(multiscale, oz.OMEZarrLabels): + sdata.add_label(key=sdata_key, label=data_tree, into=intrinsic) + elif isinstance(multiscale, oz.OMEZarrMultiscale): + sdata.add_image(key=sdata_key, image=data_tree, into=intrinsic) + else: + assert_never(multiscale) + + return sdata + +class MultiscalesDumper: + name: Final[str] + raster_data: Final[xr.DataTree] + intrinsic: Final[CsVert] + coord_systems: Final[set[CsVert]] + transform_edges: Final[Sequence[Edge]] + labels: Final[Mapping[PurePosixPath, MultiscalesDumper]] + cs_to_path: Final[Mapping[CsVert, PurePosixPath]] + + def __init__( + self, + sdata: FutureSpatialData, + key: ElementKey, + intrinsic: CsVert, + transform_edges: Sequence[Edge], # FIXME: set[Edge]? + labels: Mapping[PurePosixPath, MultiscalesDumper], + ): + coord_systems = {intrinsic} + + _labels_coord_systems = set(itertools.chain(label.coord_systems for label in labels.values())) + for e in transform_edges: + for cs in (e.origin, e.target): + if cs not in _labels_coord_systems: + coord_systems.add(cs) + + cs_to_path: dict[CsVert, PurePosixPath] = {cs: PurePosixPath("") for cs in coord_systems} + intrinsic = sdata.get_intrinsic(key) + cs_to_path[intrinsic] = PurePosixPath("") + for label_path, label_dumper in labels.items(): + if label_path.is_absolute(): + raise ValueError(f"Label path {label_path} cannot be absolute") + for label_cs, label_cs_path in label_dumper.cs_to_path.items(): + label_cs_path = "labels" / label_path / label_cs_path + if label_cs in cs_to_path: + raise ValueError(f"Coordinate system {label_cs} from {label_cs_path} is not unique") # FIXME + cs_to_path[label_cs] = label_cs_path + + for edge in transform_edges: + if edge.origin!= intrinsic and edge.target != intrinsic: + raise ValueError("Origin or target must be the intrinsic") # FIXME + for vert in (edge.origin, edge.target): + if vert not in cs_to_path: + raise ValueError(f"Reference to {vert} is dangling") # FIXME + + self.name = key.name + raster_data = sdata.get_image(key) or sdata.get_label(key) + if raster_data is None: + raise KeyError(f"Could not find element with key {key}") # FIXME + self.raster_data = raster_data + self.intrinsic = intrinsic + self.coord_systems = coord_systems + self.transform_edges = transform_edges + self.labels = labels + self.cs_to_path = cs_to_path + + super().__init__() + + def contains(self, cs: CsVert) -> bool: + if cs in self.coord_systems: + return True + for label in self.labels.values(): + if label.contains(cs): + return True + return False + + def all_coord_systems(self) -> Iterator[CsVert]: + yield from self.coord_systems + for label in self.labels.values(): + yield from label.all_coord_systems() + + def _make_ngff_cs_ident(self, csvert: CsVert) -> ozm06ct.CoordinateSystemIdentifier: + path = self.cs_to_path[csvert] + return ozm06ct.CoordinateSystemIdentifier( + name=csvert.name, + path=None if path == PurePosixPath(".") else path.as_posix(), + ) + + def dump(self) -> oz.OMEZarrMultiscale: + image_data_tree: xr.DataTree = self.raster_data + scale0 = image_data_tree["scale0"]["image"] # FIXME: using current silent convention + assert not isinstance(scale0, xr.DataTree) + assert isinstance(scale0.data, (da.Array, np.ndarray)) + + omz_image = oz.OMEZarrImage( + data=scale0.data, + axes=[str(d) for d in scale0.dims], # FIXME: safer alternative? + axes_units={ + ax.name: ax.unit + for ax in self.intrinsic.coord_system.axes + if ax.unit is not None + }, + axes_types={ax.name: ax.type for ax in self.intrinsic.coord_system.axes}, + name=self.name, + ) + + omz_transforms: list[ozm06ct.AnyTransform] = [ + edge.data["transform"].to_ngff_transform( + name=edge.data["name"], + input=self._make_ngff_cs_ident(edge.origin), + output=self._make_ngff_cs_ident(edge.target), + ) + for edge in self.transform_edges + ] + + dumped_labels: list[oz.OMEZarrLabels] = [] + for label in self.labels.values(): + dumped_label = label.dump() + dumped_labels.append( + oz.OMEZarrLabels(image=dumped_label.images[0]) # FIXME: use from_pyramid, no cs or transfers?? + ) + + all_scale_factors: list[dict[str, int]] = [] + for scale in self.raster_data.values(): # FIXME: is order guaranteed? do we need it? + img = scale["image"] # FIXME: using old convention + assert not isinstance(img, xr.DataTree) + scale_factors: dict[str, int] = {} + for ax_name in img.dims: + assert isinstance(ax_name, str) + scale = scale0.sizes[ax_name] / img.sizes[ax_name] + if not scale.is_integer(): + raise ValueError(f"scale factors must be ints for now; found {scale}") + scale_factors[ax_name] = int(scale) + all_scale_factors.append(scale_factors) + + # FIXME: use from_pyramid once it's available + return oz.OMEZarrMultiscale( + scale_factors=all_scale_factors[1:], # scale0 is implicit + image=omz_image, + labels=dumped_labels, + default_coordinate_system_name=self.intrinsic.provenance.name, + coordinate_systems=[cs.to_ngff_cs() for cs in self.coord_systems.difference({self.intrinsic})], + coordinate_transformations=tuple(omz_transforms), + ) + +class SceneDumper: + multiscales: Final[ Mapping[PurePosixPath, MultiscalesDumper] ] + coord_systems: Final[ Sequence[CsVert] ] + transform_edges: Final[ Sequence[Edge] ] + cs_to_path: Final[ Mapping[CsVert, PurePosixPath] ] + + def __init__( + self, + *, + multiscales: Mapping[PurePosixPath, MultiscalesDumper], + coord_systems: Sequence[CsVert], + transform_edges: Sequence[Edge], + ) -> None: + cs_to_path: dict[CsVert, PurePosixPath] = {cs: PurePosixPath("") for cs in coord_systems} + for ms_path, ms_dumper in multiscales.items(): + if ms_path.is_absolute(): + raise ValueError(f"Raster name cannot be absolute: {ms_path}") + for raster_cs, raster_cs_path in ms_dumper.cs_to_path.items(): + raster_cs_path = ms_path / raster_cs_path + if raster_cs in cs_to_path: + raise ValueError(f"Coordinate system {raster_cs} from {raster_cs_path} is not unique") + cs_to_path[raster_cs] = raster_cs_path + + for edge in transform_edges: + for vert in (edge.origin, edge.target): + if vert not in cs_to_path: + raise ValueError(f"Reference to {vert} is dangling") # FIXME + + self.multiscales = multiscales + self.coord_systems = coord_systems + self.transform_edges = transform_edges + self.cs_to_path = cs_to_path + super().__init__() + + def _make_ngff_cs_ident(self, csvert: CsVert) -> ozm06ct.CoordinateSystemIdentifier: + path = self.cs_to_path[csvert] + return ozm06ct.CoordinateSystemIdentifier( + name=csvert.name, + path=None if path == PurePosixPath(".") else path.as_posix(), + ) + + def dump(self) -> oz.OMEZarrScene: + omz_transforms: list[ozm06ct.AnyTransform] = [ + edge.data["transform"].to_ngff_transform( + name=edge.data["name"], + input=self._make_ngff_cs_ident(edge.origin), + output=self._make_ngff_cs_ident(edge.target), + ) + for edge in self.transform_edges + ] + + return oz.OMEZarrScene( + images={path.as_posix(): r.dump() for path, r in self.multiscales.items()}, + coordinate_systems=[cs.to_ngff_cs() for cs in self.coord_systems], + coordinate_transformations=omz_transforms, + ) + def _read_multiscale( - store: str | Path, raster_type: Literal["image", "labels"], reader_format: Format + store: str | Path, raster_type: ELEMENT_TYPE_RASTER, reader_format: Format ) -> DataArray | DataTree: assert isinstance(store, str | Path) - assert raster_type in ["image", "labels"] + assert raster_type in get_args(ELEMENT_TYPE_RASTER), ValueError( + f"{raster_type} is not a valid raster type. Must be one of {get_args(ELEMENT_TYPE_RASTER)}." + ) nodes: list[Node] = [] image_loc = ZarrLocation(store, fmt=reader_format) @@ -214,7 +758,7 @@ def _read_multiscale( transformations = _get_transformations_from_ngff_dict(encoded_ngff_transformations) # if image, read channels metadata channels: list[Any] | None = None - if raster_type == "image": + if raster_type == ELEMENT_TYPE.IMAGE: if legacy_channels_metadata is not None: channels = [d["label"] for d in legacy_channels_metadata["channels"]] if omero_metadata is not None: @@ -229,7 +773,7 @@ def _read_multiscale( data = loaded_node.array(resolution=datasets[0]) si = DataArray( data, - name="image", + name=ELEMENT_TYPE.IMAGE, dims=axes, coords={"c": channels} if channels is not None else {}, ) @@ -264,9 +808,7 @@ def _get_multiscale_nodes(image_nodes: list[Node], nodes: list[Node]) -> list[No return nodes -def _get_raster_element_group( - raster_type: Literal["image", "labels"], 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 @@ -285,7 +827,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): @@ -299,7 +841,7 @@ def _get_raster_element_group( def _write_raster( - raster_type: Literal["image", "labels"], + raster_type: ELEMENT_TYPE_RASTER, raster_data: DataArray | DataTree, group: zarr.Group, name: str, @@ -332,12 +874,13 @@ def _write_raster( metadata Additional metadata for the raster element """ - if raster_type not in ["image", "labels"]: - raise ValueError(f"{raster_type} is not a valid raster type. Must be 'image' or 'labels'.") + assert raster_type in get_args(ELEMENT_TYPE_RASTER), ValueError( + f"{raster_type} is not a valid raster type. Must be one of {get_args(ELEMENT_TYPE_RASTER)}." + ) # "name" and "label_metadata" are only used for labels. "name" is written in write_multiscale_ngff() but ignored in # write_image_ngff() (possibly an ome-zarr-py bug). We only use "name" to ensure correct group access in the # ome-zarr API. - if raster_type == "labels": + if raster_type == ELEMENT_TYPE.LABELS: metadata["name"] = name metadata["label_metadata"] = label_metadata @@ -367,7 +910,7 @@ def _write_raster( raise ValueError("Not a valid labels object") group = _get_raster_element_group(raster_type, group, name) - if raster_type == "image": + if raster_type == ELEMENT_TYPE.IMAGE: # ome-zarr-py >= 0.18 no longer writes the omero channel metadata, so we write it ourselves. overwrite_channel_names(group, raster_data) if ATTRS_KEY not in group.attrs: @@ -463,7 +1006,7 @@ def _update_dict_v3(d: dict[str, Any]) -> None: def _write_raster_dataarray( - raster_type: Literal["image", "labels"], + raster_type: ELEMENT_TYPE_RASTER, group: zarr.Group, element_name: str, raster_data: DataArray, @@ -493,7 +1036,7 @@ def _write_raster_dataarray( metadata Additional metadata for the raster element """ - write_single_scale_ngff = write_image_ngff if raster_type == "image" else write_labels_ngff + write_single_scale_ngff = write_image_ngff if raster_type == ELEMENT_TYPE.IMAGE else write_labels_ngff data = raster_data.data transformations = _get_transformations(raster_data) @@ -532,7 +1075,7 @@ def _write_raster_dataarray( def _write_raster_datatree( - raster_type: Literal["image", "labels"], + raster_type: ELEMENT_TYPE_RASTER, group: zarr.Group, element_name: str, raster_data: DataTree, @@ -562,7 +1105,9 @@ def _write_raster_datatree( metadata Additional metadata for the raster element """ - write_multi_scale_ngff = write_multiscale_ngff if raster_type == "image" else write_multiscale_labels_ngff + write_multi_scale_ngff = ( + write_multiscale_ngff if raster_type == ELEMENT_TYPE.IMAGE else write_multiscale_labels_ngff + ) data = get_pyramid_levels(raster_data, attr="data") list_of_input_axes: list[Any] = get_pyramid_levels(raster_data, attr="dims") assert len(set(list_of_input_axes)) == 1 @@ -632,7 +1177,7 @@ def write_image( ) _write_raster( - raster_type="image", + raster_type=ELEMENT_TYPE.IMAGE, raster_data=image, group=group, name=name, @@ -659,7 +1204,7 @@ def write_labels( ) _write_raster( - raster_type="labels", + raster_type=ELEMENT_TYPE.LABELS, raster_data=labels, group=group, name=name, diff --git a/src/spatialdata/_io/io_zarr.py b/src/spatialdata/_io/io_zarr.py index 186d7c974..cee3aba7d 100644 --- a/src/spatialdata/_io/io_zarr.py +++ b/src/spatialdata/_io/io_zarr.py @@ -5,7 +5,7 @@ from collections.abc import Callable from json import JSONDecodeError from pathlib import Path -from typing import Any, Literal, cast +from typing import Any, Literal, cast, get_args import zarr.storage from anndata import AnnData @@ -28,7 +28,14 @@ from spatialdata._io.io_shapes import _read_shapes from spatialdata._io.io_table import _read_table from spatialdata._logging import logger -from spatialdata._types import JSONValue, Raster_T +from spatialdata._types import ( + ELEMENT_TYPE, + ELEMENT_TYPE_RASTER, + ELEMENT_TYPE_VECTOR, + GROUP_NAME, + JSONValue, + Raster_T, +) def _read_zarr_group_spatialdata_element( @@ -37,8 +44,8 @@ def _read_zarr_group_spatialdata_element( sdata_version: Literal["0.1", "0.2"], selector: set[str], read_func: Callable[..., Any], - group_name: Literal["images", "labels", "shapes", "points", "tables"], - element_type: Literal["image", "labels", "shapes", "points", "tables"], + group_name: GROUP_NAME, + element_type: ELEMENT_TYPE, element_container: (dict[str, Raster_T] | dict[str, DaskDataFrame] | dict[str, GeoDataFrame] | dict[str, AnnData]), on_bad_files: Literal[BadFileHandleMethod.ERROR, BadFileHandleMethod.WARN], ) -> None: @@ -76,14 +83,14 @@ def _read_zarr_group_spatialdata_element( ValueError, ), ): - if element_type in ["image", "labels"]: + if element_type in get_args(ELEMENT_TYPE_RASTER): reader_format = get_raster_format_for_read(elem_group, sdata_version) element = read_func( elem_group_path, element_type, reader_format, ) - elif element_type in ["shapes", "points", "tables"]: + elif element_type in get_args(ELEMENT_TYPE_VECTOR) + (ELEMENT_TYPE.TABLES,): element = read_func(elem_group_path) else: raise ValueError(f"Unknown element type {element_type}") @@ -204,19 +211,19 @@ def read_zarr( # we could make this more readable. One can get lost when looking at this dict and iteration over the items group_readers: dict[ - Literal["images", "labels", "shapes", "points", "tables"], + GROUP_NAME, tuple[ Callable[..., Any], - Literal["image", "labels", "shapes", "points", "tables"], + ELEMENT_TYPE, dict[str, Raster_T] | dict[str, DaskDataFrame] | dict[str, GeoDataFrame] | dict[str, AnnData], ], ] = { - # ome-zarr-py needs a kwargs that has "image" has key. So here we have "image" and not "images" - "images": (_read_multiscale, "image", images), - "labels": (_read_multiscale, "labels", labels), - "points": (_read_points, "points", points), - "shapes": (_read_shapes, "shapes", shapes), - "tables": (_read_table, "tables", tables), + # ome-zarr-py needs a kwargs that has "image" as key. So here we have "image" and not "images" + GROUP_NAME.IMAGES: (_read_multiscale, ELEMENT_TYPE.IMAGE, images), + GROUP_NAME.LABELS: (_read_multiscale, ELEMENT_TYPE.LABELS, labels), + GROUP_NAME.POINTS: (_read_points, ELEMENT_TYPE.POINTS, points), + GROUP_NAME.SHAPES: (_read_shapes, ELEMENT_TYPE.SHAPES, shapes), + GROUP_NAME.TABLES: (_read_table, ELEMENT_TYPE.TABLES, tables), } for group_name, ( read_func, @@ -295,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", diff --git a/src/spatialdata/_types.py b/src/spatialdata/_types.py index e41de9fe7..46c3f19c5 100644 --- a/src/spatialdata/_types.py +++ b/src/spatialdata/_types.py @@ -1,11 +1,22 @@ from __future__ import annotations -from typing import Any +from enum import StrEnum +from typing import Any, Literal import numpy as np from xarray import DataArray, DataTree -__all__ = ["ArrayLike", "ColorLike", "DTypeLike", "JSONValue", "Raster_T"] +__all__ = [ + "ELEMENT_TYPE", + "ELEMENT_TYPE_RASTER", + "ELEMENT_TYPE_VECTOR", + "GROUP_NAME", + "ArrayLike", + "ColorLike", + "DTypeLike", + "JSONValue", + "Raster_T", +] from numpy.typing import DTypeLike, NDArray @@ -21,8 +32,29 @@ 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 # a non-string key as a string. type JSONValue = dict[str, JSONValue] | list[JSONValue] | str | int | float | bool | None + + +class ELEMENT_TYPE(StrEnum): + IMAGE = "image" + LABELS = "labels" + SHAPES = "shapes" + POINTS = "points" + TABLES = "tables" + + +class GROUP_NAME(StrEnum): + IMAGES = "images" + LABELS = "labels" + SHAPES = "shapes" + POINTS = "points" + TABLES = "tables" + + +ELEMENT_TYPE_RASTER = Literal[ELEMENT_TYPE.IMAGE, ELEMENT_TYPE.LABELS] +ELEMENT_TYPE_VECTOR = Literal[ELEMENT_TYPE.POINTS, ELEMENT_TYPE.SHAPES] diff --git a/src/spatialdata/_utils.py b/src/spatialdata/_utils.py index a0086b7b2..f5ec67865 100644 --- a/src/spatialdata/_utils.py +++ b/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") 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..84eaa25cc --- /dev/null +++ b/src/spatialdata/transformations/graph/edge.py @@ -0,0 +1,1217 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from typing import Final +from warnings import warn + +import numpy as np +import ome_zarr_models.v06.coordinate_transforms as ozm06ct + +import spatialdata._core.transformation_manager.exceptions as ex +from spatialdata._types import ArrayLike +import spatialdata.transformations.graph.vert as vert + + +class BaseTransformationEdge(ABC): + """Base class for all the transformations defined by the NGFF specification.""" + + input: Final[vert.CoordSystem] + output: Final[vert.CoordSystem] + + def __init__( + self, + *, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + self.input = input + self.output = output + 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) -> BaseTransformationEdge | None: + """Return the inverse of the transformation if it exists""" + + @abstractmethod + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + + @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: ArrayLike) -> None: + """ + Validate if the shape of the points (coordinates to be transformed) are consistent with the input size of the + transformation. + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + input_size = len(self.input.axes) + if len(points.shape) != 2 or points.shape[1] != input_size: + raise ex.UnexpectedShapeError( + array_name="points", + expected_shape=f"(, {self.input.num_axes})", + array_shape=points.shape, + ) + + # order of the composition: self is applied first, then the transformation passed as argument + def compose_with(self, transformation: BaseTransformationEdge) -> BaseTransformationEdge: + """ + 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. + + Raises + ------ + IncompatibleCoordSystemsError + if this transformation's output coordinate system doesn't match `transformation`'s input + """ + return SequenceEdge(transformations=[self, transformation]) + + @abstractmethod + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.AnyTransform: + pass + + +class AffineEdge(BaseTransformationEdge): + """The Affine transformation from the NGFF specification.""" + + linear: Final[ArrayLike] + translation: Final[ArrayLike] + affine: Final[ArrayLike] + + def __init__( + self, + *, + linear: ArrayLike, + translation: ArrayLike | None = None, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + """ + Parameters + ---------- + 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. + + Raises + ------ + UnexpectedShapeError + if `linear`'s shape isn't (output.num_axes, input.num_axes) or + if `translation`'s shape isn't (output.num_axes,) + """ + 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 ex.UnexpectedShapeError( + array_name="linear", array_shape=linear.shape, expected_shape=expected_linear_shape + ) + expected_translation_shape = (num_outputs,) + if translation.shape != expected_translation_shape: + raise ex.UnexpectedShapeError( + array_name="translation", array_shape=translation.shape, expected_shape=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) + + 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, + *, + affine_matrix: ArrayLike, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> AffineEdge: + """Creates an AffineEdge from a raw affine matrix in homogenous coordinates + + Parameters + ---------- + 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. + + Raises + ------ + UnexpectedShapeError + if `affine_matrix`'s shape isn't (output.num_axes + 1, input.num_axes + 1) + """ + return AffineEdge( + linear=affine_matrix[:-1, :-1], + translation=affine_matrix[:-1, -1], + input=input, + output=output, + ) + + def inverse(self) -> BaseTransformationEdge | None: + cond = np.linalg.cond(self.affine) + if cond > 1e10: + warn( + f"Inverting an affine matrix with large condition number ({cond:.2e})", + RuntimeWarning, + stacklevel=2, + ) + 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], + input=self.output, + output=self.input, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + self._validate_transform_points_shapes(points) + return (points @ self.linear.T) + self.translation + + def to_affine(self) -> AffineEdge: + return AffineEdge(input=self.input, output=self.output, linear=self.linear, translation=self.translation) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.Affine: + return ozm06ct.Affine( + name=name, + affine=tuple(tuple(row) for row in self.affine[:-1, :]), + input=input, + output=output, + ) + + +class IdentityEdge(BaseTransformationEdge): + """The Identity transformation from the NGFF specification.""" + + def __init__( + self, + *, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + """ + Parameters + ---------- + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + + Raises + ------ + IncompatibleCoordSystemsError + if `input` and `output` don't have the same number of axes + """ + if input.num_axes != output.num_axes: + raise ex.IncompatibleCoordSystemsError( + input=input, output=output, message="Axes must have the same number of dimensions" + ) + super().__init__(input=input, output=output) + + def inverse(self) -> BaseTransformationEdge: + return IdentityEdge(input=self.output, output=self.input) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + self._validate_transform_points_shapes(points) + return points + + def to_affine(self) -> AffineEdge: + return AffineEdge( + linear=np.identity(self.input.num_axes), + input=self.input, + output=self.output, + ) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.Identity: + return ozm06ct.Identity( + name=name, + input=input, + output=output, + ) + + +class MapAxisEdge(BaseTransformationEdge): + """The MapAxis transformation from the NGFF specification.""" + + def __init__( + self, + *, + input: vert.CoordSystem, + output: vert.CoordSystem, + input_to_output: Mapping[vert.Axis, vert.Axis], + ) -> None: + """ + Parameters + ---------- + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. Must + have the same number of axes as `input` + input_to_output + A mapping from + + Raises + ------ + IncompatibleCoordSystemsError + if `input` and `output` don't have the number of axes + """ + + if input.num_axes != output.num_axes: + raise ex.IncompatibleCoordSystemsError( + input=input, output=output, message="Input and output must have the same number of axes" + ) + + unmapped_inputs = set(input.axes) + unmapped_outputs = set(output.axes) + for inp_ax, out_ax in input_to_output.items(): + try: + unmapped_inputs.remove(inp_ax) + except KeyError as e: + raise ex.MissingAxisError(axis=inp_ax, cs=self.input) from e + try: + unmapped_outputs.remove(out_ax) + except KeyError as e: + raise ex.MissingAxisError(axis=out_ax, cs=self.output) from e + + for ax in unmapped_inputs: + raise ex.UnmappedAxisError(axis=ax, cs=self.input) + for ax in unmapped_outputs: + raise ex.UnmappedAxisError(axis=ax, cs=self.output) + + self.input_to_output = input_to_output + super().__init__(input=input, output=output) + + def __repr__(self) -> str: + s = super().__repr__() + "\n" + s += "\n".join(f" {out.name} <- {inp.name}\n" for out, inp in self.input_to_output.items()) + return s + + def inverse(self) -> BaseTransformationEdge: + return MapAxisEdge( + input=self.output, + output=self.input, + input_to_output={out: inp for inp, out in self.input_to_output.items()}, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + self._validate_transform_points_shapes(points) + output_to_input = {out_ax: in_ax for in_ax, out_ax in self.input_to_output.items()} + new_indices = [self.input.axes.index(output_to_input[out_ax]) for out_ax in self.output.axes] + assert len(new_indices) == self.output.num_axes + mapped = points[:, new_indices] + assert isinstance(mapped, np.ndarray) + return mapped + + def to_affine(self) -> AffineEdge: + linear = np.zeros((self.output.num_axes, self.input.num_axes)) + for inp_ax, out_ax in self.input_to_output.items(): + inp_idx = self.input.axes.index(inp_ax) + out_idx = self.output.axes.index(out_ax) + linear[out_idx, inp_idx] = 1 + return AffineEdge(linear=linear, input=self.input, output=self.output) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.MapAxis: + mapAxis: list[int] = [self.input.axes.index(out_ax) for out_ax in self.output.axes] + return ozm06ct.MapAxis( + name=name, + mapAxis=tuple(mapAxis), + input=input, + output=output, + ) + + +class ProjectAxisEdge(BaseTransformationEdge): + dropped_inputs: Final[set[vert.Axis]] + created_outputs: Final[set[vert.Axis]] + + def __init__( + self, + *, + input: vert.CoordSystem, + output: vert.CoordSystem, + dropped_inputs: set[vert.Axis], + created_outputs: set[vert.Axis], + ) -> None: + """ + Parameters + ---------- + dropped_inputs + axes in `input` that will be dropped by this transformation + created_inputs + axes in `output` that will be set to 0 + + Raises + ------ + MissingAxisError + axis in `dropped_inputs` not in `input` + axis in `created_outputs` not in `output` + IncompatibleCoordSystemsError + when input can't be mapped to output given dropped_inputs and created_outputs + """ + + for axis in dropped_inputs: + if axis not in input.axes: + raise ex.MissingAxisError(axis=axis, cs=input) + for axis in created_outputs: + if axis not in output.axes: + raise ex.MissingAxisError(axis=axis, cs=output) + if input.num_axes - len(dropped_inputs) + len(created_outputs) != output.num_axes: + message = f"Can't map from {input} to {output}" + if dropped_inputs: + message += f" dropping {dropped_inputs}" + if created_outputs: + message += f" creating {created_outputs}" + raise ex.IncompatibleCoordSystemsError( + input=input, + output=output, + message=message, + ) + + self.dropped_inputs = set(dropped_inputs) + self.created_outputs = set(created_outputs) + super().__init__(input=input, output=output) + + def to_affine(self) -> AffineEdge: + linear = np.zeros((self.output.num_axes, self.input.num_axes), dtype=float) + input_indices = iter(range(self.input.num_axes)) + for out_idx, out_ax in enumerate(self.output.axes): + if out_ax in self.created_outputs: + continue + in_idx = next(input_indices) + if in_idx in self.dropped_inputs: + continue + linear[out_idx, in_idx] = 1 + + return AffineEdge(input=self.input, output=self.output, linear=linear) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + return self.to_affine().transform_points(points) + + def inverse(self) -> BaseTransformationEdge | 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(), + ) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.ProjectAxis: + return ozm06ct.ProjectAxis( + name=name, + 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, + input=input, + output=output, + ) + + +class TranslationEdge(BaseTransformationEdge): + """The Translation transformation from the NGFF specification.""" + + def __init__( + self, + *, + translation: ArrayLike, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + """ + Parameters + ---------- + 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. + + Raises + ------ + IncompatibleCoordSystemsError + If the input and output have different number of dimensions + UnexpectedShapeError + if `translation`'s shape doesn't match `input` + """ + if input.num_axes != output.num_axes: + raise ex.IncompatibleCoordSystemsError( + input=input, output=output, message="Number of input and output axes must be the same" + ) + expected_translation_shape = (input.num_axes,) + if translation.shape != expected_translation_shape: + raise ex.UnexpectedShapeError( + array_name="translation", + array_shape=translation.shape, + expected_shape=expected_translation_shape, + ) + self.translation = translation + super().__init__(input=input, output=output) + + def __repr__(self) -> str: + return super().__repr__() + str(self.translation) + + def inverse(self) -> BaseTransformationEdge: + return TranslationEdge( + translation=-self.translation, + input=self.output, + output=self.input, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + 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, + ) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.Translation: + return ozm06ct.Translation( + name=name, + input=input, + output=output, + translation=tuple(self.translation), + ) + + +class ScaleEdge(BaseTransformationEdge): + """The Scale transformation from the NGFF specification.""" + + def __init__( + self, + *, + scale: ArrayLike, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + """ + Parameters + ---------- + scale + A vector specifying the scale along each axis of `input`. + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + + Raises + ------ + UnexpectedShapeError + If scale doesn't have the same number of elements as input has axes + IncompatibleCoordSystemsError + If input and output have different number of axes + """ + expected_scale_shape = (input.num_axes,) + if scale.shape != expected_scale_shape: + raise ex.UnexpectedShapeError(array_name="scale", array_shape=scale.shape, expected_shape=expected_scale_shape) + if input.num_axes != output.num_axes: + raise ex.IncompatibleCoordSystemsError( + input=input, output=output, message="input and output must have same number of dimensions" + ) + self.scale = scale + super().__init__(input=input, output=output) + + def __repr__(self) -> str: + return super().__repr__() + str(self.scale) + + def inverse(self) -> ScaleEdge | None: + if any(s == 0 for s in self.scale): + return None + new_scale = 1 / self.scale + return ScaleEdge( + scale=new_scale, + input=self.output, + output=self.input, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + 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, + ) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.Scale: + return ozm06ct.Scale( + name=name, + input=input, + output=output, + scale=tuple(self.scale), + ) + + +class RotationEdge(BaseTransformationEdge): + """The Rotation transformation from the NGFF specification.""" + + rotation: Final[ArrayLike] + + def __init__( + self, + *, + linear_matrix: ArrayLike, + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + """ + Parameters + ---------- + linear_matrix + And orthonormal matrix of shape (output.num_axes, input.num_axes) + input + Input coordinate system of the transformation. + output + Output coordinate system of the transformation. + Raises + ------ + UnexpectedShapeError + if linear_matrix's shape isn't (output.num_axes, input.num_axes) + IncompatibleCoordSystemsError + if input and output don't have the same number of axes + DeterminantDifferentFromOne + if linear_matrix doesn't have determinant ~= 1 + NotOrthonormal + if linear_matrix is not orthonormal + """ + if input.num_axes != output.num_axes: + raise ex.IncompatibleCoordSystemsError( + input=input, output=output, message="input and output should have the same numbe rof axes" + ) + expected_shape = (output.num_axes, input.num_axes) + if linear_matrix.shape != expected_shape: + raise ex.UnexpectedShapeError( + array_name="linear_matrix", array_shape=linear_matrix.shape, expected_shape=expected_shape + ) + if not np.allclose(linear_matrix.T @ linear_matrix, np.identity(input.num_axes)): + raise ex.NotOrthonormalError(matrix=linear_matrix) + if not np.isclose(np.linalg.det(linear_matrix), 1.0): + raise ex.DeterminantDifferentFromOne(matrix=linear_matrix) + linear_matrix.flags.writeable = False + self.rotation = linear_matrix + super().__init__(input=input, output=output) + + def __repr__(self) -> str: + s = super().__repr__() + "\n" + s += "\n".join(str(row) for row in self.rotation) + return s + + def inverse(self) -> BaseTransformationEdge: + return RotationEdge( + linear_matrix=self.rotation.T, + input=self.output, + output=self.input, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + self._validate_transform_points_shapes(points) + return points @ self.rotation.T + + def to_affine(self) -> AffineEdge: + return AffineEdge( + linear=self.rotation, + input=self.input, + output=self.output, + ) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.Rotation: + return ozm06ct.Rotation( + name=name, + input=input, + output=output, + rotation=tuple(tuple(row) for row in self.rotation), + ) + + +class SequenceEdge(BaseTransformationEdge): + """The Sequence transformation from the NGFF specification.""" + + def __init__( + self, + *, + transformations: Sequence[BaseTransformationEdge], + ) -> None: + """ + Init the NgffSequence object. + + Parameters + ---------- + transformations + The transformations which compose the sequence. + Raises + ------ + EmptyTransformSequence + if `transformations` is empty + IncompatibleCoordSystemsError + if any item in `transformations` is incompatible with its neighbors + """ + if len(transformations) == 0: + raise ex.EmptyTransformSequenceError() + previous_transf = transformations[0] + for transf_idx, current_transf in enumerate(transformations[1:], start=1): + if previous_transf.output != current_transf.input: + raise ex.IncompatibleCoordSystemsError( + input=current_transf.input, + output=previous_transf.output, + message=( + f"Output of transformation #{transf_idx - 1} is different " + f"from input of transformation #{transf_idx}" + ), + ) + previous_transf = current_transf + self.transformations = transformations + super().__init__( + input=transformations[0].input, + output=transformations[-1].output, + ) + + 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 | None: + inverted: list[BaseTransformationEdge] = [] + for t in reversed(self.transformations): + inv = t.inverse() + if inv is None: + return None + inverted.append(inv) + return SequenceEdge(transformations=inverted) + + 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, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + return self.to_affine().transform_points(points) # FIXME + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.Sequence: + inner_transfs: list[ozm06ct.AnyTransform] = [] + for idx, t in enumerate(self.transformations): + inp = input if idx == 0 else None + out = output if idx == (len(self.transformations) - 1) else None + inner_transfs.append(t.to_ngff_transform(name=None, input=inp, output=out)) + + return ozm06ct.Sequence( + name=name, + input=input, + output=output, + transformations=tuple(inner_transfs), + ) + + +class ByDimensionEdge(BaseTransformationEdge): + """The ByDimension transformation from the NGFF specification.""" + + transformations: Final[Sequence[BaseTransformationEdge]] + + def __init__( + self, + *, + transformations: Sequence[BaseTransformationEdge], + input: vert.CoordSystem, + output: vert.CoordSystem, + ) -> None: + """ + 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. + Raises + ------ + MissingAxisError + if any input axis from `transformations` is not present in `input` or + if any output axis from `transformations` is not present in `output`. + AxisRedefinitionError + if an output axis is specified by more than one item of `transformations` + UnmappedAxisError + if axis of `output` is not covered by any item of `transformations` + """ + # 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 + defined_output_axes: set[str] = set() + for t in transformations: + for ax in t.input.axes: + if ax not in input.axes: + raise ex.MissingAxisError(axis=ax, cs=input) + for ax in t.output.axes: + if ax not in output.axes: + raise ex.MissingAxisError(axis=ax, cs=output) + if ax.name in defined_output_axes: + raise ex.AxisRedefinitionError(axis=ax) + defined_output_axes.add(ax.name) + for ax in output.axes: + if ax.name not in defined_output_axes: + raise ex.UnmappedAxisError(axis=ax, cs=output) + + self.transformations = tuple(transformations) + super().__init__(input=input, output=output) + + 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) -> BaseTransformationEdge | None: + inverse_transformations: list[BaseTransformationEdge] = [] + 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, + output=self.input, + ) + + def transform_points(self, points: ArrayLike) -> ArrayLike: + """ + Transform points (coordinates). + + Raises + ------ + UnexpectedShapeError + if `points`'s shape is incompatible with this transformation's input shape + """ + 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, + ) + + def to_ngff_transform( + self, + *, + name: str | None, + input: ozm06ct.CoordinateSystemIdentifier | None, + output: ozm06ct.CoordinateSystemIdentifier | None, + ) -> ozm06ct.ByDimension: + raise NotImplementedError() + # by_dim_transfs: list[ozm06ct.ByDimension] = [] + # 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( + # ozm06ct.ByDimension( + # # inputAxes=input_axes, + # # outputAxes=output_axes, + # transformations=t.to_ngff_transform(name=None, input=None, output=None), + # input=None, + # output=None, + # ) + # ) + # return ozm06ct.ByDimension( + # name=name, + # input=input, + # output=output, + # transformations=tuple(by_dim_transfs), + # ) + + +def parse_identity( + model: ozm06ct.Identity, + *, + input: vert.CoordSystem, + output: vert.CoordSystem | None, +) -> IdentityEdge: + """Parse an `Identity` NGFF transformation model into an `IdentityEdge`.""" + if output is None: + output = vert.CoordSystem.create_dummy(axes=input.num_axes) + return IdentityEdge(input=input, output=output) + + +def parse_translation( + model: ozm06ct.Translation, + *, + input: vert.CoordSystem, + output: vert.CoordSystem | None, +) -> TranslationEdge: + """Parse a `Translation` NGFF transformation model into a `TranslationEdge`.""" + if output is None: + output = vert.CoordSystem.create_dummy(axes=input.num_axes) + return TranslationEdge( + translation=np.asarray(model.translation, dtype=float), + input=input, + output=output, + ) + + +def parse_project_axis( + model: ozm06ct.ProjectAxis, *, input: vert.CoordSystem, output: vert.CoordSystem | None +) -> ProjectAxisEdge: + """Parse a `ProjectAxis` NGFF transformation model into a `ProjectAxisEdge`.""" + if output is None: + num_dropped_inputs = len(model.droppedInputs or ()) + num_created_outputs = len(model.createdOutputs or ()) + num_output_axes = input.num_axes - num_dropped_inputs + num_created_outputs + output = vert.CoordSystem.create_dummy(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, + ) + + +def parse_scale( + model: ozm06ct.Scale, + *, + input: vert.CoordSystem, + output: vert.CoordSystem | None, +) -> ScaleEdge: + """Parse a `Scale` NGFF transformation model into a `ScaleEdge`.""" + return ScaleEdge( + scale=np.asarray(model.scale, dtype=float), + input=input, + output=output or vert.CoordSystem.create_dummy(axes=input.num_axes), + ) + + +def parse_map_axis(model: ozm06ct.MapAxis, *, input: vert.CoordSystem, output: vert.CoordSystem | None) -> MapAxisEdge: + """Parse a `MapAxis` NGFF transformation model into a `MapAxisEdge`.""" + if output is None: + output = vert.CoordSystem.create_dummy(axes=[input.axes[i] for i in model.mapAxis]) + input_to_output = { + input.axes[inp_idx]: output_axis for inp_idx, output_axis in zip(model.mapAxis, output.axes, strict=True) + } + return MapAxisEdge(input=input, output=output, input_to_output=input_to_output) + + +def parse_affine(model: ozm06ct.Affine, *, input: vert.CoordSystem, output: vert.CoordSystem | None) -> AffineEdge: + """Parse an `Affine` NGFF transformation model into an `AffineEdge`.""" + num_output_axes = len(model.affine_matrix) # spec doesn't save last row + if output is None: + output = vert.CoordSystem.create_dummy(axes=num_output_axes) + affine_array = np.asarray(model.affine_matrix, dtype=float) + return AffineEdge(linear=affine_array[:, :-1], translation=affine_array[:, -1], input=input, output=output) + + +def parse_rotation(model: ozm06ct.Rotation, *, input: vert.CoordSystem, output: vert.CoordSystem | None) -> RotationEdge: + """Parse a `Rotation` NGFF transformation model into a `RotationEdge`.""" + num_output_axes = len(model.rotation_matrix) + if output is None: + output = vert.CoordSystem.create_dummy(axes=num_output_axes) + return RotationEdge( + linear_matrix=np.asarray(model.rotation_matrix, dtype=float), + input=input, + output=output, + ) + + +def parse_sequence( + model: ozm06ct.Sequence, + *, + input: vert.CoordSystem, + output: vert.CoordSystem | None, +) -> SequenceEdge: + """Parse a `Sequence` NGFF transformation model into a `SequenceEdge`.""" + parsed_inners: list[BaseTransformationEdge] = [] + + parsed = parse_ngff_transf( + input=input, + model=model.transformations[0], + output=output if len(model.transformations) == 1 else None, + ) + parsed_inners.append(parsed) + + for t in model.transformations[1:-1]: + parsed = parse_ngff_transf(input=parsed.output, output=None, 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(transformations=parsed_inners) + + +def parse_by_dimension( + model: ozm06ct.ByDimension, + *, + input: vert.CoordSystem, + output: vert.CoordSystem | None, +) -> ByDimensionEdge: + """Parse a `ByDimension` NGFF transformation model into a `ByDimensionEdge`.""" + if output is None: + max_out_idx = max(ax_idx for t in model.transformations for ax_idx in t.outputAxes) + output = vert.CoordSystem.create_dummy(axes=max_out_idx + 1) + + piecewise_transforms: list[BaseTransformationEdge] = [] + for t in model.transformations: + inp_axes = [input.axes[i] for i in t.inputAxes] + partial_input = vert.CoordSystem.create_dummy(axes=inp_axes) + + out_axes = [output.axes[i] for i in t.outputAxes] + partial_out = vert.CoordSystem.create_dummy(axes=out_axes) + + 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, + transformations=piecewise_transforms, + ) + + +def parse_ngff_transf( + input: vert.CoordSystem, + model: ozm06ct.AnyTransform, + output: vert.CoordSystem | None, +) -> BaseTransformationEdge: + """Parse an NGFF coordinate transformation model into a `BaseTransfEdge`""" + if isinstance(model, ozm06ct.Identity): + return parse_identity(model, input=input, output=output) + elif isinstance(model, ozm06ct.Translation): + return parse_translation(model, input=input, output=output) + elif isinstance(model, ozm06ct.Scale): + return parse_scale(model, input=input, output=output) + elif isinstance(model, ozm06ct.MapAxis): + return parse_map_axis(model, input=input, output=output) + elif isinstance(model, ozm06ct.Affine): + return parse_affine(model, input=input, output=output) + elif isinstance(model, ozm06ct.Rotation): + return parse_rotation(model, input=input, output=output) + elif isinstance(model, ozm06ct.Sequence): + return parse_sequence(model, input=input, output=output) + elif isinstance(model, ozm06ct.ByDimension): + return parse_by_dimension(model, input=input, output=output) + elif isinstance(model, ozm06ct.ProjectAxis): + return parse_project_axis(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..1bd1d5dce --- /dev/null +++ b/src/spatialdata/transformations/graph/vert.py @@ -0,0 +1,550 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from pathlib import PurePosixPath +from typing import TYPE_CHECKING, Callable, Final, Literal, NamedTuple, Type, TypedDict, assert_never + +import networkx as nx +import ome_zarr_models.v06.coordinate_transforms as ozm06ct + +from spatialdata.transformations.graph.edge import parse_ngff_transf + +if TYPE_CHECKING: + from spatialdata.transformations.graph.edge import BaseTransformationEdge + +import ome_zarr_models.v06.coordinate_transforms as ozm06ct + + +class AxisParsingException(Exception): + pass + + +@dataclass(frozen=True) +class Axis: + """Representation of a coordinate system axis""" + + name: Final[str] + type: Final[Literal["space", "channel"]] + discrete: bool | None + "this ios just for round-tripping NGFF for now and maybe shouldbe moved out of here" + unit: Final[str | None] = None + "unit of the axis. For a set of valid options see https://ngff.openmicroscopy.org/" + long_name: Final[str | None] = None + "a longer, human-friendly name for this axis" + + 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, + discrete=self.discrete, + ) + + def __hash__(self) -> int: + return hash((self.name, self.type, self.unit, self.long_name)) + + def __repr__(self) -> str: + return f"Axis(name={self.name}, type={self.type})" + + def __eq__(self, value: Axis, /) -> bool: # pyright: ignore[reportIncompatibleMethodOverride] + if not isinstance(value, Axis): # pyright: ignore[reportUnnecessaryIsInstance] + 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: ozm06ct.Axis) -> Axis: + """ + Parse an `Axis` from an ome-zarr-models axis model. + + Raises + ------ + AxisParsingException + if `model` doesn't have a name, has a type other than "channel" or "space", or has a unit + that isn't a string or None + """ + + name = model.name + if name is None: + raise AxisParsingException("Axis doesn't have a name") + if model.type is None: + raise AxisParsingException("Axis doesn't have a type") + if model.type not in ("channel", "space"): + raise AxisParsingException(f"Can't handle axis of type {model.type}") + if not isinstance(model.unit, (str, type(None))): + raise AxisParsingException("Can't handle axis unit") + return Axis( + name=name, + type=model.type, + unit=model.unit, + long_name=model.longName, + discrete=model.discrete, + ) + + def to_model(self) -> ozm06ct.Axis: + return ozm06ct.Axis( + discrete=self.discrete, + longName=self.long_name, + name=self.name, + type=self.type, + unit=self.unit, + ) + + +class CoordSystemParsingException(Exception): + pass + + +class DuplicateAxisNameError(Exception): + def __init__(self, *, axis_name: str) -> None: + self.axis_name = axis_name + super().__init__(f"Axis name '{axis_name}' is used more than once") + + +@dataclass(frozen=True) +class CoordSystem: + """ + A coordinate space that attaches semantics to coordinates via `axes` + """ + + axes: Final[tuple[Axis, ...]] + + def __post_init__(self) -> None: + """ + Raises + ------ + DuplicateAxisNameError + if `axes` contains axes with duplicate names + """ + seen_names: set[str] = set() + for axis in self.axes: + if axis.name in seen_names: + raise DuplicateAxisNameError(axis_name=axis.name) + seen_names.add(axis.name) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.axes})" + + def __hash__(self) -> int: + return hash(self.axes) + + def __eq__(self, value: CoordSystem, /) -> bool: # pyright: ignore[reportIncompatibleMethodOverride] + return self.axes == value.axes + + @classmethod + def try_from_model(cls, model: ozm06ct.CoordinateSystem) -> CoordSystem: + """ + Parse a `CoordSystem` from an ome-zarr-models coordinate system model. + + Raises + ------ + AxisParsingException + if any axis in `model.axes` fails to parse + DuplicateAxisNameError + if `model.axes` contains axes with duplicate names + """ + return CoordSystem( + axes=tuple(Axis.try_from_model(axis) for axis in model.axes), + ) + + @classmethod + def try_from_model_or_default[T](cls, model: ozm06ct.CoordinateSystem | None, *, default: T) -> CoordSystem | T: + if model is not None: + return CoordSystem.try_from_model(model) + return default + + def to_model(self, name: str) -> ozm06ct.CoordinateSystem: + return ozm06ct.CoordinateSystem( + name=name, + axes=tuple(ax.to_model() for ax in self.axes), + ) + + @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"] + + @classmethod + def create_dummy(cls, *, axes: int | Sequence[Axis]) -> CoordSystem: + if isinstance(axes, int): + axes = tuple( + Axis( + name=f"axis_{ax_idx}", + type="space", # FIXME + discrete=None, + ) + for ax_idx in range(axes) + ) + else: + axes = tuple(axes) + + return CoordSystem( + axes=axes, + ) + + +@dataclass(frozen=True) +class NgffCsPath: + group_path: Final[PurePosixPath] + name: Final[str] + + def __post_init__(self) -> None: + assert not self.group_path.is_absolute() + + def __str__(self) -> str: + return f"{self.__class__.__name__}(path={self.group_path}, name={self.name})" + + def lives_under(self, other: NgffCsPath) -> bool: + if not self.group_path.is_relative_to(other.group_path): + return False + if self.group_path == other.group_path: + return False + return True + + def lives_over(self, other: NgffCsPath) -> bool: + return other.lives_under(self) + + def lives_in_same_level_as(self, other: NgffCsPath) -> bool: + return self.group_path == other.group_path + + +@dataclass +class UserCreated: + name: str + +type Provenance = NgffCsPath | UserCreated + + +class CsVert: + id: uuid.UUID + coord_system: CoordSystem + provenance: Provenance + + def __init__(self, cs: CoordSystem, provenance: Provenance) -> None: + self.id = uuid.uuid4() + self.coord_system = cs + self.provenance = provenance + super().__init__() + + def __hash__(self) -> int: + return hash(self.id) + + def __eq__(self, value: CsVert, /) -> bool: # pyright: ignore [reportIncompatibleMethodOverride] + if not isinstance(value, CsVert): # pyright: ignore [reportUnnecessaryIsInstance] + return False + return self.id == value.id + + def __str__(self) -> str: + return f"CsVert(id={self.id}, coord_system={self.coord_system}, provenance={self.provenance})" + + def to_ngff_cs(self) -> ozm06ct.CoordinateSystem: + return self.coord_system.to_model(name=self.name) + + @property + def name(self) -> str: + return self.provenance.name + + @property + def namespace(self) -> PurePosixPath: + if isinstance(self.provenance, NgffCsPath): + return self.provenance.group_path + if isinstance(self.provenance, UserCreated): # pyright: ignore [reportUnnecessaryIsInstance] + return PurePosixPath("") + assert_never(self.provenance) + + def to_ngff_cs_id(self, relative_to: PurePosixPath) -> ozm06ct.CoordinateSystemIdentifier: + ns = self.namespace.relative_to(relative_to) + return ozm06ct.CoordinateSystemIdentifier( + name=self.provenance.name, + path=None if ns == PurePosixPath("") else str(ns), + ) + + +class EdgeData(TypedDict): # NOTE: TypedDict because networkx uses dicts + name: str | None + transform: BaseTransformationEdge + +def hash_edge_data(ed: EdgeData) -> int: + return hash((ed["name"], ed["transform"])) + + +class Edge(NamedTuple): + origin: CsVert + target: CsVert + data: EdgeData + + def __hash__(self) -> int: + return hash((self.origin, self.target)) + + def points_down_ngff_hierarchy(self) -> bool: + if not isinstance(self.origin.provenance, NgffCsPath): + return False + if not isinstance(self.target.provenance, NgffCsPath): + return False + return self.origin.provenance.lives_over(self.target.provenance) + + def points_up_ngff_hierarchy(self) -> bool: + if not isinstance(self.origin.provenance, NgffCsPath): + return False + if not isinstance(self.target.provenance, NgffCsPath): + return False + return self.target.provenance.lives_over(self.origin.provenance) + + def is_horizontal_in_ngff(self) -> bool: + "Returns true is edge is from an NGFF hierachy and doesn't point up or down the tree" + if not isinstance(self.origin.provenance, NgffCsPath): + return False + if not isinstance(self.target.provenance, NgffCsPath): + return False + return self.origin.provenance.lives_in_same_level_as(self.target.provenance) + + def to_model(self, relative_to: PurePosixPath) -> ozm06ct.AnyTransform: + input = self.origin.to_ngff_cs_id(relative_to) + output = self.target.to_ngff_cs_id(relative_to) + return self.data["transform"].to_ngff_transform( + name=self.data["name"], + input=input, + output=output, + ) + + def vert_that_isnt(self, vert: CsVert) -> CsVert: + if self.origin == vert: + return self.target + if self.target == vert: + return self.origin + raise ValueError(f"Edge does not contain {vert}") + +type VertFilter = Callable[[CsVert, Spaces], bool] + +class Verts: + def __init__(self, verts: Iterator[CsVert], graph: Spaces, filters: Sequence[VertFilter]) -> None: + self.verts = verts + self.graph = graph + self.filters: Sequence[VertFilter] = filters + super().__init__() + + def __iter__(self) -> Verts: + return self + + def __next__(self) -> CsVert: + for vert in self.verts: + if all(f(vert, self.graph) for f in self.filters): + return vert + raise StopIteration() + + def filter(self, filter: Callable[[CsVert, Spaces], bool]) -> Verts: + return Verts(verts=self.verts, graph=self.graph, filters=[*self.filters, filter]) + + def ngff_siblings(self, *, group_path: PurePosixPath) -> Verts: + def filter(vert: CsVert, graph: Spaces) -> bool: + if not isinstance(vert.provenance, NgffCsPath): + return False + return vert.provenance.group_path == group_path + + return self.filter(filter) + + def connected_to(self, other: CsVert) -> Verts: + def filter(vert: CsVert, graph: Spaces) -> bool: + for _ in graph.get_edges_between(vert, other): + return True + return False + return self.filter(filter) + + def has_single_neighbor(self) -> Verts: + def filter(v: CsVert, g: Spaces) -> bool: + return len(g.all_neighbors(v)) == 1 + + return self.filter(filter) + + def collect[T](self, f: Callable[[Iterator[CsVert]], T]) -> T: + return f(self) + + def any(self) -> bool: + for _ in self.verts: + return True + return False + +type EdgeFilter = Callable[[Edge, Spaces], bool] + +class Edges: + def __init__( + self, edges: Iterator[Edge], graph: Spaces, filters: Sequence[EdgeFilter] + ) -> None: + self.edges = edges + self.graph = graph + self.filters = filters + + def __iter__(self) -> Edges: + return self + + def __next__(self) -> Edge: + for e in self.edges: + if all(f(e, self.graph) for f in self.filters): + return e + raise StopIteration() + + def filter(self, filter: Callable[[Edge, Spaces], bool]) -> Edges: + return Edges(edges=self.edges, graph=self.graph, filters=[*self.filters, filter]) + + def directed(self, *, fro: CsVert, to: CsVert) -> Edges: + def filter(edge: Edge, g: Spaces) -> bool: + return edge.origin == fro and edge.target == to + return self.filter(filter) + + def transf_is(self, classes: tuple[Type[BaseTransformationEdge], ...]) -> Edges: + def filter(edge: Edge, g: Spaces) -> bool: + return isinstance(edge.data["transform"], classes) + return self.filter(filter) + + def containing(self, cs: CsVert) -> Edges: + def filter(edge: Edge, graph: Spaces) -> bool: + return edge.origin == cs or edge.target == cs + return self.filter(filter) + + def between(self, cs1: CsVert, cs2: CsVert) -> Edges: + def filter(edge: Edge, graph: Spaces) -> bool: + if (edge.origin, edge.target) == (cs1, cs2): + return True + if (edge.target, edge.origin) == (cs2, cs1): + return True + return False + return self.filter(filter) + + def ngff_siblings(self) -> Edges: + def filter(edge: Edge, graph: Spaces) -> bool: + return edge.is_horizontal_in_ngff() + return self.filter(filter=filter) + + def map[T](self, f: Callable[[Edge], T]) -> Iterator[T]: + return map(f, self.edges) + + def collect[T](self, f: Callable[[Iterator[Edge]], T]) -> T: + return f(self) + + def any(self) -> bool: + for _ in self.edges: + return True + return False + + +class Spaces: + def __init__(self) -> None: + self._graph: nx.MultiDiGraph[CsVert, dict[str, None], EdgeData] = nx.MultiDiGraph() + super().__init__() + + def incorporate_ngff_graph( + self, + group_path: PurePosixPath, + coord_systems: Sequence[ozm06ct.CoordinateSystem], + transforms: Sequence[ozm06ct.AnyTransform], + ) -> None: + for cs in coord_systems: + parsed_cs = CoordSystem.try_from_model(cs) + self.add_cs(cs=parsed_cs, provenance=NgffCsPath(group_path=group_path, name=cs.name)) + + for transf in transforms: + assert transf.input is not None # as per the spec + assert transf.output is not None # as per the spec + assert transf.input.name is not None # as per the spec + assert transf.output.name is not None # as per the spec + input_path = NgffCsPath(group_path / (transf.input.path or ""), transf.input.name) + output_path = NgffCsPath(group_path / (transf.output.path or ""), transf.output.name) + + input = self.find_node(ngff_path=input_path) + output = self.find_node(ngff_path=output_path) + parsed = parse_ngff_transf(input=input.coord_system, output=output.coord_system, model=transf) + self.add_transformation(name=transf.name, input=input, output=output, transf=parsed) + + def incorporate(self, other: Spaces): + # FIXME: maybe there are some nodes or edges that cannot be incorporated. + # We might need to checkk ALL OF THEM before partially updating self, + # so we can succeed or fail in a transaction-like manner + + for vert in other.verts(): + self._graph.add_node(vert) # FIXME: + for edge in other.edges(): + self.add_transformation( + input=edge.origin, output=edge.target, transf=edge.data["transform"], name=edge.data["name"] + ) + + def all_neighbors(self, cs: CsVert) -> set[CsVert]: + import itertools + return set(itertools.chain(self._graph.successors(cs), self._graph.predecessors(cs))) + + def add_cs(self, *, cs: CoordSystem, provenance: Provenance) -> CsVert: + node = CsVert(cs, provenance) + self._graph.add_node(node) + return node + + def add_transformation(self, *, input: CsVert, output: CsVert, transf: BaseTransformationEdge, name: str | None) -> None: + assert input.coord_system == transf.input + assert output.coord_system == transf.output + _ = self._graph.add_edge(input, output, **EdgeData(name=name, transform=transf)) + + def find_direct_edges(self, *, fro: CsVert, to: CsVert) -> Iterator[EdgeData]: + for src_vert, tgt_vert, edge_data in self._graph.edges.data(): + if self._graph.nodes[src_vert] != fro: + continue + if self._graph.nodes[tgt_vert] != to: + continue + yield edge_data + + def get_edges_between(self, a: CsVert, b: CsVert) -> Iterator[Edge]: + for src_vert, tgt_vert, edge_data in self._graph.edges.data(): + if (src_vert, tgt_vert) in ((a,b), (b,a)): + yield Edge(data=edge_data, origin=src_vert, target=tgt_vert) + + def find_node(self, ngff_path: NgffCsPath) -> CsVert: + for node in self._graph.nodes: + provenance = node.provenance + if not isinstance(provenance, NgffCsPath): + continue + if provenance == ngff_path: + return node + raise KeyError(f"could not find CS with key {ngff_path}") + + def verts(self) -> Verts: + return Verts(filters=[], graph=self, verts=iter(self._graph.nodes)) + + def edges(self) -> Edges: + def edges(): + for src_node, tgt_node, edge_data in self._graph.edges.data(): + yield Edge(origin=src_node, target=tgt_node, data=edge_data) + return Edges(edges=edges(), filters=[], graph=self) diff --git a/tests/core/transformation_manager/conftest.py b/tests/core/transformation_manager/conftest.py new file mode 100644 index 000000000..696935996 --- /dev/null +++ b/tests/core/transformation_manager/conftest.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 + +""" +Fixtures for transformation manager tests. +""" + +from __future__ import annotations + +import pytest + +from spatialdata.transformations import Affine, Scale, Translation +from spatialdata.transformations.ngff.ngff_coordinate_system import NgffAxis, NgffCoordinateSystem + + +def get_ngff_coodinate_system(cs_name: str) -> NgffCoordinateSystem: + return NgffCoordinateSystem( + name=cs_name, + axes=[NgffAxis(name="x", type="space", unit="micrometer"), NgffAxis(name="y", type="space", unit="micrometer")], + ) + + +@pytest.fixture +def one_point_graph() -> tuple[list[NgffCoordinateSystem], list[Translation]]: + """Fixture providing a single point graph with one coordinate system and one transformation.""" + coordinate_systems = [get_ngff_coodinate_system("cs1")] + transformations = [Translation(translation=[1.0, 2.0], axes=("x", "y"))] + return coordinate_systems, transformations + + +@pytest.fixture +def fully_connected_two_point_graph() -> tuple[list[NgffCoordinateSystem], list[Translation]]: + """Fixture providing a fully connected two-point graph with two coordinate systems and transformations.""" + coordinate_systems = [get_ngff_coodinate_system("cs1"), get_ngff_coodinate_system("cs2")] + transformations = [ + Translation(translation=[-1.0, -2.0], axes=("x", "y")), + ] + return coordinate_systems, transformations + + +@pytest.fixture +def four_point_graph() -> tuple[list[NgffCoordinateSystem], list[Scale | Translation]]: + """Fixture providing a four-point graph with four coordinate systems and transformations.""" + coordinate_systems = [ + get_ngff_coodinate_system("cs1"), + get_ngff_coodinate_system("cs2"), + get_ngff_coodinate_system("cs3"), + get_ngff_coodinate_system("cs4"), + ] + transformations = [ + Translation(translation=[1.0, 2.0], axes=("x", "y")), # cs1 -> cs2 + Translation(translation=[3.0, 4.0], axes=("x", "y")), # cs2 -> cs3 + Scale( + scale=[ + 2, + 2, + ], + axes=("x", "y"), + ), # cs3 -> cs4; + Translation(translation=[4.0, 6.0], axes=("x", "y")), # cs1 -> cs3 (consistent with cs1->cs2 and cs2->cs3) + ] + return coordinate_systems, transformations + + +@pytest.fixture +def five_point_graph() -> tuple[list[NgffCoordinateSystem], list[Scale | Translation | Affine]]: + """Fixture providing a five-point graph with five coordinate systems and five transformations.""" + + coordinate_systems = [ + get_ngff_coodinate_system("cs1"), + get_ngff_coodinate_system("cs2"), + get_ngff_coodinate_system("cs3"), + get_ngff_coodinate_system("cs4"), + get_ngff_coodinate_system("cs5"), + ] + transformations = [ + Translation(translation=[1.0, 2.0], axes=("x", "y")), # cs1 -> cs2, cs4 -> cs3 + Translation(translation=[3.0, 4.0], axes=("x", "y")), # cs2 -> cs3, cs1 -> cs4 + Translation(translation=[4.0, 6.0], axes=("x", "y")), # cs1 -> cs3 (consistent with cs1->cs2 and cs2->cs3) + Scale( + scale=[ + 2, + 2, + ], + axes=("x", "y"), + ), # cs3 -> cs5; + Affine( + matrix=[[2, 0, 0], [0, 2, 0], [0, 0, 1]], + input_axes=("x", "y"), + output_axes=("x", "y"), + ), # cs3 -> cs5 + ] + return coordinate_systems, transformations diff --git a/tests/core/transformation_manager/test_transformation_manager.py b/tests/core/transformation_manager/test_transformation_manager.py new file mode 100644 index 000000000..c1f5d8300 --- /dev/null +++ b/tests/core/transformation_manager/test_transformation_manager.py @@ -0,0 +1,643 @@ +#!/usr/bin/env python3 + +""" +Unit tests for the TransformationManager class in _core/transformation_manager/__init__.py. + +This module tests all the functionality of the TransformationManager class to achieve 100% coverage. +""" + +from __future__ import annotations + +import pytest + +from spatialdata._core.transformation_manager import TRANSFORM_KEY, TransformationManager +from spatialdata._core.transformation_manager.exceptions import ( + CannotRemoveCoordinateSystemError, + CoordinateSystemAlreadyExistsError, + CoordinateSystemNotFoundError, + ElementAlreadyExistsError, + ElementNotRegisteredToAnyCoordinateSystemError, + InvalidPathError, + TransformationNotFoundError, + TransformationPathAmbiguousError, + TransformationPathAmbiguousMultipleEdgeExpectedError, + TransformationPathAmbiguousNoEdgeExpectedError, + TransformationPathNotFoundError, + TransformationPathNotSimple, +) +from spatialdata.transformations.transformations import Sequence + + +def test_initialization(): + """Test that TransformationManager initializes correctly.""" + + tm = TransformationManager() + assert len(tm._graph.nodes()) == 0 + assert len(tm._graph.edges()) == 0 + assert tm._element_to_cs_mapping == {} + + +def test_add_coordinate_system(one_point_graph): + """Test adding a coordinate system.""" + + tm = TransformationManager() + [cs1], _ = one_point_graph + + tm.add_coordinate_system(cs1) + assert cs1 in tm._graph.nodes() + + +def test_add_coordinate_system_duplicate(one_point_graph): + """Test adding an already existing coordinate system""" + tm = TransformationManager() + [cs1], _ = one_point_graph + + tm.add_coordinate_system(cs1) + with pytest.raises(CoordinateSystemAlreadyExistsError, match=f"Coordinate system '{cs1.name}' already exists"): + tm.add_coordinate_system(cs1) + + +def test_remove_coordinate_system(one_point_graph): + """Test removing a coordinate system.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + + tm.add_coordinate_system(cs1) + tm.remove_coordinate_system(cs1) + + assert cs1 not in tm._graph.nodes() + + +def test_remove_coordinate_system_nonexistent(one_point_graph): + """Test that removing a non-existent coordinate system raises CoordinateSystemNotFoundError.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + + # Add the coordinate system first + tm.add_coordinate_system(cs1) + + # Remove it + tm.remove_coordinate_system(cs1) + + # Try to remove it again - should raise CoordinateSystemNotFoundError + with pytest.raises(CoordinateSystemNotFoundError): + tm.remove_coordinate_system(cs1) + + +def test_remove_coordinate_system_with_associated_transformations(fully_connected_two_point_graph): + """ + Test that removing a coordinate system with associated transformations raises CannotRemoveCoordinateSystemError. + """ + tm = TransformationManager() + [cs1, cs2], [transformation] = fully_connected_two_point_graph + + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + tm.add_transformation(cs1, cs2, transformation) + + # Should raise CannotRemoveCoordinateSystem when trying to remove a coordinate system with transformations + with pytest.raises(CannotRemoveCoordinateSystemError, match="Cannot remove coordinate system"): + tm.remove_coordinate_system(cs1) + + # Should raise CannotRemoveCoordinateSystem when trying to remove a coordinate system with transformations + with pytest.raises(CannotRemoveCoordinateSystemError, match="Cannot remove coordinate system"): + tm.remove_coordinate_system(cs2) + + +def test_remove_coordinate_system_with_belonging_elements(one_point_graph): + """Test that removing a coordinate system with element associations raises CannotRemoveCoordinateSystemError.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + tm.add_coordinate_system(cs1) + tm.add_element("image1", cs1) + + with pytest.raises(CannotRemoveCoordinateSystemError, match="Cannot remove coordinate system"): + tm.remove_coordinate_system(cs1) + + +def test_list_coordinate_systems(fully_connected_two_point_graph): + """Test listing all coordinate systems.""" + tm = TransformationManager() + [cs1, cs2], _ = fully_connected_two_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + systems = tm.list_coordinate_systems() + assert len(systems) == 2 + assert cs1 in systems + assert cs2 in systems + + +def test_add_element(one_point_graph): + """Test adding an element with an existing coordinate system.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + tm.add_coordinate_system(cs1) + tm.add_element("image1", cs1) + + mapping = tm._element_to_cs_mapping + assert "image1" in mapping + assert mapping["image1"] == cs1 + + +def test_add_element_duplicate(one_point_graph): + """Test that adding a duplicate element raises ElementAlreadyExistsError.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + tm.add_coordinate_system(cs1) + element_name = "image1" + tm.add_element(element_name, cs1) + + # Try to add the same element again - should raise ElementAlreadyExistsError + with pytest.raises( + ElementAlreadyExistsError, match=f"Element '{element_name}' already exists in the transformation manager" + ): + tm.add_element(element_name, cs1) + + +def test_unset_element(one_point_graph): + """Test unsetting an element.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + tm.add_coordinate_system(cs1) + tm.add_element("image1", cs1) + tm._unset_element("image1") + + assert "image1" not in tm._element_to_cs_mapping + + +def test_unset_element_nonexistent(one_point_graph): + """Test that unsetting a non-existent element raises ElementNotFoundError.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + tm.add_coordinate_system(cs1) + element_name = "image1" + + # Try to unset non-existent element + with pytest.raises(ElementNotRegisteredToAnyCoordinateSystemError): + tm._unset_element(element_name) + + +def test_add_transformation(fully_connected_two_point_graph): + """Test adding a transformation between coordinate systems.""" + + tm = TransformationManager() + [cs1, cs2], [transform] = fully_connected_two_point_graph + + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + tm.add_transformation(cs1, cs2, transform) + + assert tm._graph.has_edge(cs1, cs2) + # Get the first edge key and check the transformation + edge_key = tm._get_edge_key((cs1, cs2, transform)) + assert tm._graph[cs1][cs2][edge_key][TRANSFORM_KEY] == transform + + +def test_add_transformation_existing_edge(fully_connected_two_point_graph): + """Test adding a transformation between coordinate systems that already have an edge.""" + tm = TransformationManager() + [cs1, cs2], [transform] = fully_connected_two_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + tm.add_transformation(cs1, cs2, transform) + assert tm._graph.has_edge(cs1, cs2) + + # Add same transformation again between the same coordinate systems + tm.add_transformation(cs1, cs2, transform) + # Should have just rewritten the edge + assert tm._graph.has_edge(cs1, cs2) + assert len(tm._graph[cs1][cs2]) == 1 + + +def test_add_transformation_nonexistent_cs(fully_connected_two_point_graph): + """Test that adding a transformation with non-existent coordinate systems raises CoordinateSystemNotFoundError.""" + tm = TransformationManager() + [cs1, cs2], [transform] = fully_connected_two_point_graph + + expected_msg = "Coordinate system .* not found in the transformation manager." + with pytest.raises(CoordinateSystemNotFoundError, match=expected_msg): + tm.add_transformation(cs1, cs2, transform) + + # Add one coordinate system + tm.add_coordinate_system(cs1) + + with pytest.raises(CoordinateSystemNotFoundError, match=expected_msg): + tm.add_transformation(cs1, cs2, transform) + + +def test_get_direct_transformations(fully_connected_two_point_graph): + """Test getting an existing transformation.""" + tm = TransformationManager() + [cs1, cs2], [transform] = fully_connected_two_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + tm.add_transformation(cs1, cs2, transform) + + retrieved = tm.get_direct_transformations(cs1, cs2) + assert retrieved == [transform] + + +def test_get_direct_transformations_nonexistent(fully_connected_two_point_graph): + """Test getting a non-existent transformation.""" + tm = TransformationManager() + coordinate_systems, _transformations = fully_connected_two_point_graph + cs1, cs2 = coordinate_systems + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + retrieved = tm.get_direct_transformations(cs1, cs2) + assert retrieved == [] + + +def test_remove_all_transformations_between_coordinate_systems(fully_connected_two_point_graph, mocker): + """Test removing all transformations between coordinate systems.""" + + tm = TransformationManager() + [cs1, cs2], [transform] = fully_connected_two_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + tm.add_transformation(cs1, cs2, transform) + + spy = mocker.spy(tm._graph, "remove_edge") + tm.remove_all_transformations_between_coordinate_systems(cs1, cs2) + assert spy.call_count == 1 + assert not tm._graph.has_edge(cs1, cs2) + + +def test_remove_all_transformation_nonexistent(fully_connected_two_point_graph, mocker): + """Test that removing non-existent transformations between coordinate systems raises TransformationNotFoundError.""" + + tm = TransformationManager() + [cs1, cs2], _ = fully_connected_two_point_graph + tm._graph.add_node(cs1) + tm._graph.add_node(cs2) + + spy = mocker.spy(tm._graph, "remove_edge") + tm.remove_all_transformations_between_coordinate_systems(cs1, cs2) + assert spy.call_count == 0 + + +def test_remove_specific_transformation_between_coordinate_systems(five_point_graph): + """Test removing specific transformation between coordinate systems.""" + + tm = TransformationManager() + [_cs1, _cs2, cs3, _cs4, cs5], [_transform1, _transform2, _transform3, transform4, transform5] = five_point_graph + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs5) + + tm.add_transformation(cs3, cs5, transform4) + tm.add_transformation(cs3, cs5, transform5) + + tm.remove_specific_transformation(cs3, cs5, transform4) + edge_key = tm._get_edge_key((cs3, cs5, transform4)) + assert not tm._graph.has_edge(cs3, cs5, key=edge_key) + + +def test_remove_specific_transformation_between_coordinate_systems_non_existent(five_point_graph): + """ + Test that removing non-existent specific transformation between coordinate systems raises + TransformationNotFoundError. + """ + + tm = TransformationManager() + [_cs1, _cs2, cs3, _cs4, cs5], [_transform1, _transform2, _transform3, transform4, _transform5] = five_point_graph + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs5) + + with pytest.raises( + TransformationNotFoundError, match=f"Transformation from '{cs3.name}' to '{cs5.name}' not found" + ): + tm.remove_specific_transformation(cs3, cs5, transform4) + + +def test_get_all_shortest_transformation_sequences_direct(four_point_graph): + """Test getting all shortest transformation sequences for a direct transformation path.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4], [transform1, transform2, _transform3, transform4] = four_point_graph + + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs3, transform4) + + sequences = tm.get_all_shortest_transformation_sequences(cs1, cs3) + assert len(sequences) == 1 + assert Sequence([transform4]) in sequences + + +def test_get_all_shortest_transformation_sequences_indirect_one_path(four_point_graph): + """Test getting all shortest transformation sequences for one indirect transformation path.""" + tm = TransformationManager() + [cs1, cs2, cs3, _cs4], [transform1, transform2, _transform3, _transform4] = four_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + + sequences = tm.get_all_shortest_transformation_sequences(cs1, cs3) + assert len(sequences) == 1 + assert Sequence([transform1, transform2]) in sequences + + +def test_get_all_shortest_transformation_sequences_indirect_two_paths(five_point_graph): + """Test getting all shortest transformation sequences for two indirect transformation paths.""" + tm = TransformationManager() + ([cs1, cs2, cs3, cs4, _cs5], [transform1, transform2, _transform3, _transform4, _transform5]) = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + + sequences = tm.get_all_shortest_transformation_sequences(cs1, cs3) + assert len(sequences) == 2 + assert Sequence([transform1, transform2]) in sequences + assert Sequence([transform2, transform1]) in sequences + + +def test_get_all_shortest_transformation_sequences_no_path(four_point_graph): + """Test that getting all shortest transformation sequences with no path raises TransformationPathNotFoundError.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4], [transform1, transform2, _transform3, _transform4] = four_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + + expected_error_msg = f"No transformation path found from {cs1.name} to {cs4.name}" + with pytest.raises(TransformationPathNotFoundError, match=expected_error_msg): + tm.get_all_shortest_transformation_sequences(cs1, cs4) + + +def test_get_all_shortest_transformation_sequences_multiple_paths_multiple_edges_success(five_point_graph): + """Test getting all shortest transformation sequences, with multiple path and multiple edges between nodes.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4, cs5], [transform1, transform2, transform3, transform4, transform5] = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + tm.add_coordinate_system(cs5) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + # tm.add_transformation(cs1, cs3, transform3) + tm.add_transformation(cs3, cs5, transform4) + tm.add_transformation(cs3, cs5, transform5) + + expected_intermediate_edges = [(cs3, cs5, transform4)] + + sequences = tm.get_all_shortest_transformation_sequences( + cs1, cs5, expected_intermediate_edges=expected_intermediate_edges + ) + assert len(sequences) == 2 + assert Sequence([transform1, transform2, transform4]) in sequences + assert Sequence([transform2, transform1, transform4]) in sequences + + +def test_get_all_shortest_transformation_sequences_multiple_paths_multiple_edges_failure(five_point_graph): + """Test getting all shortest transformation sequences, with multiple path and multiple edges between nodes.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4, cs5], [transform1, transform2, transform3, transform4, transform5] = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + tm.add_coordinate_system(cs5) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + tm.add_transformation(cs1, cs3, transform3) + tm.add_transformation(cs3, cs5, transform4) + tm.add_transformation(cs3, cs5, transform5) + + expected_msg = f"Transformation Path ambiguous from {cs1.name} to {cs5.name}." + with pytest.raises(TransformationPathAmbiguousError, match=expected_msg): + tm.get_all_shortest_transformation_sequences(cs1, cs5) + + +def test_get_all_transformation_sequences(four_point_graph): + """Test getting all transformation sequences.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4], [transform1, transform2, transform3, transform4] = four_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs3, cs4, transform3) + tm.add_transformation(cs1, cs3, transform4) + + sequences = tm.get_all_transformation_sequences(cs1, cs4) + assert len(sequences) == 2 + assert Sequence([transform1, transform2, transform3]) in sequences + assert Sequence([transform4, transform3]) in sequences + + +def test_get_all_transformation_sequences_multiple_paths(five_point_graph): + """Test getting all transformation sequences, with multiple path and multiple edges between nodes.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4, _cs5], [transform1, transform2, transform3, _transform4, _transform5] = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + tm.add_transformation(cs1, cs3, transform3) + + sequences = tm.get_all_transformation_sequences(cs1, cs3) + assert len(sequences) == 3 + assert Sequence([transform1, transform2]) in sequences + assert Sequence([transform2, transform1]) in sequences + assert Sequence([transform3]) in sequences + + +def test_get_all_transformation_sequences_multiple_paths_multiple_edges_success(five_point_graph): + """Test getting all transformation sequences, with multiple paths and multiple edges between nodes.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4, cs5], [transform1, transform2, transform3, transform4, transform5] = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + tm.add_coordinate_system(cs5) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + tm.add_transformation(cs1, cs3, transform3) + tm.add_transformation(cs3, cs5, transform4) + tm.add_transformation(cs3, cs5, transform5) + + expected_intermediate_edges = [(cs3, cs5, transform4)] + sequences = tm.get_all_transformation_sequences(cs1, cs5, expected_intermediate_edges=expected_intermediate_edges) + assert len(sequences) == 3 + assert Sequence([transform1, transform2, transform4]) in sequences + assert Sequence([transform2, transform1, transform4]) in sequences + assert Sequence([transform3, transform4]) in sequences + + +def test_get_all_transformation_sequences_multiple_paths_multiple_edges_no_edge_expected(five_point_graph): + """Test getting all transformation sequences with multiple paths and multiple edges, no edge expected.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4, cs5], [transform1, transform2, transform3, transform4, transform5] = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + tm.add_coordinate_system(cs5) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + tm.add_transformation(cs1, cs3, transform3) + tm.add_transformation(cs3, cs5, transform4) + tm.add_transformation(cs3, cs5, transform5) + + expected_msg = f"Transformation Path ambiguous from {cs3.name} to {cs5.name}." + with pytest.raises(TransformationPathAmbiguousNoEdgeExpectedError, match=expected_msg): + tm.get_all_transformation_sequences(cs1, cs5) + + +def test_get_all_transformation_sequences_multiple_paths_multiple_edges_multiple_expected(five_point_graph): + """Test getting all transformation sequences with multiple paths and multiple edges, multiple edges expected.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4, cs5], [transform1, transform2, transform3, transform4, transform5] = five_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + tm.add_coordinate_system(cs5) + + tm.add_transformation(cs1, cs2, transform1) + tm.add_transformation(cs2, cs3, transform2) + tm.add_transformation(cs1, cs4, transform2) + tm.add_transformation(cs4, cs3, transform1) + tm.add_transformation(cs1, cs3, transform3) + tm.add_transformation(cs3, cs5, transform4) + tm.add_transformation(cs3, cs5, transform5) + + expected_intermediate_edges = [ + (cs3, cs5, transform4), + (cs3, cs5, transform5), + ] + + expected_msg = f"Transformation Path ambiguous from {cs3.name} to {cs5.name}." + with pytest.raises(TransformationPathAmbiguousMultipleEdgeExpectedError, match=expected_msg): + tm.get_all_transformation_sequences(cs1, cs5, expected_intermediate_edges=expected_intermediate_edges) + + +def test_get_transformation_sequences_from_simple_paths_after_disambiguation_with_non_simple_paths( + fully_connected_two_point_graph, +): + """Test _get_transformation_sequences_from_simple_paths_after_disambiguation with non-simple paths.""" + tm = TransformationManager() + [cs1, cs2], _ = fully_connected_two_point_graph + + tm.add_coordinate_system(cs1) + + paths = [[cs1, cs2, cs1, cs2]] + + with pytest.raises(TransformationPathNotSimple): + tm._get_transformation_sequences_from_simple_paths_after_disambiguation(paths, ()) + + +def test_get_transformation_sequences_from_simple_paths_after_disambiguation_with_single_node_path( + fully_connected_two_point_graph, +): + """Test _get_transformation_sequences_from_simple_paths_after_disambiguation with a path with only one node.""" + tm = TransformationManager() + [cs1, cs2], [transformation] = fully_connected_two_point_graph + + tm.add_coordinate_system(cs1) + + paths = [[cs1]] + with pytest.raises(InvalidPathError): + tm._get_transformation_sequences_from_simple_paths_after_disambiguation(paths, ()) + + +def test_get_all_transformation_sequences_no_path(four_point_graph): + """Test getting all transformation sequences.""" + tm = TransformationManager() + [cs1, cs2, cs3, cs4], [transform1, transform2, transform3, transform4] = four_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + tm.add_coordinate_system(cs3) + tm.add_coordinate_system(cs4) + + expected_error_msg = f"No transformation path found from {cs1.name} to {cs4.name}" + with pytest.raises(TransformationPathNotFoundError, match=expected_error_msg): + tm.get_all_shortest_transformation_sequences(cs1, cs4) + + +def test_get_element_coordinate_system_success(one_point_graph): + """Test successfully getting an element's coordinate system.""" + tm = TransformationManager() + [cs1], _ = one_point_graph + tm.add_coordinate_system(cs1) + tm.add_element("image1", cs1) + + # This should successfully return the coordinate system + result = tm.get_element_coordinate_system("image1") + assert result == cs1 + + +def test_repr_method_comprehensive(fully_connected_two_point_graph): + """Test TransformationManager.__repr__ method comprehensively.""" + + # Test empty TransformationManager + tm_empty = TransformationManager() + repr_empty = repr(tm_empty) + assert "TransformationManager" in repr_empty + assert "coordinate_systems=[]" in repr_empty + assert "coordinate_transforms=[]" in repr_empty + assert "elements=[]" in repr_empty + + # Test TransformationManager with data + tm = TransformationManager() + [cs1, cs2], [transform] = fully_connected_two_point_graph + tm.add_coordinate_system(cs1) + tm.add_coordinate_system(cs2) + + tm.add_transformation(cs1, cs2, transform) + tm.add_element("image1", cs1) + tm.add_element("image2", cs2) + + repr_str = repr(tm) + assert "TransformationManager" in repr_str + assert "cs1" in repr_str + assert "cs2" in repr_str + assert repr_str.find(repr(transform)) + assert "image1" in repr_str + assert "image2" in repr_str diff --git a/tests/io/test_ngff_06.py b/tests/io/test_ngff_06.py new file mode 100644 index 000000000..492e4a51f --- /dev/null +++ b/tests/io/test_ngff_06.py @@ -0,0 +1,384 @@ +# pyright: standard + +from __future__ import annotations + +from pathlib import PurePosixPath +from typing import Iterator, assert_type + +import numpy as np +import xarray as xr +import ome_zarr as oz +import ome_zarr_models.v06.coordinate_transforms as ozm06trans +import ome_zarr_models.v06.multiscales as ozm06ms + +from spatialdata._io.io_raster import ElementKey, FutureSpatialData, try_parse_ngff06_multiscale, try_parse_ngff06_scene +from spatialdata.transformations.graph.edge import IdentityEdge, ScaleEdge +from spatialdata.transformations.graph.vert import NgffCsPath +from tests.conftest import SEED + + +def test_parse_multiscale(): + full_shape = (256, 256) + rng = np.random.default_rng(seed=SEED) + + data = rng.random(full_shape) + omz_img = oz.OMEZarrImage(data=data, axes="yx") + + labels_name = "my_labels" + label_data = oz.OMEZarrImage( + data=rng.integers(0, 256, size=(256, 256), dtype=np.uint8), axes="yx", name=labels_name + ) + omz_label = oz.OMEZarrLabels( + image=label_data, + ) + + scale_factors = (2, 4, 8, 16) + ms = oz.OMEZarrMultiscale( + image=omz_img, + scale_factors=scale_factors, + default_coordinate_system_name="intrinsic_cs", + coordinate_systems=[ + ozm06trans.CoordinateSystem( + name="another_cs", + axes=( + ozm06trans.Axis(name="i", type="space"), + ozm06trans.Axis(name="j", type="space"), + ), + ) + ], + coordinate_transformations=( + ozm06trans.Scale( + scale=(10.0, 20.0), + input=ozm06trans.CoordinateSystemIdentifier(name="intrinsic_cs"), + output=ozm06trans.CoordinateSystemIdentifier(name="another_cs"), + ), + ), + # default_coordinate_system_name="physical", + labels=[omz_label], + ) + + sdata = try_parse_ngff06_multiscale(group_path=PurePosixPath("."), multiscale=ms) + + base_path = PurePosixPath(".") + intrinsic_cs = sdata.find_node( + ngff_path=NgffCsPath(group_path=base_path, name="intrinsic_cs") + ) + another_cs = sdata.find_node( + ngff_path=NgffCsPath(group_path=base_path, name="another_cs") + ) + transfs = list(sdata.edges().directed(fro=intrinsic_cs, to=another_cs)) + assert len(transfs) == 1 + scale = transfs.pop().data["transform"] + assert isinstance(scale, ScaleEdge) + assert scale.input == intrinsic_cs.coord_system + assert scale.output == another_cs.coord_system + np.testing.assert_equal(scale.scale, [10.0, 20.0]) + + _ = sdata.find_node( + ngff_path=NgffCsPath(group_path=base_path / "labels" / labels_name, name="physical") + ) + + # labels_data_tree = data_trees[base_path / "labels" / labels_name] + image_data_tree = sdata.get_image(ElementKey(namespace=base_path, name=omz_img.name)) + assert image_data_tree is not None + for xr_scale_node, ngff_scale, ngff_meta in zip( + image_data_tree.children.values(), ms.images, ms.metadata.datasets, strict=True + ): + xr_scale = xr_scale_node["image"] + assert not isinstance(xr_scale, xr.DataTree) + assert xr_scale.shape == ngff_scale.data.shape + + coords_x = xr_scale.coords["x"] + coords_y = xr_scale.coords["y"] + + seq = ngff_meta.coordinateTransformations[0] + assert isinstance(seq, ozm06trans.Sequence) + scale = seq.transformations[0] + assert isinstance(scale, ozm06trans.Scale) + translate = seq.transformations[1] + assert isinstance(translate, ozm06trans.Translation) + + start_indices = (0, 0) + start = (coords_x[0], coords_y[0]) + + end_indices = (xr_scale.shape[0] - 1, xr_scale.shape[1] - 1) + end = (coords_x[-1], coords_y[-1]) + + for indices, point in [(start_indices, start), (end_indices, end)]: + expected = np.asarray(scale.scale) * indices + translate.translation + assert np.allclose(expected, point) + + sliced_data_tree = image_data_tree.sel(x=slice(0, 256), y=slice(0, 256), method="nearest") + assert sliced_data_tree.equals(image_data_tree) + + quarter_tree = image_data_tree.sel(x=slice(0, 128), y=slice(0, 128), method="nearest") + for xr_scale_node, scale_factor in zip(quarter_tree.children.values(), [1, *scale_factors], strict=True): + xr_scale = xr_scale_node["image"] + expected_shape = np.asarray(full_shape) / 2 / scale_factor + np.testing.assert_equal(xr_scale.shape, expected_shape) + +def test_parse_scene(): + full_shape = (256, 256) + rng = np.random.default_rng(seed=SEED) + + labels_name = "my_labels" + omz_label = oz.OMEZarrLabels( + image=oz.OMEZarrImage( + data=rng.integers(0, 256, size=full_shape, dtype=np.uint8), axes="yx", name=labels_name + ), + ) + + omz_img1 = oz.OMEZarrImage(data=rng.random(full_shape), axes="yx") + multiscales1 = oz.OMEZarrMultiscale( + image=omz_img1, + scale_factors=(2, 4), + default_coordinate_system_name="intrinsic_cs1", + labels=[omz_label], + ) + + omz_img2 = oz.OMEZarrImage(data=rng.random((1, *full_shape)), axes="cyx") + multiscales2 = oz.OMEZarrMultiscale( + image=omz_img2, + scale_factors=[{"x": s, "y": s} for s in (2, 4)], + default_coordinate_system_name="intrinsic_cs2", + coordinate_systems=[ + ozm06trans.CoordinateSystem( + name="cs_in_multiscales2", + axes=( + ozm06trans.Axis(name="c", type="channel"), + ozm06trans.Axis(name="i", type="space"), + ozm06trans.Axis(name="j", type="space"), + ), + ) + ], + coordinate_transformations=( + ozm06trans.Scale( + scale=(1.0, 10.0, 20.0), + input=ozm06trans.CoordinateSystemIdentifier(name="intrinsic_cs2"), + output=ozm06trans.CoordinateSystemIdentifier(name="cs_in_multiscales2"), + ), + ), + ) + + root_cs_2d = ozm06trans.CoordinateSystem( + name="root_cs_2d", + axes=( + ozm06trans.Axis(name="u", type="space"), + ozm06trans.Axis(name="v", type="space"), + ), + ) + root_cs_3d = ozm06trans.CoordinateSystem( + name="root_cs_3d", + axes=( + ozm06trans.Axis(name="axis_0", type="channel"), + ozm06trans.Axis(name="axis_1", type="space"), + ozm06trans.Axis(name="axis_2", type="space"), + ), + ) + scene = oz.OMEZarrScene( + coordinate_systems=[root_cs_2d, root_cs_3d], + coordinate_transformations=[ + ozm06trans.Scale( + scale=(3.0, 4.0), + input=ozm06trans.CoordinateSystemIdentifier(name=root_cs_2d.name), + output=ozm06trans.CoordinateSystemIdentifier(path="multiscales1", name="intrinsic_cs1"), + ), + ozm06trans.Identity( + input=ozm06trans.CoordinateSystemIdentifier(name=root_cs_3d.name), + output=ozm06trans.CoordinateSystemIdentifier(path="multiscales2", name="cs_in_multiscales2"), + ), + ], + images={"multiscales1": multiscales1, "multiscales2": multiscales2}, + ) + + base_path = PurePosixPath(".") + sdata = try_parse_ngff06_scene(group_path=base_path, scene=scene) + + ms1_path = base_path / "multiscales1" + ms2_path = base_path / "multiscales2" + + # scene-level coordinate systems + root_2d = sdata.find_node(ngff_path=NgffCsPath(group_path=base_path, name="root_cs_2d")) + root_3d = sdata.find_node(ngff_path=NgffCsPath(group_path=base_path, name="root_cs_3d")) + + # multiscales-level coordinate systems are namespaced under their scene path + intrinsic_cs1 = sdata.find_node(ngff_path=NgffCsPath(group_path=ms1_path, name="intrinsic_cs1")) + intrinsic_cs2 = sdata.find_node(ngff_path=NgffCsPath(group_path=ms2_path, name="intrinsic_cs2")) + cs_in_ms2 = sdata.find_node(ngff_path=NgffCsPath(group_path=ms2_path, name="cs_in_multiscales2")) + label_cs = sdata.find_node( + ngff_path=NgffCsPath(group_path=ms1_path / "labels" / labels_name, name="physical") + ) + + # scene-level transforms point into the multiscales + root_to_ms1 = sdata.edges().directed(fro=root_2d, to=intrinsic_cs1).collect(list) + assert len(root_to_ms1) == 1 + scale = root_to_ms1[0].data["transform"] + assert isinstance(scale, ScaleEdge) + np.testing.assert_equal(scale.scale, [3.0, 4.0]) + + root_to_ms2 = sdata.edges().directed(fro=root_3d, to=cs_in_ms2).collect(list) + assert len(root_to_ms2) == 1 + assert isinstance(root_to_ms2[0].data["transform"], IdentityEdge) + + # multiscales-level transforms survive being nested in a scene + ms2_transfs = list(sdata.edges().directed(fro=intrinsic_cs2, to=cs_in_ms2)) + assert len(ms2_transfs) == 1 + ms2_scale = ms2_transfs[0].data["transform"] + assert isinstance(ms2_scale, ScaleEdge) + np.testing.assert_equal(ms2_scale.scale, [1.0, 10.0, 20.0]) + + # images and labels are keyed by their scene path and attached to their intrinsic CSs + img1_key = ElementKey(namespace=ms1_path, name=omz_img1.name) + img2_key = ElementKey(namespace=ms2_path, name=omz_img2.name) + label_key = ElementKey(namespace=ms1_path / "labels" / labels_name, name=labels_name) + assert sdata.get_intrinsic(img1_key) == intrinsic_cs1 + assert sdata.get_intrinsic(img2_key) == intrinsic_cs2 + assert sdata.get_intrinsic(label_key) == label_cs + + for key, ms in [(img1_key, multiscales1), (img2_key, multiscales2)]: + data_tree = sdata.get_image(key) + assert data_tree is not None + for xr_scale_node, ngff_scale in zip(data_tree.children.values(), ms.images, strict=True): + xr_scale = xr_scale_node["image"] + assert not isinstance(xr_scale, xr.DataTree) + assert xr_scale.shape == ngff_scale.data.shape + + label_tree = sdata.get_label(label_key) + assert label_tree is not None + assert label_tree["scale0"]["image"].shape == full_shape + + +def test_dumping_multiscales_to_ngff(): + full_shape = (256, 256) + rng = np.random.default_rng(seed=SEED) + + img_data1 = rng.random(full_shape) + omz_img1 = oz.OMEZarrImage(data=img_data1, axes="yx") + + labels_name1 = "my_labels" + label_data1 = oz.OMEZarrImage( + data=rng.integers(0, 256, size=(256, 256), dtype=np.uint8), axes="yx", name=labels_name1 + ) + omz_label1 = oz.OMEZarrLabels( + image=label_data1, + ) + + scale_factors = (2, 4, 8, 16) + multiscales1 = oz.OMEZarrMultiscale( + image=omz_img1, + scale_factors=scale_factors, + default_coordinate_system_name="intrinsic_cs1", + coordinate_systems=[ + ozm06trans.CoordinateSystem( + name="cs_in_multiscales1", + axes=( + ozm06trans.Axis(name="i", type="space"), + ozm06trans.Axis(name="j", type="space"), + ), + ) + ], + coordinate_transformations=( + ozm06trans.Scale( + scale=(10.0, 20.0), + input=ozm06trans.CoordinateSystemIdentifier(name="intrinsic_cs1"), + output=ozm06trans.CoordinateSystemIdentifier(name="cs_in_multiscales1"), + ), + ), + # default_coordinate_system_name="physical", + labels=[omz_label1], + ) + + + img_data2 = rng.random((1, 256, 256)) + omz_img2 = oz.OMEZarrImage(data=img_data2, axes="cyx") + multiscales2 = oz.OMEZarrMultiscale( + image=omz_img2, + scale_factors=[ + {"x": s, "y": s} for s in scale_factors + ], + default_coordinate_system_name="intrinsic_cs2", + coordinate_systems=[ + ozm06trans.CoordinateSystem( + name="cs_in_multiscales2", + axes=( + ozm06trans.Axis(name="c", type="channel"), + ozm06trans.Axis(name="i", type="space"), + ozm06trans.Axis(name="j", type="space"), + ), + ) + ], + coordinate_transformations=( + ozm06trans.Scale( + scale=(10.0, 20.0, 1.0), + input=ozm06trans.CoordinateSystemIdentifier(name="intrinsic_cs2"), + output=ozm06trans.CoordinateSystemIdentifier(name="cs_in_multiscales2"), + ), + ), + labels=[omz_label1], + ) + + root_cs = ozm06trans.CoordinateSystem( + name="root_cs", + axes=( + ozm06trans.Axis(name="axis_0", type="space"), + ozm06trans.Axis(name="axis_1", type="space"), + ozm06trans.Axis(name="axis_2", type="channel"), + ), + ) + ome_identity = ozm06trans.Identity( + input=ozm06trans.CoordinateSystemIdentifier(name=root_cs.name), + output=ozm06trans.CoordinateSystemIdentifier(path="multiscales2", name="cs_in_multiscales2"), + ) + scene = oz.OMEZarrScene( + coordinate_systems=[root_cs], + coordinate_transformations=[ome_identity], + images={"multiscales1": multiscales1, "multiscales2": multiscales2}, + ) + + sdata = try_parse_ngff06_scene(group_path=PurePosixPath("."), scene=scene) + + dumped_scene = sdata.default_scene_dumper().dump() + + assert(set(dumped_scene.coordinate_systems) == set(scene.coordinate_systems)) + + for ct in dumped_scene.coordinate_transformations: + assert ct in scene.coordinate_transformations + # assert set(dumped_scene.coordinate_transformations) == set(scene.coordinate_transformations) + + assert dumped_scene.images.keys() == scene.images.keys() + for ms_path, original_ms in scene.images.items(): + dumped_ms = dumped_scene.images[ms_path] + for dumped_img, original_img in zip(dumped_ms.images, original_ms.images, strict=True): + assert dumped_img.data.shape == original_img.data.shape + + for (dumped_scale, dumped_translation), (original_scale, original_translation) in zip( + _get_scales_and_translations(dumped_ms.metadata), + _get_scales_and_translations(original_ms.metadata), + strict=True, + ): + np.testing.assert_allclose(dumped_scale.scale, original_scale.scale) + np.testing.assert_allclose(dumped_translation.translation, original_translation.translation) + + +def _get_scales_and_translations( + multiscale: ozm06ms.Multiscale, +) -> Iterator[tuple[ozm06trans.Scale, ozm06trans.Translation]]: + """Returns each dataset's scale and translation, defaulting to trivial ones if missing.""" + scale_noop = ozm06trans.Scale(scale=(1.0,) * multiscale.ndim) + translation_noop = ozm06trans.Translation(translation=(0.0,) * multiscale.ndim) + + for dataset in multiscale.datasets: + transform = dataset.coordinateTransformations[0] + if isinstance(transform, ozm06trans.Scale): + yield transform, translation_noop + continue + if isinstance(transform, ozm06trans.Identity): + yield scale_noop, translation_noop + continue + assert_type(transform, ozm06trans.Sequence) + scale, translate = transform.transformations + assert isinstance(scale, ozm06trans.Scale) + assert isinstance(translate, ozm06trans.Translation) + yield scale, translate + + diff --git a/tests/transformations/edge/__init__.py b/tests/transformations/edge/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/transformations/edge/conftest.py b/tests/transformations/edge/conftest.py new file mode 100644 index 000000000..5ff56a219 --- /dev/null +++ b/tests/transformations/edge/conftest.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from spatialdata.transformations.graph.vert import Axis, CoordSystem + +x_axis = Axis(name="x", type="space", unit="micrometer") +y_axis = Axis(name="y", type="space", unit="micrometer") +z_axis = Axis(name="z", type="space", unit="micrometer") +c_axis = Axis(name="c", type="channel") + +a_axis = Axis(name="a", type="space", unit="micrometer") +b_axis = Axis(name="b", type="space", unit="micrometer") +c_axis = Axis(name="c", type="space", unit="micrometer") + +x_cs = CoordSystem(name="x", axes=(x_axis,), virtual=False) +y_cs = CoordSystem(name="y", axes=(y_axis,), virtual=False) +z_cs = CoordSystem(name="z", axes=(z_axis,), virtual=False) +c_cs = CoordSystem(name="c", axes=(c_axis,), virtual=False) +xy_cs = CoordSystem( + name="xy", + axes=( + x_axis, + y_axis, + ), + virtual=False, +) +yx_cs = CoordSystem( + name="yx", + axes=( + y_axis, + x_axis, + ), + virtual=False, +) +xyz_cs = CoordSystem( + name="xyz", + axes=( + x_axis, + y_axis, + z_axis, + ), + virtual=False, +) +zyx_cs = CoordSystem( + name="zyx", + axes=( + z_axis, + y_axis, + x_axis, + ), + virtual=False, +) +xyc_cs = CoordSystem( + name="xyc", + axes=( + x_axis, + y_axis, + c_axis, + ), + virtual=False, +) +cyx_cs = CoordSystem( + name="cyx", + axes=( + c_axis, + y_axis, + x_axis, + ), + virtual=False, +) + +abc_cs = CoordSystem(name="abc", axes=(a_axis, b_axis, c_axis), virtual=False) +cba_cs = CoordSystem( + name="cba", + axes=( + c_axis, + b_axis, + a_axis, + ), + virtual=False, +) +cba_cs = CoordSystem(name="cba", axes=(c_axis, b_axis, a_axis), virtual=False) diff --git a/tests/transformations/edge/test_edge_parsing.py b/tests/transformations/edge/test_edge_parsing.py new file mode 100644 index 000000000..a20acbaa3 --- /dev/null +++ b/tests/transformations/edge/test_edge_parsing.py @@ -0,0 +1,226 @@ +# pyright: strict + +from __future__ import annotations + +import numpy as np +import ome_zarr_models.v06.coordinate_transforms as ozm06trans + +from spatialdata.transformations.graph.edge import ( + AffineEdge, + ByDimensionEdge, + CsGen, + MapAxisEdge, + RotationEdge, + ScaleEdge, + SequenceEdge, + TranslationEdge, + parse_affine, + parse_by_dimension, + parse_identity, + parse_map_axis, + parse_project_axis, + parse_rotation, + parse_scale, + parse_sequence, + parse_translation, +) +from tests.transformations.edge.conftest import ( + xy_cs, + xyz_cs, + yx_cs, + z_axis, +) + +# The ome-zarr-models transform models used below all leave `input`/`output` at their +# default of None: those fields identify coordinate systems by name/path in the NGFF +# metadata and aren't consumed by the parse_* functions, which instead take the already +# resolved `CoordSystem` (or a `CsGen`) as separate arguments. Only the transform-specific +# fields (e.g. `translation`, `scale`, `name`) are exercised here. + +POINTS_2D = np.array([[1.0, 2.0], [3.0, 4.0]]) + + +class TestParseIdentity: + def test_parses_name(self): + model = ozm06trans.Identity(name="my-identity") + edge = parse_identity(model, input=xy_cs, out=xy_cs) + assert edge.name == "my-identity" + assert edge.input is xy_cs + assert edge.output is xy_cs + + def test_csgen_output_mirrors_input(self): + model = ozm06trans.Identity(name="my-identity") + edge = parse_identity(model, input=xy_cs, out=CsGen(base_name="gen")) + assert edge.output.axes_names == xy_cs.axes_names + assert edge.output.virtual is True + + +class TestParseTranslation: + def test_parses_translation_vector(self): + model = ozm06trans.Translation(name="my-translation", translation=(1.0, 2.0)) + edge = parse_translation(model, input=xy_cs, out=xy_cs) + np.testing.assert_allclose(edge.translation, [1.0, 2.0]) + assert edge.name == model.name + assert edge.input is xy_cs + assert edge.output is xy_cs + + def test_csgen_output_mirrors_input(self): + model = ozm06trans.Translation(translation=(1.0, 2.0)) + edge = parse_translation(model, input=xy_cs, out=CsGen(base_name="gen")) + assert edge.output.axes_names == xy_cs.axes_names + assert edge.output.virtual is True + + +class TestParseProjectAxis: + def test_parses_dropped_inputs(self): + model = ozm06trans.ProjectAxis(name="drop-z", droppedInputs=(2,)) + edge = parse_project_axis(model, input=xyz_cs, output=xy_cs) + assert edge.dropped_inputs == {z_axis} + assert edge.created_outputs == set() + assert edge.name == "drop-z" + + def test_parses_created_outputs(self): + model = ozm06trans.ProjectAxis(name="create-z", createdOutputs=(2,)) + edge = parse_project_axis(model, input=xy_cs, output=xyz_cs) + assert edge.created_outputs == {z_axis} + assert edge.dropped_inputs == set() + + def test_csgen_output_computes_dimensionality(self): + model = ozm06trans.ProjectAxis(droppedInputs=(2,)) + edge = parse_project_axis(model, input=xyz_cs, output=CsGen(base_name="gen")) + assert edge.output.num_axes == 2 + assert edge.dropped_inputs == {z_axis} + + +class TestParseScale: + def test_parses_scale_vector_and_name(self): + model = ozm06trans.Scale(name="my-scale", scale=(2.0, 3.0)) + edge = parse_scale(model, input=xy_cs, out=xy_cs) + assert isinstance(edge, ScaleEdge) + np.testing.assert_allclose(edge.scale, [2.0, 3.0]) + assert edge.name == "my-scale" + + def test_csgen_output_mirrors_input(self): + model = ozm06trans.Scale(scale=(2.0, 3.0)) + edge = parse_scale(model, input=xy_cs, out=CsGen(base_name="gen")) + assert edge.output.axes_names == xy_cs.axes_names + assert edge.output.virtual is True + + +class TestParseMapAxis: + def test_parses_with_explicit_output(self): + model = ozm06trans.MapAxis(name="my-map-axis", mapAxis=(1, 0)) + edge = parse_map_axis(model, input=xy_cs, out=yx_cs) + assert isinstance(edge, MapAxisEdge) + assert edge.name == "my-map-axis" + assert edge.output is yx_cs + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[2.0, 1.0], [4.0, 3.0]])) + + def test_csgen_output_is_permuted_input_axes(self): + model = ozm06trans.MapAxis(mapAxis=(1, 0)) + edge = parse_map_axis(model, input=xy_cs, out=CsGen(base_name="gen")) + assert edge.output.axes_names == yx_cs.axes_names + assert edge.output.virtual is True + + +class TestParseAffine: + def test_parses_linear_translation_and_name(self): + model = ozm06trans.Affine(name="my-affine", affine=((2.0, 0.0, 1.0), (0.0, 3.0, 5.0))) + edge = parse_affine(model, input=xy_cs, output=xy_cs) + assert isinstance(edge, AffineEdge) + np.testing.assert_allclose(edge.linear, [[2.0, 0.0], [0.0, 3.0]]) + np.testing.assert_allclose(edge.translation, [1.0, 5.0]) + assert edge.name == "my-affine" + + def test_csgen_output_derived_from_matrix_row_count(self): + model = ozm06trans.Affine(affine=((2.0, 0.0, 1.0), (0.0, 3.0, 5.0))) + edge = parse_affine(model, input=xy_cs, output=CsGen(base_name="gen")) + assert edge.output.num_axes == 2 + assert edge.output.virtual is True + + +class TestParseRotation: + def test_parses_rotation_matrix_and_name(self): + model = ozm06trans.Rotation(name="my-rotation", rotation=((0.0, -1.0), (1.0, 0.0))) + edge = parse_rotation(model, input=xy_cs, out=xy_cs) + assert isinstance(edge, RotationEdge) + np.testing.assert_allclose(edge.rotation, [[0.0, -1.0], [1.0, 0.0]]) + assert edge.name == "my-rotation" + + def test_csgen_output_derived_from_matrix_row_count(self): + model = ozm06trans.Rotation(rotation=((0.0, -1.0), (1.0, 0.0))) + edge = parse_rotation(model, input=xy_cs, out=CsGen(base_name="gen")) + assert edge.output.num_axes == 2 + assert edge.output.virtual is True + + +class TestParseSequence: + def test_parses_name_and_composes_inner_transformations_in_order(self): + inner_translation = ozm06trans.Translation(translation=(1.0, 2.0)) + inner_scale = ozm06trans.Scale(scale=(2.0, 3.0)) + model = ozm06trans.Sequence(name="my-sequence", transformations=(inner_translation, inner_scale)) + edge = parse_sequence(model, input=xy_cs, output=xy_cs) + assert isinstance(edge, SequenceEdge) + assert edge.name == "my-sequence" + assert edge.input is xy_cs + assert edge.output is xy_cs + assert [type(t) for t in edge.transformations] == [TranslationEdge, ScaleEdge] + expected = (POINTS_2D + np.array([1.0, 2.0])) * np.array([2.0, 3.0]) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), expected) + + def test_single_inner_transformation_uses_given_output_directly(self): + inner_translation = ozm06trans.Translation(translation=(1.0, 2.0)) + model = ozm06trans.Sequence(transformations=(inner_translation,)) + edge = parse_sequence(model, input=xy_cs, output=xy_cs) + assert len(edge.transformations) == 1 + assert edge.transformations[0].output is xy_cs + + def test_multiple_inner_transformations_use_virtual_intermediate_coord_systems(self): + inner_translation = ozm06trans.Translation(translation=(1.0, 2.0)) + inner_scale = ozm06trans.Scale(scale=(2.0, 3.0)) + model = ozm06trans.Sequence(transformations=(inner_translation, inner_scale)) + edge = parse_sequence(model, input=xy_cs, output=xy_cs) + intermediate = edge.transformations[0].output + assert intermediate is not xy_cs + assert intermediate.virtual is True + assert intermediate.axes_names == xy_cs.axes_names + assert edge.transformations[1].output is xy_cs + + def test_csgen_output_mirrors_input(self): + inner_translation = ozm06trans.Translation(translation=(1.0, 2.0)) + inner_scale = ozm06trans.Scale(scale=(2.0, 3.0)) + model = ozm06trans.Sequence(transformations=(inner_translation, inner_scale)) + edge = parse_sequence(model, input=xy_cs, output=CsGen(base_name="gen")) + assert edge.output.axes_names == xy_cs.axes_names + assert edge.output.virtual is True + + +class TestParseByDimension: + def test_parses_name_and_splits_axes_between_inner_transformations(self): + scale_x = ozm06trans.Scale(scale=(2.0,)) + translate_y = ozm06trans.Translation(translation=(5.0,)) + model = ozm06trans.ByDimension( + name="my-by-dim", + transformations=( + ozm06trans.ByDimensionTransform(transformation=scale_x, inputAxes=(0,), outputAxes=(0,)), + ozm06trans.ByDimensionTransform(transformation=translate_y, inputAxes=(1,), outputAxes=(1,)), + ), + ) + edge = parse_by_dimension(model, input=xy_cs, output=xy_cs) + assert isinstance(edge, ByDimensionEdge) + assert edge.name == "my-by-dim" + assert [type(t) for t in edge.transformations] == [ScaleEdge, TranslationEdge] + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[2.0, 7.0], [6.0, 9.0]])) + + def test_csgen_output_derived_from_max_output_axis_index(self): + scale_x = ozm06trans.Scale(scale=(2.0,)) + translate_y = ozm06trans.Translation(translation=(5.0,)) + model = ozm06trans.ByDimension( + transformations=( + ozm06trans.ByDimensionTransform(transformation=scale_x, inputAxes=(0,), outputAxes=(0,)), + ozm06trans.ByDimensionTransform(transformation=translate_y, inputAxes=(1,), outputAxes=(1,)), + ), + ) + edge = parse_by_dimension(model, input=xy_cs, output=CsGen(base_name="gen")) + assert edge.output.num_axes == 2 + assert edge.output.virtual is True diff --git a/tests/transformations/edge/test_edge_transformations.py b/tests/transformations/edge/test_edge_transformations.py new file mode 100644 index 000000000..2427da620 --- /dev/null +++ b/tests/transformations/edge/test_edge_transformations.py @@ -0,0 +1,374 @@ +# pyright: strict + +from __future__ import annotations + +import numpy as np +import pytest + +from spatialdata._core.transformation_manager.exceptions import ( + AxisRedefinitionError, + DeterminantDifferentFromOne, + EmptyTransformSequenceError, + IncompatibleCoordSystemsError, + MissingAxisError, + NotOrthonormalError, + UnexpectedShapeError, + UnmappedAxisError, +) +from spatialdata.transformations.graph.edge import ( + AffineEdge, + ByDimensionEdge, + IdentityEdge, + MapAxisEdge, + ProjectAxisEdge, + RotationEdge, + ScaleEdge, + SequenceEdge, + TranslationEdge, +) +from tests.transformations.edge.conftest import ( + a_axis, + abc_cs, + b_axis, + c_axis, + x_axis, + x_cs, + xy_cs, + xyc_cs, + xyz_cs, + y_axis, + y_cs, + yx_cs, + z_axis, + zyx_cs, +) + +POINTS_2D = np.array([[1.0, 2.0], [3.0, 4.0]]) +POINTS_3D = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + + +class TestAffineEdge: + def test_constructor_rejects_wrong_linear_shape(self): + with pytest.raises(UnexpectedShapeError): + AffineEdge(linear=np.eye(3), input=xy_cs, output=xy_cs) + + def test_constructor_rejects_wrong_translation_shape(self): + with pytest.raises(UnexpectedShapeError): + AffineEdge(linear=np.eye(2), translation=np.zeros(3), input=xy_cs, output=xy_cs) + + def test_from_affine_matrix(self): + # fmt: off + affine_matrix = np.array([ + [2.0, 0.0, 1.0], + [0.0, 3.0, 5.0], + [0.0, 0.0, 1.0] + ]) + linear = [ + [2.0, 0.0], + [0.0, 3.0], + ] + translation = [ + 1.0, + 5.0 + ] + # fmt: on + edge = AffineEdge.from_affine_matrix(name=None, affine_matrix=affine_matrix, input=xy_cs, output=xy_cs) + np.testing.assert_equal(edge.linear, linear) + np.testing.assert_equal(edge.translation, translation) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[3.0, 11.0], [7.0, 17.0]])) + + def test_transform_points_scale_and_translate(self): + edge = AffineEdge( + linear=np.array([[2.0, 0.0], [0.0, 3.0]]), translation=np.array([1.0, 1.0]), input=xy_cs, output=xy_cs + ) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[3.0, 7.0], [7.0, 13.0]])) + + def test_transform_points_2d_to_3d(self): + linear = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) + edge = AffineEdge(linear=linear, input=xy_cs, output=xyz_cs) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[1.0, 2.0, 3.0], [3.0, 4.0, 7.0]])) + + def test_transform_points_rejects_wrong_shape(self): + edge = AffineEdge(linear=np.eye(2), input=xy_cs, output=xy_cs) + with pytest.raises(UnexpectedShapeError): + # expecting (n, 2), gets (2,) + edge.transform_points(np.asarray([1.0, 2.0])) + with pytest.raises(UnexpectedShapeError): + # expecting (n, 2), gets (n, 3) + edge.transform_points(POINTS_3D) + + def test_inverse_roundtrips(self): + edge = AffineEdge( + linear=np.array([[2.0, 0.0], [0.0, 4.0]]), translation=np.array([1.0, -1.0]), input=xy_cs, output=xy_cs + ) + inv = edge.inverse() + assert inv is not None + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D) + + def test_inverse_returns_none_for_singular_matrix(self): + edge = AffineEdge(linear=np.zeros((2, 2)), input=xy_cs, output=xy_cs) + assert edge.inverse() is None + + +class TestIdentityEdge: + def test_constructor_rejects_mismatched_number_of_axes(self): + with pytest.raises(IncompatibleCoordSystemsError): + IdentityEdge(None, input=xy_cs, output=xyz_cs) + + def test_transform_points_is_a_noop(self): + edge = IdentityEdge(None, input=xy_cs, output=xy_cs) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), POINTS_2D) + + def test_to_affine_matches_transform_points(self): + edge = IdentityEdge(None, input=xy_cs, output=xy_cs) + affine = edge.to_affine() + np.testing.assert_allclose(affine.linear, np.eye(2)) + np.testing.assert_allclose(affine.transform_points(POINTS_2D), POINTS_2D) + + def test_inverse_is_still_the_identity(self): + edge = IdentityEdge(None, input=xy_cs, output=xy_cs) + inv = edge.inverse() + np.testing.assert_allclose(inv.transform_points(POINTS_2D), POINTS_2D) + + +class TestMapAxisEdge: + def test_transform_points_swaps_axes(self): + edge = MapAxisEdge( + input=xyz_cs, + output=abc_cs, + input_to_output={x_axis: b_axis, y_axis: c_axis, z_axis: a_axis}, + ) + np.testing.assert_equal(edge.transform_points(POINTS_3D), np.array([[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]])) + + def test_transform_points_permutation_of_three_axes(self): + edge = MapAxisEdge( + input=xyz_cs, output=abc_cs, input_to_output={x_axis: c_axis, y_axis: b_axis, z_axis: a_axis} + ) + np.testing.assert_allclose(edge.transform_points(POINTS_3D), np.array([[3.0, 2.0, 1.0], [6.0, 5.0, 4.0]])) + + def test_to_affine_matches_transform_points(self): + edge = MapAxisEdge(input=xy_cs, output=xy_cs, input_to_output={x_axis: y_axis, y_axis: x_axis}) + affine = edge.to_affine() + np.testing.assert_allclose(affine.transform_points(POINTS_2D), edge.transform_points(POINTS_2D)) + + def test_inverse_roundtrips(self): + edge = MapAxisEdge( + input=xyz_cs, output=zyx_cs, input_to_output={x_axis: z_axis, y_axis: y_axis, z_axis: x_axis} + ) + inv = edge.inverse() + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_3D)), POINTS_3D) + + +class TestProjectAxisEdge: + def test_constructor_rejects_dropped_axis_missing_from_input(self): + with pytest.raises(MissingAxisError): + ProjectAxisEdge(input=xy_cs, output=xy_cs, dropped_inputs={z_axis}, created_outputs=set()) + + def test_constructor_rejects_created_axis_missing_from_output(self): + with pytest.raises(MissingAxisError): + ProjectAxisEdge(input=xy_cs, output=xy_cs, dropped_inputs=set(), created_outputs={z_axis}) + + def test_general_coord_system_incompatibility(self): + with pytest.raises(IncompatibleCoordSystemsError): + ProjectAxisEdge(input=xy_cs, output=xyz_cs, dropped_inputs=set(), created_outputs=set()) + with pytest.raises(IncompatibleCoordSystemsError): + ProjectAxisEdge(input=xyz_cs, output=xy_cs, dropped_inputs=set(), created_outputs=set()) + + def test_to_affine_for_identity_case(self): + edge = ProjectAxisEdge(input=xy_cs, output=xy_cs, dropped_inputs=set(), created_outputs=set()) + affine = edge.to_affine() + np.testing.assert_allclose(affine.linear, np.eye(2)) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), POINTS_2D) + + def test_dropping_one_axis(self): + edge = ProjectAxisEdge(input=xyz_cs, output=xy_cs, dropped_inputs={z_axis}, created_outputs=set()) + affine = edge.to_affine() + assert affine is not None + np.testing.assert_allclose(edge.transform_points(POINTS_3D), np.array([[1.0, 2.0], [4.0, 5.0]])) + + def test_creating_one_axis(self): + edge = ProjectAxisEdge(input=xy_cs, output=xyz_cs, dropped_inputs=set(), created_outputs={z_axis}) + affine = edge.to_affine() + assert affine is not None + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[1.0, 2.0, 0.0], [3.0, 4.0, 0.0]])) + + def test_inverse_none_when_axes_are_dropped_or_created(self): + dropped = ProjectAxisEdge(input=xyz_cs, output=xy_cs, dropped_inputs={z_axis}, created_outputs=set()) + assert dropped.inverse() is None + created = ProjectAxisEdge(input=xy_cs, output=xyz_cs, dropped_inputs=set(), created_outputs={z_axis}) + assert created.inverse() is None + + def test_inverse_roundtrip(self): + edge = ProjectAxisEdge(input=xy_cs, output=yx_cs, dropped_inputs=set(), created_outputs=set()) + inv = edge.inverse() + assert inv is not None + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D) + + +class TestTranslationEdge: + def test_constructor_rejects_mismatched_number_of_axes(self): + with pytest.raises(IncompatibleCoordSystemsError): + TranslationEdge(translation=np.array([1.0, 2.0]), input=xy_cs, output=xyz_cs) + + def test_transform_points_adds_translation(self): + edge = TranslationEdge(translation=np.array([10.0, 20.0]), input=xy_cs, output=xy_cs) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[11.0, 22.0], [13.0, 24.0]])) + + def test_to_affine_matches_transform_points(self): + edge = TranslationEdge(translation=np.array([10.0, 20.0]), input=xy_cs, output=xy_cs) + affine = edge.to_affine() + assert affine is not None + np.testing.assert_allclose(affine.transform_points(POINTS_2D), edge.transform_points(POINTS_2D)) + + def test_inverse_roundtrips(self): + edge = TranslationEdge(translation=np.array([10.0, 20.0]), input=xy_cs, output=xy_cs) + inv = edge.inverse() + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D) + + +class TestScaleEdge: + def test_constructor_rejects_wrong_scale_shape(self): + with pytest.raises(UnexpectedShapeError): + ScaleEdge(scale=np.array([1.0, 2.0, 3.0]), input=xy_cs, output=xy_cs) + + def test_constructor_rejects_mismatched_number_of_axes(self): + with pytest.raises(IncompatibleCoordSystemsError): + ScaleEdge(scale=np.array([1.0, 2.0]), input=xy_cs, output=xyz_cs) + + def test_transform_points(self): + edge = ScaleEdge(scale=np.array([2.0, 4.0]), input=xy_cs, output=xy_cs) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[2.0, 8.0], [6.0, 16.0]])) + + def test_to_affine_matches_transform_points(self): + edge = ScaleEdge(scale=np.array([2.0, 4.0]), input=xy_cs, output=xy_cs) + affine = edge.to_affine() + np.testing.assert_allclose(affine.transform_points(POINTS_2D), edge.transform_points(POINTS_2D)) + + def test_inverse_roundtrips(self): + edge = ScaleEdge(scale=np.array([2.0, 4.0]), input=xy_cs, output=xy_cs) + inv = edge.inverse() + assert inv is not None + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D) + + def test_inverse_returns_none_when_scale_is_zero(self): + edge = ScaleEdge(scale=np.array([0.0, 4.0]), input=xy_cs, output=xy_cs) + assert edge.inverse() is None + + +class TestRotationEdge: + def test_constructor_rejects_mismatched_number_of_axes(self): + with pytest.raises(IncompatibleCoordSystemsError): + RotationEdge(linear_matrix=np.eye(2), input=xy_cs, output=xyz_cs) + + def test_constructor_rejects_wrong_shape(self): + with pytest.raises(UnexpectedShapeError): + RotationEdge(linear_matrix=np.eye(3), input=xy_cs, output=xy_cs) + + def test_constructor_rejects_matrix_with_det_diff_from_1(self): + with pytest.raises(DeterminantDifferentFromOne): + RotationEdge(linear_matrix=np.array([[1.0, 0.0], [0.0, -1.0]]), input=xy_cs, output=xy_cs) + + def test_constructor_rejects_non_orthonormal_matrix(self): + with pytest.raises(NotOrthonormalError): + RotationEdge(linear_matrix=np.array([[1.0, 0.0], [0.0, 2.0]]), input=xy_cs, output=xy_cs) + + def test_transform_points_rotates_90_degrees(self): + edge = RotationEdge(linear_matrix=np.array([[0.0, -1.0], [1.0, 0.0]]), input=xy_cs, output=xy_cs) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[-2.0, 1.0], [-4.0, 3.0]])) + + def test_to_affine_matches_transform_points(self): + edge = RotationEdge(linear_matrix=np.array([[0.0, -1.0], [1.0, 0.0]]), input=xy_cs, output=xy_cs) + affine = edge.to_affine() + np.testing.assert_allclose(affine.transform_points(POINTS_2D), edge.transform_points(POINTS_2D)) + + def test_inverse_roundtrips(self): + edge = RotationEdge(linear_matrix=np.array([[0.0, -1.0], [1.0, 0.0]]), input=xy_cs, output=xy_cs) + inv = edge.inverse() + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D) + + +class TestSequenceEdge: + def test_constructor_rejects_empty_sequence(self): + with pytest.raises(EmptyTransformSequenceError): + SequenceEdge(transformations=[]) + + def test_constructor_rejects_incompatible_neighbors(self): + first = TranslationEdge(translation=np.array([1.0, 2.0]), input=xy_cs, output=xy_cs) + second = TranslationEdge(translation=np.array([1.0, 2.0, 3.0]), input=xyz_cs, output=xyz_cs) + with pytest.raises(IncompatibleCoordSystemsError): + SequenceEdge(transformations=[first, second]) + + def test_transform_points_composes_in_order(self): + translate = TranslationEdge(translation=np.array([1.0, 2.0]), input=xy_cs, output=xy_cs) + scale = ScaleEdge(scale=np.array([3.0, 4.0]), input=xy_cs, output=xy_cs) + edge = SequenceEdge(transformations=[translate, scale]) + expected = (POINTS_2D + np.array([1.0, 2.0])) * np.array([3.0, 4.0]) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), expected) + + def test_to_affine_matches_transform_points(self): + translate = TranslationEdge(translation=np.array([1.0, 2.0]), input=xy_cs, output=xy_cs) + scale = ScaleEdge(scale=np.array([3.0, 4.0]), input=xy_cs, output=xy_cs) + edge = SequenceEdge(transformations=[translate, scale]) + affine = edge.to_affine() + np.testing.assert_allclose(affine.transform_points(POINTS_2D), edge.transform_points(POINTS_2D)) + + def test_inverse_roundtrips(self): + translate = TranslationEdge(translation=np.array([1.0, 2.0]), input=xy_cs, output=xy_cs) + scale = ScaleEdge(scale=np.array([3.0, 4.0]), input=xy_cs, output=xy_cs) + edge = SequenceEdge(transformations=[translate, scale]) + inv = edge.inverse() + assert inv is not None + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D) + + def test_inverse_returns_none_if_any_component_is_not_invertible(self): + scale = ScaleEdge(scale=np.array([0.0, 4.0]), input=xy_cs, output=xy_cs) + translate = TranslationEdge(translation=np.array([1.0, 2.0]), input=xy_cs, output=xy_cs) + edge = SequenceEdge(transformations=[scale, translate]) + assert edge.inverse() is None + + +class TestByDimensionEdge: + def test_constructor_rejects_input_axis_missing_from_overall_input(self): + sub = IdentityEdge(None, input=xyc_cs, output=xyc_cs) + with pytest.raises(MissingAxisError): + ByDimensionEdge(transformations=[sub], input=xy_cs, output=xyc_cs) + + def test_constructor_rejects_output_axis_missing_from_overall_output(self): + sub = IdentityEdge(None, input=xy_cs, output=xy_cs) + with pytest.raises(MissingAxisError): + ByDimensionEdge(transformations=[sub], input=xy_cs, output=x_cs) + + def test_constructor_rejects_output_axis_defined_more_than_once(self): + first = IdentityEdge(None, input=x_cs, output=x_cs) + second = IdentityEdge(None, input=x_cs, output=x_cs) + with pytest.raises(AxisRedefinitionError): + ByDimensionEdge(transformations=[first, second], input=x_cs, output=x_cs) + + def test_constructor_rejects_unmapped_output_axis(self): + sub = IdentityEdge(None, input=x_cs, output=x_cs) + try: + ByDimensionEdge(transformations=[sub], input=xy_cs, output=xy_cs) + raise AssertionError(f"Expected {UnmappedAxisError.__name__} to be raised") + except UnmappedAxisError as e: + assert e.axis == y_axis + + def test_transform_points_applies_each_transformation_per_axis(self): + scale_x = ScaleEdge(scale=np.array([2.0]), input=x_cs, output=x_cs) + translate_y = TranslationEdge(translation=np.array([5.0]), input=y_cs, output=y_cs) + edge = ByDimensionEdge(transformations=[scale_x, translate_y], input=xy_cs, output=xy_cs) + np.testing.assert_allclose(edge.transform_points(POINTS_2D), np.array([[2.0, 7.0], [6.0, 9.0]])) + + def test_to_affine_matches_transform_points(self): + scale_x = ScaleEdge(scale=np.array([2.0]), input=x_cs, output=x_cs) + translate_y = TranslationEdge(translation=np.array([5.0]), input=y_cs, output=y_cs) + edge = ByDimensionEdge(transformations=[scale_x, translate_y], input=xy_cs, output=xy_cs) + affine = edge.to_affine() + np.testing.assert_allclose(affine.transform_points(POINTS_2D), edge.transform_points(POINTS_2D)) + + def test_inverse_roundtrips(self): + scale_x = ScaleEdge(scale=np.array([2.0]), input=x_cs, output=x_cs) + translate_y = TranslationEdge(translation=np.array([5.0]), input=y_cs, output=y_cs) + edge = ByDimensionEdge(transformations=[scale_x, translate_y], input=xy_cs, output=xy_cs) + inv = edge.inverse() + assert inv is not None + np.testing.assert_allclose(inv.transform_points(edge.transform_points(POINTS_2D)), POINTS_2D)