From 706baf0fb08f88007433fa70c0d8a02d3920bd02 Mon Sep 17 00:00:00 2001 From: lorenzo Date: Fri, 31 Jul 2026 13:58:28 +0200 Subject: [PATCH 1/2] add process lock to dask store ops --- CHANGELOG.md | 6 ++ src/ngio/common/_locks.py | 6 ++ src/ngio/common/_pyramid.py | 10 ++- src/ngio/io_pipes/_ops_slices.py | 8 +- tests/unit/common/test_pyramid.py | 34 +++++++++ tests/unit/io_pipes/test_dask_write_race.py | 83 +++++++++++++++++++++ 6 files changed, 142 insertions(+), 5 deletions(-) create mode 100644 src/ngio/common/_locks.py create mode 100644 tests/unit/io_pipes/test_dask_write_race.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a1ff3a37..ce77d0db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## [Unreleased] + +### Fix + +- Dask writes (`set_array(..., mode="dask")`, masked dask setters, pyramid consolidation) could silently drop data when the patch's block grid was misaligned with the target's chunk grid — or, for sharded targets, its shard grid: `da.store(lock=False)` let two blocks read-modify-write the same chunk concurrently. All dask store flushes now serialise behind a shared lock; block compute stays parallel. + ## [v1.0.0] First stable release. Everything deprecated in `v0.5.0` (each warned "will be removed in `ngio=0.6`") is now removed — that release became `1.0.0`. diff --git a/src/ngio/common/_locks.py b/src/ngio/common/_locks.py new file mode 100644 index 00000000..c0d45e27 --- /dev/null +++ b/src/ngio/common/_locks.py @@ -0,0 +1,6 @@ +from dask.utils import SerializableLock + +# Serialises zarr chunk read-modify-writes issued by da.store across all ngio +# dask write sites; block compute stays parallel. SerializableLock rather than +# threading.Lock so graphs survive pickling under a distributed client. +DASK_STORE_LOCK = SerializableLock() diff --git a/src/ngio/common/_pyramid.py b/src/ngio/common/_pyramid.py index 806a35c3..3a4cfad3 100644 --- a/src/ngio/common/_pyramid.py +++ b/src/ngio/common/_pyramid.py @@ -8,6 +8,7 @@ import zarr from pydantic import BaseModel, ConfigDict, model_validator +from ngio.common._locks import DASK_STORE_LOCK from ngio.common._zoom import ( InterpolationOrder, _zoom_inputs_check, @@ -45,7 +46,10 @@ def _on_disk_dask_zoom( # re-derives chunks via normalize_chunks(chunks="auto", ...) and warns # (treated as error by our filterwarnings) when the result isn't a # multiple of the zarr target's chunks. da.store writes blocks 1:1. - da.store(target_array, target, lock=False) # type: ignore + # The shared lock serialises chunk read-modify-writes: for sharded + # targets the blocks (aligned to inner chunks) share shard objects and + # would otherwise race and silently drop data. + da.store(target_array, target, lock=DASK_STORE_LOCK) # type: ignore def _on_disk_coarsen( @@ -97,8 +101,8 @@ def _on_disk_coarsen( aggregation_function, source_array, coarsening_setup, trim_excess=True ) out_target = out_target.rechunk(target.chunks) - # See _on_disk_dask_zoom for rationale. - da.store(out_target, target, lock=False) # type: ignore + # See _on_disk_dask_zoom for rationale, including the lock. + da.store(out_target, target, lock=DASK_STORE_LOCK) # type: ignore def on_disk_zoom( diff --git a/src/ngio/io_pipes/_ops_slices.py b/src/ngio/io_pipes/_ops_slices.py index aa2fa776..37e01093 100644 --- a/src/ngio/io_pipes/_ops_slices.py +++ b/src/ngio/io_pipes/_ops_slices.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict from ngio.common._dimensions import Dimensions +from ngio.common._locks import DASK_STORE_LOCK from ngio.io_pipes._ops_slices_utils import compute_slice_chunks from ngio.ome_zarr_meta.ngio_specs import Axis from ngio.utils import NgioUserWarning, NgioValueError @@ -255,7 +256,10 @@ def set_slice_as_dask( # da.store instead of da.to_zarr: see ngio.common._pyramid for the # dask>=2025.11 PerformanceWarning regression that to_zarr triggers # when the input chunks aren't a multiple of the target's chunks. - da.store(patch, zarr_array, regions=slice_tuple, lock=False) # type: ignore + # The shared lock serialises chunk read-modify-writes: blocks + # misaligned with the target chunk (or shard) grid would otherwise + # race and silently drop data. + da.store(patch, zarr_array, regions=slice_tuple, lock=DASK_STORE_LOCK) # type: ignore return # Complex case, we have exactly one tuple in the slicing tuple @@ -264,7 +268,7 @@ def set_slice_as_dask( _sub_slice = (*slice_tuple[:ax], slice(idx, idx + 1), *slice_tuple[ax + 1 :]) sub_patch = da.take(patch, indices=i, axis=ax) sub_patch = da.expand_dims(sub_patch, axis=ax) - da.store(sub_patch, zarr_array, regions=_sub_slice, lock=False) # type: ignore + da.store(sub_patch, zarr_array, regions=_sub_slice, lock=DASK_STORE_LOCK) # type: ignore ############################################################## diff --git a/tests/unit/common/test_pyramid.py b/tests/unit/common/test_pyramid.py index 9329be17..891042dd 100644 --- a/tests/unit/common/test_pyramid.py +++ b/tests/unit/common/test_pyramid.py @@ -1,6 +1,7 @@ from pathlib import Path from typing import Literal +import numpy as np import pytest import zarr @@ -28,3 +29,36 @@ def test_on_disk_zooms( target_array = zarr.create_array(target, shape=(16, 64, 64), dtype="uint8") on_disk_zoom(source_array, target_array, order=order, mode=mode) + + +@pytest.mark.parametrize("mode", ["dask", "coarsen"]) +def test_on_disk_zoom_sharded_matches_unsharded( + tmp_path: Path, mode: Literal["dask", "coarsen"] +): + """Dask writes onto a sharded target are content-equal to an unsharded one. + + The dask blocks align to the target's inner chunks while write atomicity + is per shard, so concurrent partial-shard writes race without the shared + `da.store` lock. The unsharded layout is race-free by alignment and + serves as the reference. + """ + rng = np.random.default_rng(0) + source_data = rng.integers(0, 255, size=(16, 128, 128), dtype="uint8") + source_array = zarr.create_array( + tmp_path / "source.zarr", shape=(16, 128, 128), dtype="uint8" + ) + source_array[...] = source_data + + results = {} + for name, shards in (("plain", None), ("sharded", (8, 32, 32))): + target_array = zarr.create_array( + tmp_path / f"target_{mode}_{name}.zarr", + shape=(16, 64, 64), + chunks=(4, 16, 16), + shards=shards, + dtype="uint8", + ) + on_disk_zoom(source_array, target_array, order="nearest", mode=mode) + results[name] = target_array + + np.testing.assert_array_equal(results["sharded"][...], results["plain"][...]) diff --git a/tests/unit/io_pipes/test_dask_write_race.py b/tests/unit/io_pipes/test_dask_write_race.py new file mode 100644 index 00000000..39ccc990 --- /dev/null +++ b/tests/unit/io_pipes/test_dask_write_race.py @@ -0,0 +1,83 @@ +"""Deterministic repro for the `da.store(lock=False)` chunk write race. + +A misaligned dask region write makes two dask blocks read-modify-write the +same zarr chunk. The `GatedStore` below forces both reads of the contested +chunk to complete before either write, so the lost update manifests on every +run — no sleeps, no scheduling luck. +""" + +import asyncio +import contextlib + +import dask +import dask.array as da +import numpy as np +import zarr +import zarr.storage + +from ngio.io_pipes._ops_slices import ( + SlicingOps, + set_slice_as_dask, + set_slice_as_numpy, +) + + +class GatedStore(zarr.storage.WrapperStore): + """Force two concurrent RMWs of one chunk to both read before either writes.""" + + def __init__(self, store: zarr.storage.MemoryStore) -> None: + super().__init__(store) + self.armed = False + self.arrivals = 0 + self._event: asyncio.Event | None = None # created lazily on zarr's I/O loop + + async def get(self, key, prototype, byte_range=None): + if self.armed and key == "c/0": + if self._event is None: + self._event = asyncio.Event() + self.arrivals += 1 + if self.arrivals == 1: + # Wait for the second reader; under a properly locked da.store + # it can never arrive while we hold the chunk, so time out. + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(self._event.wait(), timeout=1.0) + else: + self._event.set() + return await super().get(key, prototype, byte_range) + + +def test_misaligned_dask_write_matches_numpy() -> None: + baseline = np.arange(1000, 1064, dtype="uint16") + patch = np.arange(1, 17, dtype="uint16") + slicing_ops = SlicingOps( + on_disk_axes=("x",), + on_disk_shape=(64,), + on_disk_chunks=(16,), + slicing_tuple=(slice(4, 20),), + ) + + ref_arr = zarr.create_array( + store=zarr.storage.MemoryStore(), shape=(64,), chunks=(16,), dtype="uint16" + ) + ref_arr[:] = baseline + set_slice_as_numpy(zarr_array=ref_arr, patch=patch, slicing_ops=slicing_ops) + + gated = GatedStore(zarr.storage.MemoryStore()) + arr = zarr.create_array(store=gated, shape=(64,), chunks=(16,), dtype="uint16") + arr[:] = baseline + + # Block 0 covers [4:12] (chunk 0 partial); block 1 covers [12:20] + # (chunk 0 partial and chunk 1 partial) -> both RMW chunk "c/0". + gated.armed = True + with dask.config.set(scheduler="threads", num_workers=4): + set_slice_as_dask( + zarr_array=arr, + patch=da.from_array(patch, chunks=(8,)), + slicing_ops=slicing_ops, + ) + gated.armed = False + + # Both pre- and post-fix the contested chunk is read twice; anything less + # means the chunk key drifted and the gate went vacuous. + assert gated.arrivals >= 2 + np.testing.assert_array_equal(arr[:], ref_arr[:]) From 2cbc3ffad3c3c905f8d04b783325f2044ad557b2 Mon Sep 17 00:00:00 2001 From: lorenzo Date: Fri, 31 Jul 2026 14:44:02 +0200 Subject: [PATCH 2/2] ground work for write granularity --- CHANGELOG.md | 19 +++ src/ngio/__init__.py | 2 + src/ngio/images/_abstract_image.py | 9 ++ src/ngio/io_pipes/__init__.py | 2 + src/ngio/io_pipes/_ops_slices_utils.py | 115 ++++++++++++-- src/ngio/iterators/__init__.py | 3 +- src/ngio/iterators/_abstract_iterator.py | 178 ++++++++++++++++++---- src/ngio/iterators/_image_processing.py | 5 + src/ngio/iterators/_mappers.py | 122 +++++++++++---- src/ngio/iterators/_rois_utils.py | 25 ++- src/ngio/iterators/_segmentation.py | 5 + tests/unit/io_pipes/test_chunk_rect.py | 140 +++++++++++++++++ tests/unit/iterators/test_iterators.py | 130 +++++++++++++++- tests/unit/iterators/test_mappers.py | 186 +++++++++++++++++++++++ 14 files changed, 858 insertions(+), 83 deletions(-) create mode 100644 tests/unit/io_pipes/test_chunk_rect.py create mode 100644 tests/unit/iterators/test_mappers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ce77d0db..e2c8e949 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,27 @@ ## [Unreleased] +### API Breaking Changes + +- `MapperProtocol` and `BasicMapper` are re-typed around the new `IterUnit`: a mapper now receives `(func, units)` instead of `(func, getters, setters)` and returns the results as a list ordered by `unit.index`. A unit whose `setter` is `None` is a read-only unit (compute and collect, never write), no longer an error. `map_as_numpy`/`map_as_dask` call sites are unaffected — only custom mapper implementations need updating: + +```python +def __call__(self, func, units): # was (self, func, getters, setters) + return [func(unit.getter()) for unit in units] # was: zip getters/setters, return None +``` + +### Behaviour changes + +- `by_chunks()` gained a `grid="write" | "read"` parameter, defaulting to `"write"`: tiles now align to the output image's write granularity (shard shape when sharded, chunk shape otherwise) instead of the input chunk grid, so the resulting ROIs pass `check_if_chunks_overlap`. Identical for default-derived outputs; different when the output has custom `chunks=`/`shards=`. Read-only iterators fall back to the input grid; `grid="read"` restores the old tiling. + +### Features + +- `reduce_as_numpy`/`reduce_as_dask` on all iterators: apply a function to every ROI and collect the results (any return type, e.g. a `DataFrame` per ROI) without writing anything — `results[i]` corresponds to `iterator.rois[i]` regardless of mapper execution order. +- `IterUnit`: the unit of work a mapper schedules, carrying `index`, `roi`, `getter`, `setter` and `write_footprint` (the chunk/shard rectangle the unit writes — the conflict currency for parallel mappers). + ### Fix +- `check_if_chunks_overlap`/`require_no_chunks_overlap` measured the *read* side: slicing from the getters against the input image's chunk grid. They now measure the write target — slicing from the setters against the output array's write granularity, which is the shard shape when the output is sharded (writes are atomic per shard object). This is a behaviour change: layouts that previously passed the gate (finely chunked input with a coarser-chunked or sharded output) may now correctly raise, and read-only iterators now always report no overlap. - Dask writes (`set_array(..., mode="dask")`, masked dask setters, pyramid consolidation) could silently drop data when the patch's block grid was misaligned with the target's chunk grid — or, for sharded targets, its shard grid: `da.store(lock=False)` let two blocks read-modify-write the same chunk concurrently. All dask store flushes now serialise behind a shared lock; block compute stays parallel. ## [v1.0.0] diff --git a/src/ngio/__init__.py b/src/ngio/__init__.py index e8b6ee78..d11e8850 100644 --- a/src/ngio/__init__.py +++ b/src/ngio/__init__.py @@ -38,6 +38,7 @@ BasicMapper, FeatureExtractorIterator, ImageProcessingIterator, + IterUnit, MapperProtocol, MaskedSegmentationIterator, SegmentationIterator, @@ -82,6 +83,7 @@ "Image", "ImageInWellPath", "ImageProcessingIterator", + "IterUnit", "Label", "MapperProtocol", "MaskedImage", diff --git a/src/ngio/images/_abstract_image.py b/src/ngio/images/_abstract_image.py index d44b48cd..0e938ce7 100644 --- a/src/ngio/images/_abstract_image.py +++ b/src/ngio/images/_abstract_image.py @@ -176,6 +176,15 @@ def chunks(self) -> tuple[int, ...]: """Return the chunks of the image.""" return self.zarr_array.chunks + @property + def write_granularity(self) -> tuple[int, ...]: + """Return the atomic write unit of the on-disk array. + + The shard shape when the array is sharded (writes are read-modify-writes + of whole shard objects), otherwise the chunk shape. + """ + return self.zarr_array.shards or self.zarr_array.chunks + @property def is_3d(self) -> bool: """Return True if the image is 3D.""" diff --git a/src/ngio/io_pipes/__init__.py b/src/ngio/io_pipes/__init__.py index 4fd71fd1..78450376 100644 --- a/src/ngio/io_pipes/__init__.py +++ b/src/ngio/io_pipes/__init__.py @@ -49,9 +49,11 @@ ) from ngio.io_pipes._match_shape import dask_match_shape, numpy_match_shape from ngio.io_pipes._ops_slices import SlicingInputType, SlicingOps, SlicingType +from ngio.io_pipes._ops_slices_utils import ChunkRect from ngio.io_pipes._ops_transforms import TransformProtocol __all__ = [ + "ChunkRect", "DaskGetter", "DaskGetterMasked", "DaskRoiGetter", diff --git a/src/ngio/io_pipes/_ops_slices_utils.py b/src/ngio/io_pipes/_ops_slices_utils.py index c132fada..2b7af2ad 100644 --- a/src/ngio/io_pipes/_ops_slices_utils.py +++ b/src/ngio/io_pipes/_ops_slices_utils.py @@ -176,25 +176,112 @@ def compute_slice_chunks( return {tuple(idx) for idx in product(*per_axis_chunks)} +ChunkRect: TypeAlias = tuple[tuple[int, int], ...] +"""Per-axis inclusive `(first, last)` chunk-index ranges. + +Empty selections are represented as `None` wherever a `ChunkRect` is expected, +never as a `first > last` range: an empty selection touches no chunks and must +conflict with nothing. +""" + + +def _chunk_range_for_axis( + sel: SlicingType, size: int, csize: int +) -> tuple[int, int] | None: + """The inclusive (first, last) chunk-index range touched on a single axis. + + Exact for `slice` and `int` selections (their chunk footprint is contiguous, + see `_chunk_indices_for_axis`). For `list[int]` selections this is the + bounding box — a conservative over-approximation. Returns `None` for an + empty selection. + """ + if isinstance(sel, slice): + start, stop = _normalize_slice(sel, size) + if start >= stop: # empty + return None + return start // csize, (stop - 1) // csize + + if isinstance(sel, int): + if sel < 0 or sel >= size: + raise IndexError(f"index {sel} out of bounds for axis of size {size}") + return sel // csize, sel // csize + + if isinstance(sel, list): + if not sel: + return None + for v in sel: + if not isinstance(v, int): + raise TypeError("Only integers allowed inside tuple selections") + if v < 0 or v >= size: + raise IndexError(f"index {v} out of bounds for axis of size {size}") + return min(sel) // csize, max(sel) // csize + + raise TypeError(f"Unsupported index type: {type(sel)!r}") + + +def compute_chunk_rect( + shape: tuple[int, ...], + chunks: tuple[int, ...], + slicing_tuple: tuple[SlicingType, ...], +) -> ChunkRect | None: + """Bounding rectangle, in chunk-index space, of the chunks a selection touches. + + Exact for `slice`/`int` selections; for `list[int]` selections the per-axis + bounding box is used — a conservative over-approximation that may include + chunks the selection does not touch, never fewer. + + Args: + shape: Overall array shape. + chunks: Chunk (or shard) shape — the conflict granularity. + slicing_tuple: Tuple of slices, ints, or lists of ints. + + Returns: + Per-axis inclusive `(first, last)` chunk-index ranges, or `None` when + the selection is empty on any axis (touches no chunks). + """ + if len(slicing_tuple) != len(shape): + raise NgioValueError( + f"key must have {len(shape)} items, got {len(slicing_tuple)}" + ) + + per_axis_ranges = [] + for sel, size, csize in zip(slicing_tuple, shape, chunks, strict=True): + ax_range = _chunk_range_for_axis(sel, size, csize) + if ax_range is None: + return None + per_axis_ranges.append(ax_range) + return tuple(per_axis_ranges) + + +def chunk_rects_intersect(rect_a: ChunkRect | None, rect_b: ChunkRect | None) -> bool: + """Whether two chunk rectangles share at least one chunk. + + `None` encodes an empty selection and never intersects anything. + """ + if rect_a is None or rect_b is None: + return False + if len(rect_a) != len(rect_b): + raise NgioValueError("Chunk rectangles must have the same rank") + return all( + a_first <= b_last and b_first <= a_last + for (a_first, a_last), (b_first, b_last) in zip(rect_a, rect_b, strict=True) + ) + + def check_if_chunks_overlap( slices: Iterable[tuple[SlicingType, ...]], shape: tuple[int, ...], chunks: tuple[int, ...], ) -> bool: - """Check for overlaps in a list of slicing tuples using brute-force method. + """Check whether any two slicing tuples touch a common chunk. + + This is O(n^2) in the number of slicing tuples, but each pair costs only a + per-axis range comparison — nothing is materialised. `list[int]` selections + are conservatively bounding-boxed: overlap may be reported for lists whose + exact chunk sets are disjoint, never the reverse. - This is O(n^2) and not efficient for large lists. Returns True if any overlaps are found. """ - slices_chunks = (compute_slice_chunks(shape, chunks, si) for si in slices) - for it, (si, sj) in enumerate(_pairs_stream(slices_chunks)): - if si & sj: - return True - if it == 10_000: - warnings.warn( - "Performance Warning check_for_chunks_overlaps is O(n^2) and may be " - "slow for large numbers of regions.", - NgioUserWarning, - stacklevel=2, - ) - return False + rects = (compute_chunk_rect(shape, chunks, si) for si in slices) + non_empty = (rect for rect in rects if rect is not None) + return any(chunk_rects_intersect(ri, rj) for ri, rj in _pairs_stream(non_empty)) diff --git a/src/ngio/iterators/__init__.py b/src/ngio/iterators/__init__.py index 167bc8db..2753c6e6 100644 --- a/src/ngio/iterators/__init__.py +++ b/src/ngio/iterators/__init__.py @@ -2,7 +2,7 @@ from ngio.iterators._feature import FeatureExtractorIterator from ngio.iterators._image_processing import ImageProcessingIterator -from ngio.iterators._mappers import BasicMapper, MapperProtocol +from ngio.iterators._mappers import BasicMapper, IterUnit, MapperProtocol from ngio.iterators._segmentation import ( MaskedSegmentationIterator, SegmentationIterator, @@ -12,6 +12,7 @@ "BasicMapper", "FeatureExtractorIterator", "ImageProcessingIterator", + "IterUnit", "MapperProtocol", "MaskedSegmentationIterator", "SegmentationIterator", diff --git a/src/ngio/iterators/_abstract_iterator.py b/src/ngio/iterators/_abstract_iterator.py index 4071015f..200b7794 100644 --- a/src/ngio/iterators/_abstract_iterator.py +++ b/src/ngio/iterators/_abstract_iterator.py @@ -5,8 +5,17 @@ from ngio.common import Roi from ngio.images._abstract_image import AbstractImage from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol -from ngio.io_pipes._ops_slices_utils import check_if_regions_overlap -from ngio.iterators._mappers import BasicMapper, MapperProtocol +from ngio.io_pipes._ops_slices_utils import ( + _pairs_stream, + check_if_regions_overlap, + chunk_rects_intersect, +) +from ngio.iterators._mappers import ( + BasicMapper, + IterUnit, + MapperProtocol, + compute_write_footprint, +) from ngio.iterators._rois_utils import ( by_chunks, by_yx, @@ -19,6 +28,7 @@ NumpyPipeType = TypeVar("NumpyPipeType") DaskPipeType = TypeVar("DaskPipeType") +R = TypeVar("R") class AbstractIteratorBuilder(ABC, Generic[NumpyPipeType, DaskPipeType]): @@ -53,6 +63,11 @@ def ref_image(self) -> AbstractImage: """Get the reference image for the iterator.""" return self._ref_image + @property + def output_image(self) -> AbstractImage | None: + """The image this iterator writes to, or `None` for a read-only iterator.""" + return None + def _new_from_rois(self, rois: list[Roi]) -> Self: """Create a new instance of the iterator with a different set of ROIs.""" init_kwargs = self.get_init_kwargs() @@ -104,18 +119,39 @@ def by_zyx(self, strict: bool = True) -> Self: rois = by_zyx(self.rois, self.ref_image, strict=strict) return self._new_from_rois(rois) - def by_chunks(self, overlap_xy: int = 0, overlap_z: int = 0) -> Self: - """Return a new iterator that iterates over ROIs by chunks. + def by_chunks( + self, + overlap_xy: int = 0, + overlap_z: int = 0, + grid: Literal["write", "read"] = "write", + ) -> Self: + """Return a new iterator that iterates over ROIs by storage tiles. Args: overlap_xy (int): Overlap in XY dimensions. overlap_z (int): Overlap in Z dimension. + grid: `"write"` (default) sizes the tiles by the output image's + write granularity — the shard shape when the output is sharded, + the chunk shape otherwise — so the resulting ROIs pass + `check_if_chunks_overlap`. Falls back to the input chunk grid + when the iterator is read-only. `"read"` uses the input + image's chunk grid. Returns: - SegmentationIterator: A new iterator with chunked ROIs. + A new iterator with tiled ROIs. """ + if grid == "write": + grid_image = self.output_image + elif grid == "read": + grid_image = None + else: + raise NgioValueError(f"Invalid grid {grid!r}; expected 'write' or 'read'.") rois = by_chunks( - self.rois, self.ref_image, overlap_xy=overlap_xy, overlap_z=overlap_z + self.rois, + self.ref_image, + overlap_xy=overlap_xy, + overlap_z=overlap_z, + grid_image=grid_image, ) return self._new_from_rois(rois) @@ -171,6 +207,30 @@ def _dask_setters_generator( """Return a list of dask setter functions for all ROIs.""" yield from (self.build_dask_setter(roi) for roi in self.rois) + def _numpy_units_generator( + self, with_setters: bool = True + ) -> Generator[IterUnit[NumpyPipeType]]: + """Yield one numpy `IterUnit` per ROI, in ROI order.""" + for index, roi in enumerate(self.rois): + yield IterUnit( + index=index, + roi=roi, + getter=self.build_numpy_getter(roi), + setter=self.build_numpy_setter(roi) if with_setters else None, + ) + + def _dask_units_generator( + self, with_setters: bool = True + ) -> Generator[IterUnit[DaskPipeType]]: + """Yield one dask `IterUnit` per ROI, in ROI order.""" + for index, roi in enumerate(self.rois): + yield IterUnit( + index=index, + roi=roi, + getter=self.build_dask_getter(roi), + setter=self.build_dask_setter(roi) if with_setters else None, + ) + def _read_and_write_generator( self, getters: Generator[ @@ -300,42 +360,89 @@ def iter_as_dask( """Create an iterator over the pixels of the ROIs.""" return self.iter(lazy=False, data_mode="dask", iterator_mode="readwrite") + def _require_writable_units( + self, units: list[IterUnit[NumpyPipeType]] | list[IterUnit[DaskPipeType]] + ) -> None: + """Raise if every unit is read-only: there is nothing for `map` to write.""" + if units and all(unit.setter is None for unit in units): + name = self.__class__.__name__ + raise NgioValueError( + f"{name} is read-only: map has nothing to write back. " + "Use reduce (or iter) to compute without writing." + ) + def map_as_numpy( self, func: Callable[[NumpyPipeType], NumpyPipeType], - mapper: MapperProtocol[NumpyPipeType] | None = None, + mapper: MapperProtocol[NumpyPipeType, NumpyPipeType] | None = None, ) -> None: - """Apply a transformation function to the ROI pixels.""" + """Apply a transformation function to the ROI pixels and write it back.""" if mapper is None: - _mapper = BasicMapper[NumpyPipeType]() + _mapper = BasicMapper[NumpyPipeType, NumpyPipeType]() else: _mapper = mapper - _mapper( - func=func, - getters=self._numpy_getters_generator(), - setters=self._numpy_setters_generator(), - ) + units = list(self._numpy_units_generator()) + self._require_writable_units(units) + _mapper(func, units) self.post_consolidate() def map_as_dask( self, func: Callable[[DaskPipeType], DaskPipeType], - mapper: MapperProtocol[DaskPipeType] | None = None, + mapper: MapperProtocol[DaskPipeType, DaskPipeType] | None = None, ) -> None: - """Apply a transformation function to the ROI pixels.""" + """Apply a transformation function to the ROI pixels and write it back.""" if mapper is None: - _mapper = BasicMapper[DaskPipeType]() + _mapper = BasicMapper[DaskPipeType, DaskPipeType]() else: _mapper = mapper - _mapper( - func=func, - getters=self._dask_getters_generator(), - setters=self._dask_setters_generator(), - ) + units = list(self._dask_units_generator()) + self._require_writable_units(units) + _mapper(func, units) self.post_consolidate() + def reduce_as_numpy( + self, + func: Callable[[NumpyPipeType], R], + mapper: MapperProtocol[NumpyPipeType, R] | None = None, + ) -> list[R]: + """Apply a function to every ROI and collect the results without writing. + + Units are built read-only even on writable iterators: nothing is + written and `post_consolidate` does not run. + + Returns: + One result per ROI, in ROI order: `results[i]` corresponds to + `self.rois[i]`, regardless of the mapper's execution order. + """ + if mapper is None: + _mapper = BasicMapper[NumpyPipeType, R]() + else: + _mapper = mapper + return _mapper(func, list(self._numpy_units_generator(with_setters=False))) + + def reduce_as_dask( + self, + func: Callable[[DaskPipeType], R], + mapper: MapperProtocol[DaskPipeType, R] | None = None, + ) -> list[R]: + """Apply a function to every ROI and collect the results without writing. + + Units are built read-only even on writable iterators: nothing is + written and `post_consolidate` does not run. + + Returns: + One result per ROI, in ROI order: `results[i]` corresponds to + `self.rois[i]`, regardless of the mapper's execution order. + """ + if mapper is None: + _mapper = BasicMapper[DaskPipeType, R]() + else: + _mapper = mapper + return _mapper(func, list(self._dask_units_generator(with_setters=False))) + def check_if_regions_overlap(self) -> bool: """Check if any of the ROIs overlap logically. @@ -362,27 +469,32 @@ def require_no_regions_overlap(self) -> None: raise NgioValueError("Some rois overlap.") def check_if_chunks_overlap(self) -> bool: - """Check if any of the ROIs overlap in terms of chunks. + """Check if any two ROIs write into the same chunk (or shard) of the output. + + Measured on the write target: slicing tuples come from the setters and + the grid is the output array's write granularity — the shard shape when + the output is sharded (writes are atomic per shard object), the chunk + shape otherwise. Two ROIs sharing a write unit make concurrent writes + unsafe: the read-modify-write of that unit can lose data. A read-only + iterator has no write hazard and always returns `False`. - If two ROIs cover the same chunk, they are considered to overlap in chunks. - This does not consider pixel-level overlaps. + This is O(n^2) in the number of ROIs; avoid calling it repeatedly in a + loop. Returns: bool: True if any ROIs overlap in chunks, False otherwise. """ - from ngio.io_pipes._ops_slices_utils import check_if_chunks_overlap - if len(self.rois) < 2: # Less than 2 ROIs cannot overlap return False - slicing_tuples = ( - g.slicing_ops.normalized_slicing_tuple - for g in self._numpy_getters_generator() + footprints = ( + compute_write_footprint(setter) + for setter in self._numpy_setters_generator() + if setter is not None ) - shape = self.ref_image.shape - chunks = self.ref_image.chunks - return check_if_chunks_overlap(slicing_tuples, shape, chunks) + non_empty = (footprint for footprint in footprints if footprint is not None) + return any(chunk_rects_intersect(fi, fj) for fi, fj in _pairs_stream(non_empty)) def require_no_chunks_overlap(self) -> None: """Ensure that the ROIs do not overlap in terms of chunks.""" diff --git a/src/ngio/iterators/_image_processing.py b/src/ngio/iterators/_image_processing.py index 6f365164..90b3e40b 100644 --- a/src/ngio/iterators/_image_processing.py +++ b/src/ngio/iterators/_image_processing.py @@ -86,6 +86,11 @@ def get_init_kwargs(self) -> dict: "output_transforms": self._output_transforms, } + @property + def output_image(self) -> Image: + """The image this iterator writes to.""" + return self._output + def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[np.ndarray]: return NumpyRoiGetter( zarr_array=self._input.zarr_array, diff --git a/src/ngio/iterators/_mappers.py b/src/ngio/iterators/_mappers.py index b6ede464..704f6807 100644 --- a/src/ngio/iterators/_mappers.py +++ b/src/ngio/iterators/_mappers.py @@ -1,48 +1,114 @@ """Mappers for iterators. -Mappers are classes that can be passed to the `map` method of iterators to -transform the items yielded by the iterator. - +Mappers execute a function over the units produced by an iterator and collect +the results. They can be passed to the `map_as_*`/`reduce_as_*` methods of +iterators to customize how the units are scheduled. """ from collections.abc import Callable, Iterable -from typing import Generic, Protocol, TypeVar +from dataclasses import dataclass +from typing import Any, Generic, Protocol, TypeVar, cast +from ngio.common import Roi from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol -from ngio.utils import NgioValueError +from ngio.io_pipes._ops_slices_utils import ChunkRect, compute_chunk_rect T = TypeVar("T") +R = TypeVar("R") + + +def compute_write_footprint(setter: DataSetterProtocol[Any]) -> ChunkRect | None: + """The chunk rectangle a setter will write, at write granularity. + + The granularity is the shard shape when the target array is sharded (writes + are atomic per shard object), the chunk shape otherwise — see + `AbstractImage.write_granularity`. Returns `None` when the setter's + selection is empty. + """ + granularity = setter.zarr_array.shards or setter.slicing_ops.on_disk_chunks + return compute_chunk_rect( + shape=setter.slicing_ops.on_disk_shape, + chunks=granularity, + slicing_tuple=setter.slicing_ops.normalized_slicing_tuple, + ) + + +@dataclass(frozen=True) +class IterUnit(Generic[T]): + """One schedulable unit of iterator work: read one ROI, optionally write back. + + Attributes: + index: Position in `iterator.rois` — mapper results must be returned + in this order. + roi: The ROI this unit covers. + getter: Reads the ROI's data. + setter: Writes the transformed data back, or `None` for a read-only + unit (a read-only iterator, or any `reduce_as_*` call). + """ + + index: int + roi: Roi + getter: DataGetterProtocol[T] + setter: DataSetterProtocol[T] | None + + @property + def write_footprint(self) -> ChunkRect | None: + """The chunk rectangle this unit writes, at write granularity. + + The granularity is the shard shape when the output array is sharded, + the chunk shape otherwise. `None` when the unit is read-only or its + write selection is empty — either way there is nothing to claim and + the unit conflicts with nothing. + """ + if self.setter is None: + return None + return compute_write_footprint(self.setter) + + +class MapperProtocol(Protocol[T, R]): + """Protocol for mappers. + Implementations must honour this contract: -class MapperProtocol(Protocol[T]): - """Protocol for mappers.""" + - The mapper *schedules* each unit's write; it is not required to invoke + `unit.setter` eagerly or on any particular thread, only to guarantee + that all writes are complete when `__call__` returns. + - A unit whose `setter` is `None` is read-only: compute and collect its + result, never write. + - The returned list is ordered by `unit.index` (`results[i]` corresponds + to `iterator.rois[i]`), regardless of execution order. + - `units` is typically a generator and each unit is expensive to build + (store metadata round-trips); generators are not thread-safe. A parallel + mapper must materialise and consume `units` on the dispatching thread + only, then distribute the materialised units to its workers. + """ def __call__( self, - func: Callable[[T], T], - getters: Iterable[DataGetterProtocol[T]], - setters: Iterable[DataSetterProtocol[T] | None], - ) -> None: - """Map an item to another item.""" + func: Callable[[T], R], + units: Iterable[IterUnit[T]], + ) -> list[R]: + """Apply `func` to every unit and return the results in ROI order.""" ... -class BasicMapper(Generic[T]): - """A basic mapper that simply applies a function to the data.""" +class BasicMapper(Generic[T, R]): + """Serial mapper: read, apply, write (if writable), one unit at a time.""" def __call__( self, - func: Callable[[T], T], - getters: Iterable[DataGetterProtocol[T]], - setters: Iterable[DataSetterProtocol[T] | None], - ) -> None: - """Map an item to another item.""" - for getter, setter in zip(getters, setters, strict=True): - data = getter() - data = func(data) - if setter is None: - raise NgioValueError( - "Error in BasicMapper: setter is None, " - "this iterator is read-only, mapping is not possible." - ) - setter(data) + func: Callable[[T], R], + units: Iterable[IterUnit[T]], + ) -> list[R]: + """Apply `func` to every unit and return the results in ROI order.""" + results: list[tuple[int, R]] = [] + for unit in units: + result = func(unit.getter()) + if unit.setter is not None: + # map_as_* is the only writing entry point and its func is + # Callable[[T], T], so T == R whenever setter is not None; + # reduce_as_* always builds units with setter=None. + unit.setter(cast("T", result)) + results.append((unit.index, result)) + results.sort(key=lambda item: item[0]) + return [result for _, result in results] diff --git a/src/ngio/iterators/_rois_utils.py b/src/ngio/iterators/_rois_utils.py index 17eefda2..c6caf3c8 100644 --- a/src/ngio/iterators/_rois_utils.py +++ b/src/ngio/iterators/_rois_utils.py @@ -112,13 +112,26 @@ def by_chunks( overlap_xy: int = 0, overlap_z: int = 0, overlap_t: int = 0, + grid_image: AbstractImage | None = None, ) -> list[Roi]: - """This method is a placeholder for chunked processing.""" - chunk_size = ref_image.chunks - t_axis = ref_image.axes_handler.get_index("t") - z_axis = ref_image.axes_handler.get_index("z") - y_axis = ref_image.axes_handler.get_index("y") - x_axis = ref_image.axes_handler.get_index("x") + """Tile the ROIs on a storage grid. + + By default the tiles are sized by `ref_image`'s chunk grid. When + `grid_image` is given, its write granularity (shard shape when sharded, + chunk shape otherwise) sizes the tiles instead; the ROIs themselves stay + in `ref_image`'s space. An axis present on `ref_image` but absent on + `grid_image` is left un-tiled (one tile spans it) — coarser, never unsafe. + """ + if grid_image is None: + chunk_size = ref_image.chunks + axes_handler = ref_image.axes_handler + else: + chunk_size = grid_image.write_granularity + axes_handler = grid_image.axes_handler + t_axis = axes_handler.get_index("t") + z_axis = axes_handler.get_index("z") + y_axis = axes_handler.get_index("y") + x_axis = axes_handler.get_index("x") size_x = chunk_size[x_axis] if x_axis is not None else None size_y = chunk_size[y_axis] if y_axis is not None else None diff --git a/src/ngio/iterators/_segmentation.py b/src/ngio/iterators/_segmentation.py index c00d14fc..1c2cf5d9 100644 --- a/src/ngio/iterators/_segmentation.py +++ b/src/ngio/iterators/_segmentation.py @@ -79,6 +79,11 @@ def get_init_kwargs(self) -> dict: "output_transforms": self._output_transforms, } + @property + def output_image(self) -> Label: + """The label this iterator writes to.""" + return self._output + def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[np.ndarray]: return NumpyRoiGetter( zarr_array=self._input.zarr_array, diff --git a/tests/unit/io_pipes/test_chunk_rect.py b/tests/unit/io_pipes/test_chunk_rect.py new file mode 100644 index 00000000..85256ac0 --- /dev/null +++ b/tests/unit/io_pipes/test_chunk_rect.py @@ -0,0 +1,140 @@ +import numpy as np +import pytest + +from ngio.io_pipes._ops_slices_utils import ( + check_if_chunks_overlap, + chunk_rects_intersect, + compute_chunk_rect, + compute_slice_chunks, +) +from ngio.utils import NgioValueError + + +def _rect_from_chunk_set(chunk_set: set[tuple[int, ...]]) -> tuple | None: + if not chunk_set: + return None + ndim = len(next(iter(chunk_set))) + return tuple( + (min(c[ax] for c in chunk_set), max(c[ax] for c in chunk_set)) + for ax in range(ndim) + ) + + +def _random_slice_tuple(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple: + slicing = [] + for size in shape: + kind = rng.integers(0, 3) + if kind == 0: + start = int(rng.integers(0, size)) + stop = int(rng.integers(start, size + 1)) + slicing.append(slice(start, stop)) + elif kind == 1: + slicing.append(int(rng.integers(0, size))) + else: + slicing.append(slice(None)) + return tuple(slicing) + + +def test_compute_chunk_rect_matches_sets_for_slices(): + """For slice/int selections the rect must equal the chunk set's bounding box, + and (as the per-axis footprints are contiguous) the rect intersection must + agree exactly with set intersection.""" + rng = np.random.default_rng(0) + for _ in range(200): + ndim = int(rng.integers(1, 4)) + shape = tuple(int(rng.integers(1, 40)) for _ in range(ndim)) + chunks = tuple(int(rng.integers(1, size + 1)) for size in shape) + + tuple_a = _random_slice_tuple(rng, shape) + tuple_b = _random_slice_tuple(rng, shape) + + set_a = compute_slice_chunks(shape, chunks, tuple_a) + set_b = compute_slice_chunks(shape, chunks, tuple_b) + rect_a = compute_chunk_rect(shape, chunks, tuple_a) + rect_b = compute_chunk_rect(shape, chunks, tuple_b) + + assert rect_a == _rect_from_chunk_set(set_a) + assert rect_b == _rect_from_chunk_set(set_b) + assert chunk_rects_intersect(rect_a, rect_b) == bool(set_a & set_b) + + +def test_chunk_rect_list_selection_is_conservative_only(): + """For list selections the rect may over-report overlap, never under-report.""" + rng = np.random.default_rng(1) + shape, chunks = (32,), (4,) + for _ in range(200): + list_a = sorted(rng.choice(32, size=rng.integers(1, 6), replace=False)) + list_b = sorted(rng.choice(32, size=rng.integers(1, 6), replace=False)) + tuple_a = ([int(v) for v in list_a],) + tuple_b = ([int(v) for v in list_b],) + + sets_overlap = bool( + compute_slice_chunks(shape, chunks, tuple_a) + & compute_slice_chunks(shape, chunks, tuple_b) + ) + rects_overlap = chunk_rects_intersect( + compute_chunk_rect(shape, chunks, tuple_a), + compute_chunk_rect(shape, chunks, tuple_b), + ) + if sets_overlap: + assert rects_overlap + + # Pinned over-approximation: sets are disjoint, bounding boxes intersect. + shape, chunks = (16,), (1,) + assert not ( + compute_slice_chunks(shape, chunks, ([0, 8],)) + & compute_slice_chunks(shape, chunks, ([4],)) + ) + assert chunk_rects_intersect( + compute_chunk_rect(shape, chunks, ([0, 8],)), + compute_chunk_rect(shape, chunks, ([4],)), + ) + + +def test_chunk_rect_empty_selection_is_none(): + shape, chunks = (16, 16), (4, 4) + assert compute_chunk_rect(shape, chunks, (slice(5, 5), slice(None))) is None + assert compute_chunk_rect(shape, chunks, (slice(7, 3), slice(None))) is None + assert compute_chunk_rect(shape, chunks, ([], slice(None))) is None + + rect = compute_chunk_rect(shape, chunks, (slice(None), slice(None))) + assert rect == ((0, 3), (0, 3)) + assert chunk_rects_intersect(None, rect) is False + assert chunk_rects_intersect(rect, None) is False + assert chunk_rects_intersect(None, None) is False + + +def test_chunk_rect_validation(): + with pytest.raises(NgioValueError): + compute_chunk_rect((16, 16), (4, 4), (slice(None),)) + with pytest.raises(IndexError): + compute_chunk_rect((16,), (4,), (16,)) + with pytest.raises(IndexError): + compute_chunk_rect((16,), (4,), ([3, 16],)) + with pytest.raises(TypeError): + compute_chunk_rect((16,), (4,), ([3, "x"],)) # ty: ignore[invalid-argument-type] + with pytest.raises(TypeError): + compute_chunk_rect((16,), (4,), ("x",)) # ty: ignore[invalid-argument-type] + with pytest.raises(NgioValueError): + chunk_rects_intersect(((0, 1),), ((0, 1), (0, 1))) + + +def test_check_if_chunks_overlap_rect_semantics(): + shape, chunks = (8, 64), (1, 16) + disjoint = [(0, slice(0, 16)), (0, slice(16, 32)), (1, slice(0, 16))] + assert not check_if_chunks_overlap(disjoint, shape, chunks) + + overlapping = [(0, slice(0, 16)), (0, slice(8, 24))] + assert check_if_chunks_overlap(overlapping, shape, chunks) + + # Empty selections conflict with nothing, even against themselves. + empties = [(0, slice(4, 4)), (0, slice(4, 4)), (0, slice(0, 16))] + assert not check_if_chunks_overlap(empties, shape, chunks) + + +def test_check_if_chunks_overlap_no_warning_at_scale(): + """~20k pairs of disjoint tuples: must stay silent under filterwarnings=error + (the old implementation warned at exactly 10k pairs).""" + shape, chunks = (200,), (1,) + slices = [(slice(i, i + 1),) for i in range(200)] # 19900 pairs, all disjoint + assert not check_if_chunks_overlap(slices, shape, chunks) diff --git a/tests/unit/iterators/test_iterators.py b/tests/unit/iterators/test_iterators.py index 3e257fbf..88091d7e 100644 --- a/tests/unit/iterators/test_iterators.py +++ b/tests/unit/iterators/test_iterators.py @@ -5,7 +5,7 @@ import pytest from zarr.storage import MemoryStore -from ngio import open_ome_zarr_container +from ngio import create_ome_zarr_from_array, open_ome_zarr_container from ngio.iterators import ( FeatureExtractorIterator, ImageProcessingIterator, @@ -152,3 +152,131 @@ def test_features_iterator( feat_iterator = feat_iterator.by_yx() for data, seg, _ in feat_iterator.iter_as_numpy(): assert data.shape == seg.shape + + +def _build_ome_zarr(chunks=(1, 16, 16), ngff_version="0.4"): + rng = np.random.default_rng(0) + array = rng.integers(0, 255, size=(1, 64, 64)).astype("uint8") + return create_ome_zarr_from_array( + store=MemoryStore(), + array=array, + pixelsize=1.0, + axes_names="cyx", + levels=1, + chunks=chunks, + ngff_version=ngff_version, + ) + + +def test_chunk_gate_measures_write_grid(): + """The regression the write-side gate fix targets: ROIs disjoint on the + input chunk grid but sharing chunks of a coarser-chunked output must be + flagged (the getter-sourced gate returned False here).""" + ome_zarr = _build_ome_zarr(chunks=(1, 16, 16)) + # chunks are given at the reference (cyx) rank; the channel axis is squeezed + label = ome_zarr.derive_label("coarse", chunks=(1, 32, 32)) + image = ome_zarr.get_image() + + iterator = SegmentationIterator(image, label, channel_selection=0, axes_order="yx") + iterator = iterator.grid(size_x=16, size_y=16) + + assert iterator.check_if_chunks_overlap() + with pytest.raises(NgioValueError): + iterator.require_no_chunks_overlap() + + # Same tiling against a matching output grid is safe. + matching = ome_zarr.derive_label("matching") + iterator = SegmentationIterator( + image, matching, channel_selection=0, axes_order="yx" + ) + iterator = iterator.grid(size_x=16, size_y=16) + assert not iterator.check_if_chunks_overlap() + iterator.require_no_chunks_overlap() + + +def test_chunk_gate_shard_granularity(): + """ROIs disjoint at inner-chunk granularity but landing in one shard must be + flagged: writes are atomic per shard object.""" + ome_zarr = _build_ome_zarr(chunks=(1, 16, 16), ngff_version="0.5") + image = ome_zarr.get_image() + + sharded = ome_zarr.derive_label("sharded", chunks=(1, 16, 16), shards=(1, 64, 64)) + iterator = SegmentationIterator( + image, sharded, channel_selection=0, axes_order="yx" + ) + iterator = iterator.grid(size_x=16, size_y=16) + assert iterator.check_if_chunks_overlap() + + # The identical layout without sharding is chunk-disjoint and passes. + unsharded = ome_zarr.derive_label("unsharded", chunks=(1, 16, 16)) + iterator = SegmentationIterator( + image, unsharded, channel_selection=0, axes_order="yx" + ) + iterator = iterator.grid(size_x=16, size_y=16) + assert not iterator.check_if_chunks_overlap() + + +def test_by_chunks_grid_parameter(): + """by_chunks(grid="write") tiles on the output's write granularity, so the + resulting ROIs pass the write-side gate; grid="read" restores input tiling.""" + # Default-derived output: grids match, tiling passes either way + ome_zarr = _build_ome_zarr(chunks=(1, 16, 16)) + label = ome_zarr.derive_label("default") + image = ome_zarr.get_image() + iterator = SegmentationIterator(image, label, channel_selection=0, axes_order="yx") + assert not iterator.by_chunks().check_if_chunks_overlap() + + # Coarser-chunked output: write tiling passes, read tiling correctly flags + coarse = ome_zarr.derive_label("coarse", chunks=(1, 32, 32)) + iterator = SegmentationIterator(image, coarse, channel_selection=0, axes_order="yx") + write_tiled = iterator.by_chunks() + assert len(write_tiled.rois) == 4 # 32x32 tiles from the output grid + assert not write_tiled.check_if_chunks_overlap() + read_tiled = iterator.by_chunks(grid="read") + assert len(read_tiled.rois) == 16 # 16x16 tiles from the input grid + assert read_tiled.check_if_chunks_overlap() + + # Sharded output: one tile per shard + ome_zarr_v5 = _build_ome_zarr(chunks=(1, 16, 16), ngff_version="0.5") + image_v5 = ome_zarr_v5.get_image() + sharded = ome_zarr_v5.derive_label( + "sharded", chunks=(1, 16, 16), shards=(1, 64, 64) + ) + iterator = SegmentationIterator( + image_v5, sharded, channel_selection=0, axes_order="yx" + ) + shard_tiled = iterator.by_chunks() + assert len(shard_tiled.rois) == 1 # the single 64x64 shard + assert not shard_tiled.check_if_chunks_overlap() + assert iterator.by_chunks(grid="read").check_if_chunks_overlap() + + with pytest.raises(NgioValueError): + iterator.by_chunks(grid="bogus") # ty: ignore[invalid-argument-type] + + # Read-only iterators fall back to the input grid + feature_iterator = FeatureExtractorIterator( + input_image=image, + input_label=label, + channel_selection=0, + axes_order="yx", + ) + assert len(feature_iterator.by_chunks().rois) == 16 + + +def test_chunk_gate_readonly_returns_false(): + """A read-only iterator has no write hazard: the gate returns False and + require_no_chunks_overlap does not raise, even for chunk-sharing ROIs.""" + ome_zarr = _build_ome_zarr(chunks=(1, 64, 64)) + label = ome_zarr.derive_label("label") + image = ome_zarr.get_image() + + iterator = FeatureExtractorIterator( + input_image=image, + input_label=label, + channel_selection=0, + axes_order="yx", + ) + iterator = iterator.grid(size_x=16, size_y=16) + + assert iterator.check_if_chunks_overlap() is False + iterator.require_no_chunks_overlap() diff --git a/tests/unit/iterators/test_mappers.py b/tests/unit/iterators/test_mappers.py new file mode 100644 index 00000000..48406382 --- /dev/null +++ b/tests/unit/iterators/test_mappers.py @@ -0,0 +1,186 @@ +from collections.abc import Callable, Iterable + +import numpy as np +import pandas as pd +import pytest +from zarr.storage import MemoryStore + +from ngio import create_ome_zarr_from_array +from ngio.io_pipes._ops_slices_utils import chunk_rects_intersect +from ngio.iterators import ( + FeatureExtractorIterator, + IterUnit, + SegmentationIterator, +) +from ngio.utils import NgioValueError + + +def _build_ome_zarr(chunks="auto", ngff_version="0.4"): + rng = np.random.default_rng(0) + array = rng.integers(0, 255, size=(2, 16, 16)).astype("uint8") + return create_ome_zarr_from_array( + store=MemoryStore(), + array=array, + pixelsize=1.0, + axes_names="cyx", + levels=1, + chunks=chunks, + ngff_version=ngff_version, + ) + + +def _build_feature_iterator(ome_zarr) -> FeatureExtractorIterator: + iterator = FeatureExtractorIterator( + input_image=ome_zarr.get_image(), + input_label=ome_zarr.derive_label(name="feature_label"), + channel_selection=0, + axes_order="yx", + ) + return iterator.grid(size_x=8, size_y=8) + + +def _build_segmentation_iterator(ome_zarr, label_name="seg_label", **derive_kwargs): + label = ome_zarr.derive_label(name=label_name, **derive_kwargs) + iterator = SegmentationIterator( + ome_zarr.get_image(), label, channel_selection=0, axes_order="yx" + ) + return iterator.grid(size_x=8, size_y=8), label + + +def test_reduce_as_numpy_roi_order_and_equivalence(): + iterator = _build_feature_iterator(_build_ome_zarr()) + assert len(iterator.rois) == 4 + + results = iterator.reduce_as_numpy( + lambda data: (data[2].name, float(data[0].mean())) + ) + + assert len(results) == len(iterator.rois) + for result, roi in zip(results, iterator.rois, strict=True): + assert result[0] == roi.name + + # Equivalent to the hand-rolled tutorial idiom + expected = [ + (roi.name, float(image.mean())) + for image, _, roi in iterator.iter_as_numpy() + ] + assert results == expected + + +def test_reduce_returns_dataframe(): + iterator = _build_feature_iterator(_build_ome_zarr()) + + def to_features(data) -> pd.DataFrame: + image, label, roi = data + return pd.DataFrame({"roi": [roi.name], "mean": [image.mean()]}) + + results = iterator.reduce_as_numpy(to_features) + assert all(isinstance(result, pd.DataFrame) for result in results) + table = pd.concat(results, ignore_index=True) + assert list(table["roi"]) == [roi.name for roi in iterator.rois] + + +def test_reduce_as_dask(): + iterator = _build_feature_iterator(_build_ome_zarr()) + results = iterator.reduce_as_dask( + lambda data: (data[2].name, float(data[0].mean().compute())) + ) + assert [name for name, _ in results] == [roi.name for roi in iterator.rois] + + +def test_reduce_never_writes(): + iterator, label = _build_segmentation_iterator(_build_ome_zarr()) + before = label.zarr_array[...].copy() + + results = iterator.reduce_as_numpy(lambda patch: np.full_like(patch, 7)) + + assert len(results) == len(iterator.rois) + np.testing.assert_array_equal(label.zarr_array[...], before) + + +def test_map_as_numpy_still_writes_and_returns_none(): + iterator, label = _build_segmentation_iterator(_build_ome_zarr()) + + out = iterator.map_as_numpy(lambda patch: np.full_like(patch, 7)) + + assert out is None + np.testing.assert_array_equal( + label.zarr_array[...], np.full(label.shape, 7, dtype=label.zarr_array.dtype) + ) + + +def test_map_readonly_raises_before_func_runs(): + iterator = _build_feature_iterator(_build_ome_zarr()) + calls: list[int] = [] + + def recording_func(data): + calls.append(1) + return data + + with pytest.raises(NgioValueError, match="read-only"): + iterator.map_as_numpy(recording_func) + with pytest.raises(NgioValueError, match="read-only"): + iterator.map_as_dask(recording_func) + assert calls == [] + + +def test_iter_unit_write_footprint(): + iterator, label = _build_segmentation_iterator(_build_ome_zarr()) + + units = list(iterator._numpy_units_generator()) + assert [unit.index for unit in units] == list(range(len(iterator.rois))) + for unit in units: + footprint = unit.write_footprint + assert footprint is not None + # Rank of the output (yx label), inclusive (first, last) per axis + assert len(footprint) == 2 + assert all(first <= last for first, last in footprint) + + readonly_units = list(iterator._numpy_units_generator(with_setters=False)) + assert all(unit.write_footprint is None for unit in readonly_units) + + feature_units = list( + _build_feature_iterator(_build_ome_zarr())._numpy_units_generator() + ) + assert all(unit.write_footprint is None for unit in feature_units) + + +def test_iter_unit_write_footprint_shard_granularity(): + ome_zarr = _build_ome_zarr(chunks=(2, 8, 8), ngff_version="0.5") + + iterator, _ = _build_segmentation_iterator( + ome_zarr, label_name="sharded", chunks=(2, 8, 8), shards=(2, 16, 16) + ) + footprints = [unit.write_footprint for unit in iterator._numpy_units_generator()] + # All four 8x8 tiles land in the single 16x16 shard + assert all(chunk_rects_intersect(footprints[0], fp) for fp in footprints[1:]) + + iterator, _ = _build_segmentation_iterator( + ome_zarr, label_name="unsharded", chunks=(2, 8, 8) + ) + footprints = [unit.write_footprint for unit in iterator._numpy_units_generator()] + # Without sharding each tile owns its own chunk + assert not any(chunk_rects_intersect(footprints[0], fp) for fp in footprints[1:]) + + +class ShuffledMapper: + """A MapperProtocol implementation that runs units in reverse order.""" + + def __call__(self, func: Callable, units: Iterable[IterUnit]) -> list: + materialized = list(units) + results = [] + for unit in reversed(materialized): + results.append((unit.index, func(unit.getter()))) + results.sort(key=lambda item: item[0]) + return [result for _, result in results] + + +def test_custom_shuffled_mapper_pins_roi_order(): + iterator = _build_feature_iterator(_build_ome_zarr()) + func = lambda data: (data[2].name, float(data[0].mean())) # noqa: E731 + + shuffled = iterator.reduce_as_numpy(func, mapper=ShuffledMapper()) + basic = iterator.reduce_as_numpy(func) + + assert shuffled == basic + assert [name for name, _ in shuffled] == [roi.name for roi in iterator.rois]