feat: delegate feature persistence to copick

src/copick_utils/features/skimage.py

@@ -1,11 +1,25 @@
+import tempfile
+
 import numpy as np
-import zarr
-from numcodecs import Blosc
 from skimage.feature import multiscale_basic_features
 
 from copick_utils.io.zarr import get_level_array
 
 
+def _axis_slices(origin, chunk_size, overlap, image_size):
+    target_end = min(origin + chunk_size, image_size)
+    read_start = max(origin - overlap, 0)
+    read_end = min(target_end + overlap, image_size)
+    crop_start = origin - read_start
+    crop_end = target_end - read_start
+
+    return (
+        slice(read_start, read_end),
+        slice(crop_start, crop_end),
+        slice(origin, target_end),
+    )
+
+
 def compute_skimage_features(
     tomogram,
     feature_type,
@@ -16,6 +30,9 @@ def compute_skimage_features(
     sigma_min=0.5,
     sigma_max=16.0,
     feature_chunk_size=None,
+    *,
+    chunks=None,
+    shards=None,
 ):
     """
     Processes the tomogram chunkwise and computes the multiscale basic features.
@@ -42,61 +59,55 @@ def compute_skimage_features(
     )
     num_features = test_features.shape[-1]
 
-    # Prepare output Zarr array directly in the tomogram store
+    # Preserve the existing entity-creation timing, but defer all persistence
+    # until the complete feature tensor has been assembled on disk.
     print(f"Creating new feature store with {num_features} features...")
     copick_features = tomogram.new_features(feature_type)
-    feature_store = copick_features.zarr()
-
-    # Use the provided feature chunk size if available, otherwise default to the input chunk size
-    if feature_chunk_size is None:
-        feature_chunk_size = (num_features, *chunk_size)
-    else:
-        feature_chunk_size = (num_features, *feature_chunk_size)
-
-    out_array = zarr.create(
-        shape=(num_features, *image.shape),
-        chunks=feature_chunk_size,
-        dtype="float32",
-        compressor=Blosc(cname="zstd", clevel=3, shuffle=2),
-        store=feature_store,
-        overwrite=True,
-    )
-
-    # Process each chunk
-    for z in range(0, image.shape[0], chunk_size[0]):
-        for y in range(0, image.shape[1], chunk_size[1]):
-            for x in range(0, image.shape[2], chunk_size[2]):
-                z_start = max(z - overlap, 0)
-                z_end = min(z + chunk_size[0] + overlap, image.shape[0])
-                y_start = max(y - overlap, 0)
-                y_end = min(y + chunk_size[1] + overlap, image.shape[1])
-                x_start = max(x - overlap, 0)
-                x_end = min(x + chunk_size[2] + overlap, image.shape[2])
-
-                chunk = image[z_start:z_end, y_start:y_end, x_start:x_end]
-                chunk_features = multiscale_basic_features(
-                    chunk,
-                    intensity=intensity,
-                    edges=edges,
-                    texture=texture,
-                    sigma_min=sigma_min,
-                    sigma_max=sigma_max,
-                )
-
-                # Adjust indices for overlap
-                z_slice = slice(overlap if z_start > 0 else 0, None if z_end == image.shape[0] else -overlap)
-                y_slice = slice(overlap if y_start > 0 else 0, None if y_end == image.shape[1] else -overlap)
-                x_slice = slice(overlap if x_start > 0 else 0, None if x_end == image.shape[2] else -overlap)
-
-                # Ensure contiguous array and correct slicing
-                contiguous_chunk = np.ascontiguousarray(chunk_features[z_slice, y_slice, x_slice].transpose(3, 0, 1, 2))
-
-                out_array[
-                    0:num_features,
-                    z : z + chunk_size[0],
-                    y : y + chunk_size[1],
-                    x : x + chunk_size[2],
-                ] = contiguous_chunk
+    with tempfile.TemporaryDirectory(prefix="copick-utils-skimage-") as directory:
+        out_array = np.memmap(
+            f"{directory}/features.dat",
+            dtype=np.float32,
+            mode="w+",
+            shape=(num_features, *image.shape),
+        )
+        try:
+            # Process each chunk
+            for z in range(0, image.shape[0], chunk_size[0]):
+                for y in range(0, image.shape[1], chunk_size[1]):
+                    for x in range(0, image.shape[2], chunk_size[2]):
+                        z_read, z_crop, z_output = _axis_slices(z, chunk_size[0], overlap, image.shape[0])
+                        y_read, y_crop, y_output = _axis_slices(y, chunk_size[1], overlap, image.shape[1])
+                        x_read, x_crop, x_output = _axis_slices(x, chunk_size[2], overlap, image.shape[2])
+
+                        chunk = image[z_read, y_read, x_read]
+                        chunk_features = multiscale_basic_features(
+                            chunk,
+                            intensity=intensity,
+                            edges=edges,
+                            texture=texture,
+                            sigma_min=sigma_min,
+                            sigma_max=sigma_max,
+                        )
+
+                        contiguous_chunk = np.ascontiguousarray(
+                            chunk_features[z_crop, y_crop, x_crop].transpose(3, 0, 1, 2),
+                        )
+                        out_array[:, z_output, y_output, x_output] = contiguous_chunk
+
+            storage_chunks = chunks
+            if storage_chunks is None and feature_chunk_size is not None:
+                storage_chunks = feature_chunk_size
+
+            out_array.flush()
+            copick_features.from_numpy(
+                out_array,
+                chunks=storage_chunks,
+                shards=shards,
+                dtype=np.float32,
+                overwrite=True,
+            )
+        finally:
+            del out_array
 
     print(f"Features saved under feature type '{feature_type}'")
     return copick_features

tests/test_delegated_writers.py

@@ -8,6 +8,7 @@
 from copick.impl.filesystem import CopickConfigFSSpec, CopickRootFSSpec
 from copick.util.ome import get_level_path
 from copick_utils.converters import lazy_converter
+from copick_utils.features.skimage import compute_skimage_features
 from copick_utils.io import writers
 from copick_utils.process.rescale import rescale_segmentation
 
@@ -67,6 +68,73 @@ def test_processor_reads_and_writes_through_copick_entities(run):
     _assert_canonical_volume(derived, labels)
 
 
+@pytest.mark.parametrize(
+    ("feature_chunk_size", "chunks", "shards", "expected_chunks", "expected_shards"),
+    [
+        (None, None, None, (1, 128, 128, 128), (1, 128, 128, 128)),
+        ((3, 4, 5), None, None, (1, 3, 4, 5), (1, 6, 8, 10)),
+        ((3, 4, 5), (1, 2, 3, 4), (1, 6, 6, 8), (1, 2, 3, 4), (1, 6, 6, 8)),
+    ],
+)
+def test_feature_writer_uses_core_layout_policy(
+    run,
+    feature_chunk_size,
+    chunks,
+    shards,
+    expected_chunks,
+    expected_shards,
+):
+    volume = np.arange(5 * 6 * 7, dtype=np.float32).reshape(5, 6, 7)
+    writers.tomogram(run, volume, voxel_size=10, algorithm="wbp")
+    tomogram = run.get_voxel_spacing(10).get_tomogram("wbp")
+
+    features = compute_skimage_features(
+        tomogram,
+        "skimage",
+        run.root,
+        intensity=True,
+        edges=False,
+        texture=False,
+        sigma_min=0.5,
+        sigma_max=0.5,
+        feature_chunk_size=feature_chunk_size,
+        chunks=chunks,
+        shards=shards,
+    )
+
+    group = zarr.open_group(store=features.zarr(), mode="r")
+    level = group[get_level_path(group, 0)]
+    axes = tuple(axis["name"] for axis in group.attrs["ome"]["multiscales"][0]["axes"])
+    assert group.metadata.zarr_format == 3
+    assert axes == ("feature", "z", "y", "x")
+    assert level.metadata.dimension_names == axes
+    assert level.shape == (1, 5, 6, 7)
+    assert level.chunks == expected_chunks
+    assert level.shards == expected_shards
+    np.testing.assert_array_equal(features.numpy(), level[:])
+
+
+def test_invalid_explicit_feature_layout_is_rejected_by_core(run):
+    volume = np.arange(5 * 6 * 7, dtype=np.float32).reshape(5, 6, 7)
+    writers.tomogram(run, volume, voxel_size=10, algorithm="wbp")
+    tomogram = run.get_voxel_spacing(10).get_tomogram("wbp")
+
+    with pytest.raises(ValueError, match="shard"):
+        compute_skimage_features(
+            tomogram,
+            "invalid-layout",
+            run.root,
+            intensity=True,
+            edges=False,
+            texture=False,
+            sigma_min=0.5,
+            sigma_max=0.5,
+            feature_chunk_size=(3, 4, 5),
+            chunks=(1, 2, 3, 4),
+            shards=(1, 3, 6, 8),
+        )
+
+
 def test_lazy_worker_preserves_core_writer_failure(monkeypatch):
     task = {
         "segmentation": SimpleNamespace(session_id="input"),