refactor: port fixes on main to refactor branch

Author: galenlynch

src/aind_zarr_utils/__init__.py

@@ -24,7 +24,6 @@
     neuroglancer_annotations_to_indices,
 )
 from .origin import Origin
-from .points import Points, Space
 
 # Pipeline integration
 from .pipeline_transformed import (
@@ -36,6 +35,7 @@
     neuroglancer_to_ccf_auto_metadata,
     swc_data_to_ccf_auto_metadata,
 )
+from .points import Points, Space
 from .zarr import (
     scaled_points_to_indices,
     zarr_to_ants,

src/aind_zarr_utils/asset.py

@@ -272,7 +272,7 @@ def _eagerly(
         template_used: str,
         template_base: str | os.PathLike | None,
     ) -> Asset:
-        """Internal: build an :class:`Asset` and pre-open its Zarr."""
+        """Build an :class:`Asset` and pre-open its Zarr (internal helper)."""
         kwargs: dict[str, Any] = {
             "zarr_uri": zarr_uri,
             "metadata": metadata,

src/aind_zarr_utils/io/processing.py

@@ -10,11 +10,15 @@
 
 from __future__ import annotations
 
+import logging
 from pathlib import PurePosixPath
 from typing import Any
 
 from aind_zarr_utils.io.paths import _zarr_base_name_pathlike
 
+logger = logging.getLogger(__name__)
+_KNOWN_GOOD_PIPELINE_VERSIONS = {3, 4, 5}
+
 
 def _get_processing_pipeline_data(
     processing_data: dict[str, Any],
@@ -42,8 +46,17 @@ def _get_processing_pipeline_data(
     if not ver_str:
         raise ValueError("Missing pipeline version")
     pipeline_ver = int(ver_str.split(".")[0])
-    if pipeline_ver not in set((3, 4)):
-        raise ValueError(f"Unsupported pipeline version: {pipeline_ver}")
+    if pipeline_ver not in _KNOWN_GOOD_PIPELINE_VERSIONS:
+        maxver = max(_KNOWN_GOOD_PIPELINE_VERSIONS)
+        if pipeline_ver > maxver:
+            logger.warning(
+                f"Pipeline version {pipeline_ver} is greater than max "
+                f"verified version {maxver}, results may not be accurate. "
+                "File an issue at "
+                "https://github.com/AllenNeuralDynamics/aind-zarr-utils/issues."
+            )
+        else:
+            raise ValueError(f"Unsupported pipeline version: {pipeline_ver}")
     pipeline: dict[str, Any] = processing_data.get("processing_pipeline", {})
     return pipeline
 

src/aind_zarr_utils/io/transforms.py

@@ -15,23 +15,25 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-from aind_s3_cache.s3_cache import (
-    get_local_path_for_resource,
-)
+from aind_s3_cache.s3_cache import get_local_path_for_resource
 from aind_s3_cache.uri_utils import as_pathlike, as_string, join_any
 
 from aind_zarr_utils.io.paths import _asset_from_zarr_pathlike
 from aind_zarr_utils.io.processing import (
+    _get_processing_pipeline_data,
     image_atlas_alignment_path_relative_from_processing,
 )
 
 if TYPE_CHECKING:
     from mypy_boto3_s3 import S3Client
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass(slots=True, frozen=True)
 class TransformChain:
@@ -108,7 +110,7 @@ class TemplatePaths:
     )
 }
 
-_PIPELINE_INDIVIDUAL_TRANSFORM_CHAINS: dict[int, TransformChain] = {
+_KNOWN_INDIVIDUAL_TRANSFORM_CHAINS: dict[int, TransformChain] = {
     3: TransformChain(
         fixed="template",
         moving="individual",
@@ -125,6 +127,34 @@ class TemplatePaths:
     )
 }
 
+_PIPELINE_INDIVIDUAL_TRANSFORM_CHAINS: dict[int, TransformChain] = {
+    3: _KNOWN_INDIVIDUAL_TRANSFORM_CHAINS[3],
+    4: _KNOWN_INDIVIDUAL_TRANSFORM_CHAINS[3],
+    5: _KNOWN_INDIVIDUAL_TRANSFORM_CHAINS[3],
+}
+
+
+def _resolve_individual_transform_chain(
+    pipeline_ver: int,
+) -> TransformChain:
+    """Look up the individual (LS → template) transform chain for a pipeline major version.
+
+    Falls back to the latest known chain when the version is newer than
+    anything registered.
+    """
+    chain = _PIPELINE_INDIVIDUAL_TRANSFORM_CHAINS.get(pipeline_ver)
+    if chain is not None:
+        return chain
+    max_known = max(_PIPELINE_INDIVIDUAL_TRANSFORM_CHAINS)
+    if pipeline_ver > max_known:
+        logger.warning(
+            f"No individual transform chain registered for pipeline "
+            f"version {pipeline_ver}; falling back to chain for version "
+            f"{max_known}. Results may not be accurate."
+        )
+        return _PIPELINE_INDIVIDUAL_TRANSFORM_CHAINS[max_known]
+    raise ValueError(f"No individual transform chain registered for pipeline version {pipeline_ver}")
+
 
 def pipeline_transforms(
     zarr_uri: str,
@@ -169,9 +199,12 @@ def pipeline_transforms(
         bucket,
         asset_pathlike / alignment_rel_path,
     )
+    pipeline = _get_processing_pipeline_data(processing_data)
+    pipeline_ver = int(pipeline["pipeline_version"].split(".")[0])
+    transform_chain = _resolve_individual_transform_chain(pipeline_ver)
     individual_ants_paths = TemplatePaths(
         alignment_path,
-        _PIPELINE_INDIVIDUAL_TRANSFORM_CHAINS[3],
+        transform_chain,
     )
     if template_base:
         template_ants_paths = TemplatePaths(

src/aind_zarr_utils/pipeline_domain_selector.py

@@ -23,6 +23,9 @@
     fix_corner_compute_origin as fix_corner_compute_origin,
 )
 
+from aind_zarr_utils.domain.overlays import (
+    _PIPELINE_MULTISCALE_FACTOR as _PIPELINE_MULTISCALE_FACTOR,
+)
 from aind_zarr_utils.domain.overlays import (
     FlipIndexAxesOverlay as FlipIndexAxesOverlay,
 )
@@ -44,9 +47,6 @@
 from aind_zarr_utils.domain.overlays import (
     Vec3 as Vec3,
 )
-from aind_zarr_utils.domain.overlays import (
-    _PIPELINE_MULTISCALE_FACTOR as _PIPELINE_MULTISCALE_FACTOR,
-)
 from aind_zarr_utils.domain.overlays import (
     _require_cardinal as _require_cardinal,
 )

src/aind_zarr_utils/pipeline_transformed.py

@@ -21,20 +21,35 @@
 
 import numpy as np
 import SimpleITK as sitk
-from aind_anatomical_utils.anatomical_volume import (
-    AnatomicalHeader,
-    fix_corner_compute_origin,
-)
-from aind_anatomical_utils.coordinate_systems import _OPPOSITE_AXES
 from aind_registration_utils.ants import (
     apply_ants_transforms_to_point_arr,
 )
 from numpy.typing import NDArray
-from packaging.version import Version
 
 from aind_zarr_utils.annotations import annotation_indices_to_anatomical
 from aind_zarr_utils.formats.swc import swc_data_to_indices
-from aind_zarr_utils.io.metadata import _unit_conversion as _unit_conversion  # noqa: F401  # legacy re-export
+
+# The implementations of _pipeline_anatomical_check_args,
+# _apply_pipeline_overlays_to_header, and _mimic_pipeline_anatomical_header
+# moved to ``aind_zarr_utils.image`` in commit C4. They are re-exported here
+# (along with _build_pipeline_header, the new name for the third) so existing
+# callers and test patches keep working.
+from aind_zarr_utils.image import (
+    _apply_pipeline_overlays_to_header as _apply_pipeline_overlays_to_header,
+)
+from aind_zarr_utils.image import (
+    _build_pipeline_header as _build_pipeline_header,
+)
+from aind_zarr_utils.image import (
+    _build_pipeline_header as _mimic_pipeline_anatomical_header,
+)
+from aind_zarr_utils.image import (
+    _pipeline_anatomical_check_args as _pipeline_anatomical_check_args,
+)
+from aind_zarr_utils.image import (
+    apply_pipeline_overlays as apply_pipeline_overlays,
+)
+from aind_zarr_utils.io.metadata import _unit_conversion as _unit_conversion  # legacy re-export
 
 # Re-exported from io/* so existing test patches (and downstream imports)
 # continue to find these names at ``aind_zarr_utils.pipeline_transformed.*``.
@@ -105,14 +120,11 @@
 )
 from aind_zarr_utils.pipeline_domain_selector import (
     OverlaySelector,
-    apply_overlays,
-    estimate_pipeline_multiscale,
     get_selector,
 )
 from aind_zarr_utils.zarr import (
     zarr_to_ants,
     zarr_to_sitk,
-    zarr_to_sitk_stub,
 )
 
 if TYPE_CHECKING:
@@ -123,28 +135,6 @@
 T = TypeVar("T", int, float)
 
 
-# The implementations of _pipeline_anatomical_check_args,
-# _apply_pipeline_overlays_to_header, and _mimic_pipeline_anatomical_header
-# moved to ``aind_zarr_utils.image`` in commit C4. They are re-exported here
-# (along with _build_pipeline_header, the new name for the third) so existing
-# callers and test patches keep working.
-from aind_zarr_utils.image import (
-    _apply_pipeline_overlays_to_header as _apply_pipeline_overlays_to_header,
-)
-from aind_zarr_utils.image import (
-    _build_pipeline_header as _build_pipeline_header,
-)
-from aind_zarr_utils.image import (
-    _build_pipeline_header as _mimic_pipeline_anatomical_header,
-)
-from aind_zarr_utils.image import (
-    _pipeline_anatomical_check_args as _pipeline_anatomical_check_args,
-)
-from aind_zarr_utils.image import (
-    apply_pipeline_overlays as apply_pipeline_overlays,
-)
-
-
 def base_and_pipeline_anatomical_stub(
     zarr_uri: str,
     metadata: dict,

tests/conftest.py

@@ -243,19 +243,13 @@ def mock_overlay_selector(monkeypatch):
     def mock_apply_overlays(header, overlays, meta, multiscale_no, **_kwargs):
         return header, []  # Return header unchanged with no applied overlays
 
-    # ``apply_overlays`` is referenced from three modules now: the canonical
-    # home (``aind_zarr_utils.domain.selector``), the legacy
-    # ``aind_zarr_utils.pipeline_transformed`` re-export (kept for tests
-    # that patch the old path), and ``aind_zarr_utils.image`` where the
-    # actual call sites in the C4 dispatcher live. Patch all three.
+    # ``apply_overlays`` lives at ``aind_zarr_utils.domain.selector`` (the
+    # canonical home) and is called from ``aind_zarr_utils.image`` where the
+    # C4 dispatcher routes through it. Patch both.
     monkeypatch.setattr(
         "aind_zarr_utils.domain.selector.apply_overlays",
         mock_apply_overlays,
     )
-    monkeypatch.setattr(
-        "aind_zarr_utils.pipeline_transformed.apply_overlays",
-        mock_apply_overlays,
-    )
     monkeypatch.setattr(
         "aind_zarr_utils.image.apply_overlays",
         mock_apply_overlays,

tests/test_asset.py

@@ -10,7 +10,6 @@
 from aind_zarr_utils.asset import Asset, TransformPaths
 from aind_zarr_utils.origin import Origin
 
-
 # ---------------------------------------------------------------- Origin ---
 
 

tests/test_points.py

@@ -25,7 +25,6 @@
     transform_points,
 )
 
-
 # ----------------------------------------------------------- construction ---
 
 
@@ -450,9 +449,8 @@ def test_neuroglancer_to_ccf_numerical_equivalence_new_vs_legacy(
     (via ``mock_ants_transforms``) so the result is deterministic and the
     test does not depend on the registration files being available.
     """
-    from tests.conftest import create_comprehensive_processing_data
-
     from aind_zarr_utils.pipeline_transformed import neuroglancer_to_ccf
+    from tests.conftest import create_comprehensive_processing_data
 
     processing = create_comprehensive_processing_data()