refactor: port fixes on main to refactor branch

Author: galenlynch

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,23 @@
 
 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 (
-    image_atlas_alignment_path_relative_from_processing,
-)
+    _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 +108,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 +125,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 +197,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(