Implement pad-wrapper

AGENTS.md

@@ -3,25 +3,31 @@
 This library implements segmentation and other image analysis functionality in python.
 The functionality is implemented in `elf` and its submodules.
 
-## Current task
-
-The library relies heavily on implementations from `nifty`, `vigra`, and `affogato`, which are C++ bindings with python libraries.
-However, these libraries are difficult to install (e.g. not available on PyPI), limiting accessibility.
-
-We have developed a new library `bioimage-cpp` to bundle the needed functionality and to make it more easily available.
-Now, we want to migrate `elf` to use `bioimage-cpp`to eliminate `nifty`, `vigra`, and `affogato` dependencies.
-In addition, we also want to use more efficient implementation from `bioimage-cpp` over `scipy`, `skimage` where available.
-Please refer to `bioimage-cpp`'s migration guide for how to migrate the functionality:
-https://raw.githubusercontent.com/computational-cell-analytics/bioimage-cpp/refs/heads/main/MIGRATION_GUIDE.md
-
-We will go through the sub-modules one by one to migrate functionality. For each sub-module please follow this strategy:
-- First check for sufficient test coverage and add further tests to increase it.
-- Then check for general issues or missing functionality (indicated by TODOs etc.) in the sub-module and fix them.
-- Then check for functionality to migrate.
-    - If you cannot find a matching function in `bioimage-cpp` let me know and it will be implemented there.
-- Then do the migration, run tests to ensure its success.
-
-We will go through each of these steps with planning mode.
+## Library structure
+
+All functionality lives in the `elf` package. Each submodule is self-contained and exposes its public API via its `__init__.py`. The `test` folder mirrors this layout, with one subfolder per submodule.
+
+Top-level utilities:
+- `elf/util.py`: General-purpose helpers used across submodules.
+- `elf/__version__.py`: Single source of truth for the package version.
+
+Submodules:
+- `elf/io`: Unified interface for reading/writing large microscopy data. `files.py` is the main entry point (`open_file`); the `*_wrapper.py` files adapt specific backends (zarr, n5, mrc, knossos, nifti, intern, image stacks, etc.).
+- `elf/wrapper`: Lazy array wrappers that apply transformations on-the-fly (affine, resize, caching). `base.py` defines the wrapper base class; `generic.py` provides a generic apply-function wrapper.
+- `elf/parallel`: Blockwise/parallel implementations of array operations (label, relabel, watershed, distance transform, filters, size filter, unique, stats, copy). Use these for data too large to fit in memory.
+- `elf/segmentation`: Core segmentation algorithms — (lifted) multicut, mutex watershed, GASP, clustering, watershed, plus blockwise variants (`blockwise_*_impl.py`), feature computation (`features.py`), embeddings, learning, stitching, postprocessing, and high-level pipelines in `workflows.py`.
+- `elf/evaluation`: Segmentation metrics — rand index, variation of information, dice, cremi score, and matching-based scores.
+- `elf/transformation`: Coordinate/image transformations, including affine transforms, resizing, NGFF transforms, and elastix/transformix wrappers.
+- `elf/tracking`: Graph-based tracking (motile), with MaMuT import/export and shared utilities.
+- `elf/mesh`: Mesh generation from segmentations, mesh I/O, and mesh-to-segmentation conversion.
+- `elf/skeleton`: Skeletonization (`skeletonize.py`, `thinning.py`) and skeleton I/O.
+- `elf/label_multiset`: Paintera-style label multiset data structure — creation, serialization, and the core data structure.
+- `elf/htm`: High-throughput microscopy helpers — parsing and visualization.
+- `elf/ilastik`: Interop with ilastik (e.g. carving).
+- `elf/visualisation`: Visualization helpers (edges, grids, metrics, object/size views).
+- `elf/color`: Color palettes for visualization.
+
+Note the British spelling in `elf/visualisation`. Heavy numerical routines are increasingly backed by [bioimage-cpp](https://github.com/computational-cell-analytics/bioimage-cpp) (replacing the older affogato/nifty/vigra dependencies); some functions degrade gracefully or raise if an optional backend is missing.
 
 ## Coding standards
 

elf/wrapper/__init__.py

@@ -2,4 +2,4 @@
 """
 
 from .base import SimpleTransformationWrapper, SimpleTransformationWrapperWithHalo, TransformationWrapper
-from .generic import NormalizeWrapper, ThresholdWrapper, RoiWrapper
+from .generic import NormalizeWrapper, ThresholdWrapper, RoiWrapper, PadWrapper

elf/wrapper/generic.py

@@ -48,28 +48,92 @@ class RoiWrapper(WrapperBase):
 
     Args:
         volume: The data to wrap.
-        roi: The region of interest
+        roi: The region of interest.
+        squeeze: Whether to drop singleton axes introduced by integer entries in `roi`
+            from the wrapper's shape and output. Defaults to False.
     """
-    def __init__(self, volume: ArrayLike, roi: Tuple[slice, ...]):
+    def __init__(self, volume: ArrayLike, roi: Tuple[slice, ...], squeeze: bool = False):
         super().__init__(volume)
-        self._roi, _ = normalize_index(roi, volume.shape)
+        self._roi, roi_squeeze = normalize_index(roi, volume.shape)
+        self._squeeze = squeeze
+        # Full-volume axis positions introduced as singletons by the roi.
+        self._squeeze_axes = roi_squeeze if squeeze else ()
+        # Full-volume axes that remain visible in the (reduced) wrapper shape.
+        self._kept_axes = tuple(ax for ax in range(len(self._roi)) if ax not in self._squeeze_axes)
 
     @property
     def shape(self):
-        return tuple(b.stop - b.start for b in self._roi)
+        return tuple(self._roi[ax].stop - self._roi[ax].start for ax in self._kept_axes)
+
+    @property
+    def ndim(self):
+        return len(self._kept_axes)
+
+    @property
+    def chunks(self):
+        chunks = super().chunks
+        return None if chunks is None else tuple(chunks[ax] for ax in self._kept_axes)
 
     def map_index_to_volume(self, index):
-        index = tuple(slice(ind.start + roi.start, ind.stop + roi.start)
-                      for ind, roi in zip(index, self._roi))
-        return index
+        full_index = list(self._roi)  # Default to the roi slice for every axis.
+        for ind, ax in zip(index, self._kept_axes):
+            roi = self._roi[ax]
+            full_index[ax] = slice(ind.start + roi.start, ind.stop + roi.start)
+        return tuple(full_index)
 
     def __getitem__(self, index):
         index, to_squeeze = normalize_index(index, self.shape)
         index = self.map_index_to_volume(index)
         out = self._volume[index]
-        return squeeze_singletons(out, to_squeeze)
+        squeeze = sorted(set(self._squeeze_axes) | {self._kept_axes[i] for i in to_squeeze})
+        return squeeze_singletons(out, tuple(squeeze))
 
     def __setitem__(self, index, item):
         index, _ = normalize_index(index, self.shape)
         index = self.map_index_to_volume(index)
         self._volume[index] = item
+
+
+class PadWrapper(WrapperBase):
+    """Wrapper to pad the input.
+
+    This only supports right-padding.
+
+    Args:
+        volume: The data to pad.
+        pad_shape: The shape for padding the data.
+    """
+    def __init__(self, volume: ArrayLike, pad_width: Tuple[int, ...], mode: str = "constant"):
+        assert volume.ndim == len(pad_width)
+        super().__init__(volume)
+        self._pad_width = pad_width
+        self._shape = volume.shape
+        self._mode = mode
+
+    @property
+    def shape(self):
+        return tuple(sh + pw for sh, pw in zip(self._shape, self._pad_width))
+
+    def __getitem__(self, index):
+        index, to_squeeze = normalize_index(index, self.shape)
+
+        local_pad, local_index = [], []
+        for idx, sh in zip(index, self._shape):
+            overhang_start = max(0, idx.start - sh)
+            overhang_stop = max(0, idx.stop - sh)
+            if overhang_start > 0:
+                raise NotImplementedError
+            elif overhang_stop > 0:
+                local_pad.append(overhang_stop)
+                local_index.append(slice(idx.start, sh))
+            else:
+                local_pad.append(0)
+                local_index.append(idx)
+
+        local_index = tuple(local_index)
+        out = self._volume[local_index]
+        if any(lpad > 0 for lpad in local_pad):
+            pad_width = tuple((0, lpad) for lpad in local_pad)
+            out = np.pad(out, pad_width, mode=self._mode)
+
+        return squeeze_singletons(out, to_squeeze)

test/wrapper/test_generic.py

@@ -39,6 +39,61 @@ def test_roi_wrapper(self):
             exp = x[bb]
             self.assertTrue(np.allclose(res, exp))
 
+    def test_roi_wrapper_squeeze(self):
+        from elf.wrapper import RoiWrapper
+
+        x = np.random.rand(*self.shape)
+        # Each roi introduces one or more singleton axes via integer indexing.
+        bbs = [np.s_[0], np.s_[5, 10], np.s_[3, 4, 5], np.s_[2, 10:40, :]]
+        for bb in bbs:
+            wrapped = RoiWrapper(x, bb, squeeze=True)
+            exp = x[bb]
+            self.assertEqual(wrapped.shape, np.shape(exp))
+            self.assertEqual(wrapped.ndim, len(wrapped.shape))
+            self.assertEqual(wrapped.ndim, np.ndim(exp))
+            # A fully-reduced (scalar) wrapper must be indexed with `()`, just like numpy.
+            res = wrapped[()] if wrapped.shape == () else wrapped[:]
+            self.assertEqual(np.shape(res), np.shape(exp))
+            self.assertTrue(np.array_equal(res, exp))
+
+        # Sub-indexing a wrapper with squeezed roi axes maps to the correct volume region.
+        wrapped = RoiWrapper(x, np.s_[2, 10:40, :], squeeze=True)
+        self.assertTrue(np.array_equal(wrapped[:5, :5], x[2, 10:15, 0:5]))
+
+        # Round-trip setitem writes into the correct volume region.
+        y = np.zeros(self.shape)
+        wrapped = RoiWrapper(y, np.s_[2, 10:40, :], squeeze=True)
+        val = np.random.rand(*wrapped.shape)
+        wrapped[:] = val
+        self.assertTrue(np.array_equal(y[2, 10:40, :], val))
+
+        # ndim and chunks reflect the reduced view. numpy arrays expose no chunks, so wrap a
+        # small chunked stand-in to check that the squeezed axes' chunks are dropped.
+        class _Chunked:
+            def __init__(self, arr, chunks):
+                self._arr, self.chunks = arr, chunks
+
+            shape = property(lambda self: self._arr.shape)
+            ndim = property(lambda self: self._arr.ndim)
+            dtype = property(lambda self: self._arr.dtype)
+
+            def __getitem__(self, index):
+                return self._arr[index]
+
+        chunked = _Chunked(x, (8, 64, 64))
+        wrapped = RoiWrapper(chunked, np.s_[2, 10:40, :], squeeze=True)
+        self.assertEqual(wrapped.ndim, 2)
+        self.assertEqual(wrapped.chunks, (64, 64))
+        self.assertEqual(len(wrapped.chunks), wrapped.ndim)
+
+        # squeeze=False (default) keeps the singleton axes, ndim, and full chunks.
+        wrapped = RoiWrapper(x, np.s_[0], squeeze=False)
+        self.assertEqual(wrapped.shape, (1,) + self.shape[1:])
+        self.assertEqual(wrapped.ndim, x.ndim)
+        wrapped = RoiWrapper(chunked, np.s_[0], squeeze=False)
+        self.assertEqual(wrapped.ndim, x.ndim)
+        self.assertEqual(wrapped.chunks, (8, 64, 64))
+
 
 if __name__ == '__main__':
     unittest.main()