Store precomputed skeleton caches with source backend

Author: Donglai Wei

connectomics/data/io/io.py

@@ -388,7 +388,7 @@ def save_volume(
         filename: Output filename or directory path.
         volume: Volume data to save.
         dataset: Dataset name for HDF5 format.
-        file_format: 'h5', 'tiff', 'png', 'nii', 'nii.gz'.
+        file_format: Optional override. If omitted, inferred from ``filename``.
     """
     file_format = file_format or _detect_format(filename)
 

tests/unit/test_data_factory_zarr_aux.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+
+import numpy as np
+
+from connectomics.config import Config
+from connectomics.data.io import save_volume
+from connectomics.training.lightning.data_factory import _maybe_precompute_label_aux
+
+
+def test_maybe_precompute_label_aux_reuses_existing_zarr_cache(monkeypatch, tmp_path: Path):
+    label_path = tmp_path / "data.zarr" / "seg"
+    aux_path = tmp_path / "data.zarr" / "seg_skeleton"
+    save_volume(str(label_path), np.zeros((2, 2, 2), dtype=np.uint16), file_format="zarr")
+    save_volume(str(aux_path), np.ones((2, 2, 2), dtype=np.uint16), file_format="zarr")
+
+    cfg = Config()
+    cfg.data.label_transform.targets = [{"name": "skeleton_aware_edt", "kwargs": {}}]
+    cfg.data.train.label_aux_type = "skeleton"
+
+    def fail(*args, **kwargs):
+        raise AssertionError("existing zarr cache should be reused")
+
+    monkeypatch.setattr(
+        "connectomics.data.processing.distance.precompute_skeleton_volume",
+        fail,
+    )
+
+    paths = _maybe_precompute_label_aux(
+        cfg,
+        cfg.data.train,
+        [str(label_path)],
+        split_name="train",
+    )
+
+    assert paths == [str(aux_path)]

tests/unit/test_io_zarr.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from connectomics.data.io import get_vol_shape, read_volume, save_volume, volume_exists
+from connectomics.data.processing.distance import sdt_path_for_label
+
+zarr = pytest.importorskip("zarr")
+
+
+def test_sdt_path_for_label_preserves_zarr_backend():
+    assert sdt_path_for_label("/tmp/labels.h5", mode="skeleton") == "/tmp/labels_skeleton.h5"
+    assert (
+        sdt_path_for_label("/tmp/data.zarr/seg", mode="skeleton") == "/tmp/data.zarr/seg_skeleton"
+    )
+    assert (
+        sdt_path_for_label("/tmp/data.zarr/group/seg", mode="sdt") == "/tmp/data.zarr/group/seg_sdt"
+    )
+    assert sdt_path_for_label("/tmp/data.zarr", mode="skeleton") == "/tmp/data_skeleton.zarr"
+
+
+def test_zarr_subkey_round_trip_and_shape(tmp_path: Path):
+    path = tmp_path / "data.zarr" / "seg_skeleton"
+    volume = np.arange(24, dtype=np.uint16).reshape(2, 3, 4)
+
+    save_volume(str(path), volume, file_format="zarr")
+
+    loaded = read_volume(str(path))
+    np.testing.assert_array_equal(loaded, volume)
+    assert get_vol_shape(str(path)) == volume.shape
+    assert volume_exists(str(path))
+
+
+def test_detect_format_prefers_real_suffix_inside_zarr_directory(tmp_path: Path):
+    h5_path = tmp_path / "data.zarr" / "seg_skeleton.h5"
+    h5_path.parent.mkdir(parents=True, exist_ok=True)
+    save_volume(str(h5_path), np.ones((2, 2, 2), dtype=np.uint8), file_format="h5")
+
+    loaded = read_volume(str(h5_path))
+
+    assert loaded.shape == (2, 2, 2)