wip code cleanup

Author: LucaMarconato

notebooks/transformations_demo.ipynb

@@ -9,6 +9,8 @@
    "source": [
     "import zarr\n",
     "import numpy as np\n",
+    "\n",
+    "import ngff_transformations.graph\n",
     "import ngff_transformations.transform as ngt\n",
     "\n",
     "from pathlib import Path\n",
@@ -95,9 +97,7 @@
    "id": "72648a07-c8d1-4078-80ef-2b7fb4758451",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "transformation_path, (src_coord_system, tgt_coord_system) = ngt.find_walks_in_graph(nx_graph, 'VOI-01.ome.zarr/1', None, 'overview.ome.zarr', 'anatomical')"
-   ]
+   "source": "transformation_path, (src_coord_system, tgt_coord_system) = ngff_transformations.graph.find_walks_in_graph(nx_graph, 'VOI-01.ome.zarr/1', None, 'overview.ome.zarr', 'anatomical')"
   },
   {
    "cell_type": "code",

src/ngff_transformations/__init__.py

@@ -2,8 +2,8 @@
 
 __version__ = version("ngff-transformations")
 
-from .graph import (
-    create_sequence_transformation_from_path,
+from ngff_transformations.graph import (
+    find_walks_in_graph,
     draw_graph,
     get_relative_path,
     transform_graph_to_networkx,

src/ngff_transformations/graph.py

@@ -6,8 +6,13 @@
 from ome_zarr_models._v06.coordinate_transforms import (
     CoordinateSystemIdentifier,
     Transform,
+    Sequence,
 )
 from pydantic import ValidationError
+import logging
+from ome_zarr_models._v06.coordinate_transforms import (
+    Sequence as SequenceTransformation,
+)
 
 
 def transform_graph_to_networkx(tgraph: TransformGraph) -> nx.DiGraph:
@@ -120,7 +125,9 @@ def _get_name_of_subgraph(
             f"Ambiguous coordinate system name '{cs_identifier}' found in both root and subgraph '{path_name}'. Use full identifier."
         )
     if cs_identifier not in nodes and cs_path_name not in nodes:
-        raise ValueError(f"Coordinate system '{cs_identifier}' not found in graph nodes.")
+        raise ValueError(
+            f"Coordinate system '{cs_identifier}' not found in graph nodes."
+        )
     if cs_path_name in nodes:
         return cs_path_name
     return cs_identifier
@@ -152,7 +159,9 @@ def _add_transform_and_inverse_transformation_edges(
         pass
 
 
-def draw_graph(g: nx.DiGraph, figsize: tuple[int, int] = (12, 8), with_edge_labels: bool = True) -> None:
+def draw_graph(
+    g: nx.DiGraph, figsize: tuple[int, int] = (12, 8), with_edge_labels: bool = True
+) -> None:
     """
     Draw a NetworkX graph showing all nodes and edges with their names.
 
@@ -210,7 +219,9 @@ def draw_graph(g: nx.DiGraph, figsize: tuple[int, int] = (12, 8), with_edge_labe
     plt.show()
 
 
-def get_relative_path(graph: nx.DiGraph, source_coordinate_system: str, target_coordinate_system: str) -> list[str]:
+def get_relative_path(
+    graph: nx.DiGraph, source_coordinate_system: str, target_coordinate_system: str
+) -> list[str]:
     cost_key = "cost"
     """
     Get the relative path from one node to another in the transformation graph.
@@ -244,23 +255,51 @@ def get_relative_path(graph: nx.DiGraph, source_coordinate_system: str, target_c
     return path
 
 
-def create_sequence_transformation_from_path(
+def create_sequence_transformation_from_graph_walk(
     graph: nx.DiGraph,
-    path: list[str],
-) -> list[Any]:
+    walk: list[str | CoordinateSystemIdentifier],
+) -> SequenceTransformation:
     """
-    Create a sequence of transformations from a path of coordinate systems
+    Create a sequence of transformations from a walk of coordinate systems
     in the transformation graph.
     """
     transformations = []
-    for i in range(len(path) - 1):
-        source = path[i]
-        target = path[i + 1]
+    for i in range(len(walk) - 1):
+        source = walk[i]
+        target = walk[i + 1]
         edge_transformation = graph.get_edge_data(source, target)["transformation"]
         transformations.append(edge_transformation)
 
-    from ome_zarr_models._v06.coordinate_transforms import Sequence
+    return Sequence(
+        input=walk[0], output=walk[-1], transformations=tuple(transformations)
+    )
 
-    transformations = Sequence(transformations=transformations)
 
-    return transformations
+def get_node(
+    path: str | None = None, name: str | None = None
+) -> str | CoordinateSystemIdentifier:
+    if path is None and name is None:
+        raise ValueError("Both path and name of the coordinate system cannot be None")
+    if path is None:
+        return name
+    if name is None:
+        return path
+    return CoordinateSystemIdentifier(path=path, name=name)
+
+
+def find_walks_in_graph(
+    graph, src_path, src_name, tgt_path, tgt_name
+) -> list[str | CoordinateSystemIdentifier]:
+    src_node = get_node(src_path, src_name)
+    tgt_node = get_node(tgt_path, tgt_name)
+
+    graph_walk = list(nx.all_shortest_paths(graph, src_node, tgt_node))
+    if not graph_walk:
+        raise ValueError(
+            f"No path found from {src_node} to {tgt_node} in the transformation graph."
+        )
+    if len(graph_walk) > 1:
+        logging.warning(
+            f"Multiple paths found from {src_node} to {tgt_node} in the transformation graph. Using the first one."
+        )
+    return graph_walk[0]

src/ngff_transformations/transform.py

@@ -1,6 +1,5 @@
-import networkx as nx
 import numpy as np
-from ome_zarr_models._v06.coordinate_transforms import CoordinateSystemIdentifier, Sequence
+from ome_zarr_models._v06.coordinate_transforms import Sequence
 from xarray import DataArray
 
 
@@ -9,32 +8,6 @@ def validata_point_shape(point: np.ndarray, transformation_sequence: Sequence):
         assert len(point) == transformation.ndim, "Point ndim doesn't match transformation ndim"
 
 
-def get_node(path: str | None = None, name: str | None = None) -> str | CoordinateSystemIdentifier:
-    if path is None and name is None:
-        raise ValueError("Both path and name of the coordinate system cannot be None")
-    if path is None:
-        return name
-    if name is None:
-        return path
-    return CoordinateSystemIdentifier(path=path, name=name)
-
-
-def find_walks_in_graph(graph, src_path, src_name, tgt_path, tgt_name):
-    src_node = get_node(src_path, src_name)
-    tgt_node = get_node(tgt_path, tgt_name)
-
-    graph_walk = list(nx.all_shortest_paths(graph, src_node, tgt_node))[0]
-
-    transformation_sequence = []
-    for i in range(len(graph_walk) - 1):
-        transformation_sequence.append(graph.get_edge_data(graph_walk[i], graph_walk[i + 1])["transformation"])
-
-    transformation_sequence = Sequence(
-        input=graph_walk[0], output=graph_walk[-1], transformations=transformation_sequence
-    )
-    return transformation_sequence, (graph_walk[0], graph_walk[-1])
-
-
 def transform_with_sequence3D(
     data: np.ndarray, axes: list[str], transformation_sequence: Sequence, output_axes: list[str]
 ) -> DataArray:
@@ -58,7 +31,7 @@ def transform_with_sequence3D(
     y_prime = transformed_points[:, 1].reshape(Y, X, Z)
     z_prime = transformed_points[:, 2].reshape(Y, X, Z)
 
-    return xarray.DataArray(
+    return DataArray(
         data,
         coords={
             "x_prime": (("y", "x", "z"), x_prime),
@@ -91,7 +64,7 @@ def transform_with_sequence(
     x_prime = transformed_points[:, 0].reshape(H, W)
     y_prime = transformed_points[:, 1].reshape(H, W)
 
-    return xarray.DataArray(
+    return DataArray(
         data,
         coords={
             "x_prime": (("y", "x"), x_prime),

tests/test_graph.py

@@ -8,7 +8,7 @@
 from ome_zarr_models._v06.image import Image
 
 from ngff_transformations.graph import (
-    create_sequence_transformation_from_path,
+    find_walks_in_graph,
     get_relative_path,
     transform_graph_to_networkx,
 )
@@ -67,6 +67,7 @@ def test_graph(zarr_path: Path):
 
     example_edge = list(nx_graph.edges)[0]
     path = get_relative_path(nx_graph, example_edge[0], example_edge[1])
-    sequence_transformation = create_sequence_transformation_from_path(nx_graph, path)
+    sequence_transformation = find_walks_in_graph(
+        graph=nx_graph, path)
 
     assert isinstance(sequence_transformation, Sequence)