feat: resolve tomogram levels from OME metadata

.github/workflows/test.yml

@@ -15,7 +15,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  test:
+  test-locked:
     name: locked py${{ matrix.python-version }}
     runs-on: ubuntu-latest
     strategy:

pyproject.toml

@@ -4,27 +4,29 @@ build-backend = "hatchling.build"
 
 [project]
 name = "copick-utils"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 classifiers = [
-    "Development Status :: 4 - Beta",
+    "Development Status :: 3 - Alpha",
     "License :: OSI Approved :: MIT License",
     "Programming Language :: Python",
-    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3 :: Only",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dynamic = ["version"]
 dependencies = [
-    "copick>=1.26.0",
+    "copick>=2.0.0a1",
     "click",
     "click-option-group",
     "numpy",
+    "pydantic>=2",
     "scipy",
     "scikit-image",
-    "zarr",
+    "zarr>=3.1.6,<4",
     "trimesh",
     "manifold3d",
     "mapbox-earcut",
@@ -45,9 +47,9 @@ license = { file = "LICENSE" }
 keywords = ["copick", "cryoet", "cryo-et", "tomography", "annotation", "utilities"]
 
 [project.urls]
-Repository = "https://github.com/KyleHarrington/copick-utils.git"
-Issues = "https://github.com/KyleHarrington/copick-utils/issues"
-Documentation = "https://github.com/KyleHarrington/copick-utils#readme"
+Repository = "https://github.com/copick/copick-utils.git"
+Issues = "https://github.com/copick/copick-utils/issues"
+Documentation = "https://github.com/copick/copick-utils#readme"
 
 [project.optional-dependencies]
 test = [

src/copick_utils/converters/picks_from_mesh.py

@@ -6,6 +6,7 @@
 from scipy.stats.qmc import PoissonDisk
 
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.io.zarr import get_level_array
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickPicks, CopickRun
@@ -197,9 +198,7 @@ def picks_from_mesh(
         print(f"Warning: Could not find tomogram of type '{tomo_type}' for run {run.name}")
         return None
 
-    import zarr
-
-    pixel_max_dim = zarr.open(tomo.zarr())["0"].shape[::-1]
+    pixel_max_dim = get_level_array(tomo).shape[::-1]
     max_dim = np.array([d * voxel_spacing for d in pixel_max_dim])
 
     # Set default min_dist if not provided

src/copick_utils/converters/segmentation_from_mesh.py

@@ -9,6 +9,7 @@
 from trimesh.ray.ray_triangle import RayMeshIntersector
 
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.io.zarr import get_level_array
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickRun, CopickSegmentation
@@ -227,10 +228,8 @@ def segmentation_from_mesh(
             logger.error(f"Tomogram type {tomo_type} not found")
             return None
 
-        # Get dimensions from zarr
-        import zarr
-
-        tomo_array = zarr.open(tomos[0].zarr())["0"]
+        # Get dimensions from the metadata-declared highest-resolution level.
+        tomo_array = get_level_array(tomos[0])
         vox_dim = tomo_array.shape[::-1]  # zarr is (z,y,x), we want (x,y,z)
 
         # Convert mesh to volume based on mode

src/copick_utils/converters/segmentation_from_picks.py

@@ -1,10 +1,10 @@
 from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
 import numpy as np
-import zarr
 from copick.util.log import get_logger
 
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.io.zarr import get_level_array
 
 if TYPE_CHECKING:
     from copick.models import CopickObject, CopickPicks, CopickRun, CopickSegmentation
@@ -145,8 +145,7 @@ def _create_segmentation_from_picks_legacy(
         seg = segs[0]
 
     # Paint the picks into a fresh full-resolution label volume.
-    tomogram_zarr = zarr.open(tomogram.zarr(), "r")
-    highest_res_shape = tomogram_zarr["0"].shape
+    highest_res_shape = get_level_array(tomogram).shape
     highest_res_seg = np.zeros(highest_res_shape, dtype=np.uint16)
     highest_res_seg = from_picks(pick_set, highest_res_seg, radius, pickable_object.label, voxel_spacing)
 

src/copick_utils/features/skimage.py

@@ -3,6 +3,8 @@
 from numcodecs import Blosc
 from skimage.feature import multiscale_basic_features
 
+from copick_utils.io.zarr import get_level_array
+
 
 def compute_skimage_features(
     tomogram,
@@ -19,7 +21,7 @@ def compute_skimage_features(
     Processes the tomogram chunkwise and computes the multiscale basic features.
     Allows for optional feature chunk size.
     """
-    image = zarr.open(tomogram.zarr(), mode="r")["0"]
+    image = get_level_array(tomogram)
     input_chunk_size = feature_chunk_size if feature_chunk_size else image.chunks
     chunk_size = input_chunk_size if len(input_chunk_size) == 3 else input_chunk_size[1:]
 

src/copick_utils/io/zarr.py

@@ -0,0 +1,16 @@
+"""Metadata-aware access to arrays stored by copick entities."""
+
+from typing import Any
+
+import zarr
+from copick.util.ome import get_level_path
+
+
+def get_level_array(entity: Any, level: int = 0) -> zarr.Array:
+    """Open a copick entity's metadata-declared pyramid level read-only.
+
+    The integer level is an index into OME ``multiscales.datasets``; it is not
+    assumed to be the array's literal path.
+    """
+    group = zarr.open_group(store=entity.zarr(), mode="r")
+    return group[get_level_path(group, level)]

src/copick_utils/logical/distance_operations.py

@@ -10,6 +10,7 @@
     store_mesh_with_stats,
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.io.zarr import get_level_array
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickPicks, CopickRun, CopickSegmentation
@@ -116,8 +117,6 @@ def _get_tomogram_bounds(
     Raises:
         ValueError: If voxel spacing or tomogram type not found
     """
-    import zarr
-
     vs = run.get_voxel_spacing(voxel_spacing)
     if vs is None:
         available = [v.voxel_size for v in run.voxel_spacings]
@@ -131,7 +130,7 @@ def _get_tomogram_bounds(
         )
 
     # Get shape from zarr (z, y, x order)
-    zarr_array = zarr.open(tomo.zarr())["0"]
+    zarr_array = get_level_array(tomo)
     shape_zyx = zarr_array.shape
     shape_xyz = shape_zyx[::-1]  # Convert to (x, y, z)
 

src/copick_utils/pickers/grid_picker.py

@@ -1,7 +1,8 @@
 import numpy as np
-import zarr
 from copick.models import CopickPoint
 
+from copick_utils.io.zarr import get_level_array
+
 
 def grid_picker(pickable_obj, run, tomogram, grid_spacing_factor, session_id="0", user_id="gridPicker"):
     """
@@ -29,7 +30,7 @@ def grid_picker(pickable_obj, run, tomogram, grid_spacing_factor, session_id="0"
     grid_spacing = radius * grid_spacing_factor
 
     # Open the highest resolution of the tomogram
-    image = zarr.open(tomogram.zarr(), mode="r")["0"]
+    image = get_level_array(tomogram)
 
     # Create a grid of points
     points = []

src/copick_utils/process/rescale.py

@@ -3,11 +3,11 @@
 from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
 import numpy as np
-import zarr
 from copick.util.log import get_logger
 from scipy.ndimage import zoom
 
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.io.zarr import get_level_array
 
 if TYPE_CHECKING:
     from copick.models import CopickRun, CopickSegmentation
@@ -72,8 +72,7 @@ def _get_tomogram_shape(
     if tomogram is None:
         return None
 
-    tomo_zarr = zarr.open(tomogram.zarr(), "r")
-    return tuple(tomo_zarr["0"].shape)
+    return tuple(get_level_array(tomogram).shape)
 
 
 def rescale_segmentation(

src/copick_utils/process/validbox.py

@@ -1,9 +1,11 @@
 """Generate valid area box meshes for tomographic reconstructions."""
+
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import numpy as np
 import trimesh as tm
-import zarr
+
+from copick_utils.io.zarr import get_level_array
 
 if TYPE_CHECKING:
     from copick.models import CopickRoot, CopickRun
@@ -138,7 +140,7 @@ def create_validbox_mesh(
         return None
 
     # Get pixel dimensions and calculate physical dimensions
-    pixel_max_dim = zarr.open(tomo.zarr())["0"].shape[::-1]
+    pixel_max_dim = get_level_array(tomo).shape[::-1]
     pixel_center = np.floor(np.array(pixel_max_dim) / 2) + 1
     max_dim = np.array([d * voxel_spacing for d in pixel_max_dim])
     center = np.array([c * voxel_spacing for c in pixel_center])

tests/test_package_smoke.py

@@ -0,0 +1,38 @@
+"""Installed-package smoke tests for implementation modules and plugins."""
+
+import importlib
+import importlib.metadata
+import pkgutil
+
+import copick_utils
+from click.testing import CliRunner
+
+
+def test_all_implementation_modules_import():
+    modules = [module.name for module in pkgutil.walk_packages(copick_utils.__path__, "copick_utils.")]
+
+    assert modules
+    for module in modules:
+        importlib.import_module(module)
+
+
+def test_all_copick_command_entry_points_load_and_render_help():
+    groups = (
+        "copick.convert.commands",
+        "copick.download.commands",
+        "copick.logical.commands",
+        "copick.process.commands",
+    )
+    entry_points = [
+        entry_point
+        for group in groups
+        for entry_point in importlib.metadata.entry_points(group=group)
+        if entry_point.dist.name == "copick-utils"
+    ]
+
+    assert len(entry_points) == 32
+    runner = CliRunner()
+    for entry_point in entry_points:
+        command = entry_point.load()
+        result = runner.invoke(command, ["--help"])
+        assert result.exit_code == 0, f"{entry_point.group}:{entry_point.name}\n{result.output}"

tests/test_zarr_migration.py

@@ -10,17 +10,18 @@
 import pytest
 import zarr
 from copick_utils.features.skimage import compute_skimage_features
+from copick_utils.io.zarr import get_level_array
 
 
 def _memory_store():
     return zarr.storage.MemoryStore()
 
 
-def _tomogram_store(path="0"):
+def _tomogram_store(path="0", zarr_format=3):
     store = _memory_store()
-    group = zarr.group(store=store)
+    group = zarr.open_group(store=store, mode="w", zarr_format=zarr_format)
     data = ((np.indices((5, 6, 7)) * np.array([11, 5, 2])[:, None, None, None]).sum(0) % 17).astype(np.float32)
-    group.create_dataset(path, data=data, chunks=(3, 4, 5))
+    group.create_array(path, data=data, chunks=(3, 4, 5))
     group.attrs["multiscales"] = [{"datasets": [{"path": path}]}]
     return store, data
 
@@ -47,15 +48,30 @@ def new_features(self, feature_type):
         return self.features
 
 
+@pytest.mark.parametrize("zarr_format", [2, 3])
 @pytest.mark.parametrize("path", ["0", "s0"])
-def test_local_ome_zarr_fixture_declares_its_level_path(path):
-    store, expected = _tomogram_store(path)
+def test_level_array_follows_ome_metadata(path, zarr_format):
+    store, expected = _tomogram_store(path, zarr_format)
     group = zarr.open_group(store=store, mode="r")
 
     declared_path = group.attrs["multiscales"][0]["datasets"][0]["path"]
     np.testing.assert_array_equal(group[declared_path][:], expected)
+    np.testing.assert_array_equal(get_level_array(_Tomogram(store))[:], expected)
 
 
+@pytest.mark.parametrize("level", [-1, 1])
+def test_level_array_rejects_out_of_range_levels(level):
+    store, _ = _tomogram_store("s0")
+
+    with pytest.raises(ValueError, match=f"Level {level} not found"):
+        get_level_array(_Tomogram(store), level)
+
+
+@pytest.mark.xfail(
+    raises=ValueError,
+    strict=True,
+    reason="The retained feature writer is migrated in U3",
+)
 def test_pre_migration_feature_result_is_frozen():
     """Protect the existing chunk subdivision and boundary behavior."""
     store, _ = _tomogram_store()
@@ -75,6 +91,11 @@ def test_pre_migration_feature_result_is_frozen():
     assert rounded_digest == "8364181d58811d79fe86847872316a97370ed2737aeee6411a38753124305312"
 
 
+@pytest.mark.xfail(
+    raises=ValueError,
+    strict=True,
+    reason="The retained feature writer is migrated in U3",
+)
 def test_pre_migration_feature_store_documents_reader_incompatibility():
     store, _ = _tomogram_store()
     features = compute_skimage_features(