Iterators API reference¶

def with_halo(self, x: int = 0, y: int = 0, z: int = 0, t: int = 0) -> Self:
    """Read a margin of context around each ROI, but write only the ROI.

    The function receives the grown region and must return it grown too;
    the border is cropped before the write, so tiles come out seamless.
    Margins are reference-image pixels, clipped at the borders. Write
    footprints are unchanged, so a haloed iterator parallelizes exactly
    as far as it did without one.

    Note:
        The margins do not survive a rescale: a shape-changing transform
        on the data path (a `ZoomTransform`) makes the crop refuse —
        iterate on a coarser pyramid level instead. A `MaskTransform`
        composes normally (its zoom rescales the label it reads).

    Args:
        x: Pixels added on each side along x.
        y: Pixels added on each side along y.
        z: Pixels added on each side along z.
        t: Frames added on each side along t.

    Returns:
        A new iterator reading with the halo.

    Raises:
        NgioValueError: On a read-only iterator that has not opted in
            (no write to crop the halo from; detection opts in and
            reconciles the overlap with NMS, feature extraction opts in
            and stamps `roi_index`/`roi_name` so your `coalesce` can),
            on an in-place iterator, or for a negative margin.

    Example:
        ```python
        # 8 px of context per side, disjoint writes, still parallel.
        it = iterator.by_write_units().with_halo(x=8, y=8)
        it.map(smooth, mapper=ThreadedMapper("auto"))
        ```
    """
    halo = {"x": x, "y": y, "z": z, "t": t}
    for axis_name, margin in halo.items():
        if margin < 0:
            raise NgioValueError(
                f"Halo along '{axis_name}' must be >= 0, got {margin}."
            )
    if self.output_image is None and not self._allow_readonly_halo:
        name = self.__class__.__name__
        raise NgioValueError(
            f"{name} is read-only, so there is no written region for a "
            "halo to surround. Widen the ROIs themselves instead (e.g. "
            "`by_grid(...)` with a stride smaller than the size)."
        )
    output = self.output_image
    if output is not None and _is_same_zarr_array(
        self._ref_image.zarr_array, output.zarr_array
    ):
        raise NgioValueError(
            "In-place iteration cannot take a halo: each tile's grown "
            "read covers pixels that neighbouring tiles write, so the "
            "result depends on execution order. Write to a different "
            "image or label, or drop the halo."
        )
    new_instance = self._new_from_rois(self.rois)
    new_instance._halo = {ax: m for ax, m in halo.items() if m > 0}
    return new_instance

def by_grid(
    self,
    *,
    size_x: int | None = None,
    size_y: int | None = None,
    size_z: int | None = None,
    size_t: int | None = None,
    stride_x: int | None = None,
    stride_y: int | None = None,
    stride_z: int | None = None,
    stride_t: int | None = None,
    tail: TailPolicy = "clip",
    base_name: str = "",
) -> Self:
    """Tile the current ROIs with a regular grid.

    Args:
        size_x: Tile size along x; `None` (default) means the full extent.
        size_y: Tile size along y; `None` means the full extent.
        size_z: Tile size along z; `None` means the full extent.
        size_t: Tile size along t; `None` means the full extent.
        stride_x: Step between tiles along x; `None` means `size_x`
            (adjacent tiles). A smaller stride produces overlapping tiles.
        stride_y: As `stride_x`, along y.
        stride_z: As `stride_x`, along z.
        stride_t: As `stride_x`, along t.
        tail: What happens where the grid does not divide the axis.
            `"clip"` (default) shrinks the last tile to the border;
            `"balance"` re-splits the last full tile and the tail into two
            near-equal tiles, so a thin overhang never yields a thin tile
            (needs adjacent tiles, `stride == size`); `"shift"` moves the
            last tile back to stay full-size — it then overlaps its
            neighbour, so a writing iterator needs a merge or serial
            execution; `"drop"` discards the tail and leaves that border
            uncovered.
        base_name: Prefix for the generated tile names.

    Returns:
        A new iterator over the tiled ROIs.
    """
    rois = by_grid(
        rois=self.rois,
        ref_image=self.ref_image,
        size_x=size_x,
        size_y=size_y,
        size_z=size_z,
        size_t=size_t,
        stride_x=stride_x,
        stride_y=stride_y,
        stride_z=stride_z,
        stride_t=stride_t,
        tail=tail,
        base_name=base_name,
    )
    return self._new_from_rois(rois)

def by_blocks(
    self,
    *,
    num_x: int = 1,
    num_y: int = 1,
    num_z: int = 1,
    num_t: int = 1,
    base_name: str = "",
) -> Self:
    """Divide each axis into a fixed number of near-equal blocks.

    The complement of `by_grid`: you say how many tiles, not how big.
    Block lengths along an axis differ by at most one pixel, so there is
    no tail to police — the partition is balanced by construction.

    Args:
        num_x: Number of blocks along x (default 1: no split).
        num_y: Number of blocks along y.
        num_z: Number of blocks along z.
        num_t: Number of blocks along t.
        base_name: Prefix for the generated tile names.

    Returns:
        A new iterator over the blocks.
    """
    rois = by_blocks(
        rois=self.rois,
        ref_image=self.ref_image,
        num_x=num_x,
        num_y=num_y,
        num_z=num_z,
        num_t=num_t,
        base_name=base_name,
    )
    return self._new_from_rois(rois)

def by_yx(self) -> Self:
    """Split each ROI into one 2D plane per remaining coordinate.

    A broadcasting call, not a tiling: every ROI becomes one ROI per
    `(t, z, c)` combination, each covering the full y/x extent — the
    shape for "run this 2D function on every plane".
    """
    rois = by_yx(self.rois, self.ref_image)
    return self._new_from_rois(rois)

def by_zyx(self, strict: bool = True) -> Self:
    """Split each ROI into one 3D volume per remaining coordinate.

    A broadcasting call, not a tiling: every ROI becomes one ROI per
    `(t, c)` combination, each covering the full z/y/x extent.

    Args:
        strict: Require a real z axis (present and larger than 1);
            with `False`, images without one fall back to 2D planes.
    """
    rois = by_zyx(self.rois, self.ref_image, strict=strict)
    return self._new_from_rois(rois)

def by_chunks(
    self,
    *,
    overlap_x: int = 0,
    overlap_y: int = 0,
    overlap_z: int = 0,
    overlap_t: int = 0,
) -> Self:
    """Tile the ROIs by the input image's chunk grid.

    Reads are chunk-granular, so this is the natural tiling for
    read-heavy work. For collision-free parallel *writes* use
    `by_write_units`, which tiles by the output's write granularity.

    Args:
        overlap_x: Overlap between adjacent tiles along x, in pixels.
        overlap_y: Overlap along y.
        overlap_z: Overlap along z.
        overlap_t: Overlap along t.

    Returns:
        A new iterator with tiled ROIs.
    """
    rois = by_storage_units(
        self.rois,
        self.ref_image,
        overlap_x=overlap_x,
        overlap_y=overlap_y,
        overlap_z=overlap_z,
        overlap_t=overlap_t,
        grid_image=None,
    )
    return self._new_from_rois(rois)

def by_write_units(
    self,
    *,
    overlap_x: int = 0,
    overlap_y: int = 0,
    overlap_z: int = 0,
    overlap_t: int = 0,
) -> Self:
    """Tile the ROIs by the output image's write granularity.

    The write unit is the shard shape when the output is sharded (writes
    are atomic per shard object), the chunk shape otherwise. With no
    overlap the resulting ROIs pass `check_if_write_units_overlap` by
    construction, so a parallel `map` runs as a single fully-parallel
    wave — a throughput optimization; conflicting tilings are still safe,
    just scheduled in more waves. Falls back to the input chunk grid when
    the iterator is read-only.

    Args:
        overlap_x: Overlap between adjacent tiles along x, in pixels.
        overlap_y: Overlap along y.
        overlap_z: Overlap along z.
        overlap_t: Overlap along t.

    Returns:
        A new iterator with tiled ROIs.
    """
    rois = by_storage_units(
        self.rois,
        self.ref_image,
        overlap_x=overlap_x,
        overlap_y=overlap_y,
        overlap_z=overlap_z,
        overlap_t=overlap_t,
        grid_image=self.output_image,
    )
    return self._new_from_rois(rois)

def product(self, other: list[Roi] | GenericRoiTable) -> Self:
    """Cartesian product of the current ROIs with an arbitrary list of ROIs."""
    if isinstance(other, GenericRoiTable):
        other = other.rois()
    rois = rois_product(self.rois, other)
    return self._new_from_rois(rois)

def iter(  # ty: ignore[invalid-method-override]
    self,
    lazy: bool = False,
    data_mode: Literal["numpy", "dask"] | None = None,
    # "readwrite" by design — the one deliberate LSP variance of the
    # reader/writer split.
    iterator_mode: Literal["readwrite", "readonly"] = "readwrite",
    *,
    batch_size: int | None = None,
) -> Generator:
    """Create an iterator over the pixels of the ROIs.

    A writing loop finalizes only when fully drained — on a stitching
    iterator that finalize is the resolve, so abandoning the loop
    mid-way leaves block-offset ids on disk (plus the transient scratch)
    until a re-run or a later `finalize()`. To write now and gather
    later on purpose, iterate a `for_job(0, 1)` slice and call
    `finalize()` on the unrestricted iterator when ready.

    With `batch_size` set, a writing loop yields `(patches, writers)` —
    two aligned lists of up to `batch_size` items, in ROI order — and a
    read-only loop yields the payload lists alone. Stack the patches
    yourself (raggedness is yours to handle, unlike `BatchedMapper`'s
    automatic padding), run the model once, and hand each result to its
    writer. Batches follow ROI order — under `write_order="roi"` that
    is the same later-ROI-wins order the mappers schedule, so the
    manual loop is bit-identical to `map`. Batching is numpy-only
    and eager: combining `batch_size` with `lazy=True` or
    `data_mode="dask"` raises.

    Args:
        lazy: Yield getter/setter handles instead of materialized data.
        data_mode: `"numpy"` or `"dask"` (deprecated, see note).
        iterator_mode: `"readwrite"` yields `(patch, writer)` pairs,
            `"readonly"` yields patches alone.
        batch_size: When set, yield batches of up to this many items;
            the last batch may be smaller.

    Example:
        ```python
        for patches, writers in iterator.iter(data_mode="numpy", batch_size=8):
            outs = model(np.stack(patches))
            for writer, out in zip(writers, outs):
                writer(out)
        ```

    Note:
        The dask data mode is deprecated and will be removed in ngio=1.2;
        from then on `iter()` yields numpy arrays.
    """
    return self._iter_impl(
        lazy=lazy,
        data_mode=data_mode,
        iterator_mode=iterator_mode,
        batch_size=batch_size,
    )

def iter_as_numpy(
    self,
):
    """Alias for `iter(data_mode="numpy")`."""
    return self._iter(lazy=False, data_mode="numpy", iterator_mode="readwrite")

@deprecated(
    replacement="iter_as_numpy() (or Image.get_as_dask() for a lazy array)",
    removed_in="1.2",
)
def iter_as_dask(
    self,
):
    """Alias for `iter(data_mode="dask")`.

    Deprecated: removed in ngio=1.2.
    """
    return self._iter(lazy=False, data_mode="dask", iterator_mode="readwrite")

def for_job(self, job_index: int, n_jobs: int) -> Self:
    """Restrict the iterator to one of `n_jobs` independent partitions.

    No write unit is ever shared between two partitions, so the jobs
    need no locks, no barriers, and no channel to one another — safe as
    separate SLURM array tasks, in any order. Every job must build an
    identical iterator (construction is metadata-only and the partition
    derives deterministically from it), apply `for_job` *last* in the
    builder chain, and its verbs deliberately do not finalize — a slice's
    `map`/`process`/`segment` writes only its share, and `measure`/
    `detect` bank a partial instead of joining. The gather step is the
    unrestricted iterator's `finalize()`, once, after all jobs.
    Stitched and read-only runs also need `prepare_jobs(n_jobs)` first.
    See the distributed-runs guide in the iterators docs.

    Args:
        job_index: This job's partition, `0 <= job_index < n_jobs`
            (e.g. `int($SLURM_ARRAY_TASK_ID)`). Partitions beyond the
            number of independent unit groups are empty no-ops.
        n_jobs: How many partitions the run is split into; must match
            across every job and the gather.

    Returns:
        The restricted iterator; further reshaping refuses.

    Example:
        ```python
        # array task:
        iterator = SegmentationIterator(image, label).by_chunks()
        iterator.for_job(job_index, n_jobs=n_jobs).map(func)

        # gather task, after all jobs:
        SegmentationIterator(image, label).by_chunks().finalize()
        ```
    """
    if self._partition is not None:
        _, current_index, current_n = self._partition
        raise NgioValueError(
            f"This iterator is already restricted to partition "
            f"{current_index}/{current_n}; partitions do not nest. "
            "Apply `for_job` once, to the full iterator."
        )
    if n_jobs < 1:
        raise NgioValueError(f"n_jobs must be >= 1, got {n_jobs}.")
    if not 0 <= job_index < n_jobs:
        raise NgioValueError(
            f"job_index must be in [0, {n_jobs}), got {job_index}."
        )
    self._require_splittable()
    self._validate_job_context(n_jobs)
    units = list(self._numpy_units_generator())
    partition = partition_components(write_conflict_components(units), n_jobs)
    restricted = self._new_from_rois(self.rois)
    restricted._partition = (
        frozenset(partition[job_index]),
        job_index,
        n_jobs,
    )
    return restricted

def prepare_jobs(self, n_jobs: int) -> list[JobArgs]:
    """Set up a distributed run and return its parallelization list.

    The init step of the three-phase recipe (init → parallel jobs →
    consolidate): performs the iterator's setup — always clearing stale
    scratch state from earlier runs first — and returns one `JobArgs`
    per non-empty partition. Optional for plain writers; a stitching or
    read-only run requires it (the scratch root must exist, race-free,
    before any job writes into it).

    Args:
        n_jobs: How many partitions to split into; empty ones are
            dropped from the list.

    Returns:
        One `{"job_index": ..., "n_jobs": ...}` dict per non-empty
        partition — JSON-ready, splatting into `for_job(**args)`.

    Example:
        ```python
        args_list = iterator.prepare_jobs(n_jobs=4)  # init task
        iterator.for_job(**args).map(func)  # per parallel task
        iterator.finalize()  # consolidate task
        ```
    """
    if self._partition is not None:
        raise NgioValueError(
            "prepare_jobs belongs to the unrestricted iterator; this one "
            "is already restricted to a partition."
        )
    if n_jobs < 1:
        raise NgioValueError(f"n_jobs must be >= 1, got {n_jobs}.")
    self._require_splittable()
    if self.output_image is not None:
        self._validate_write_plan()
    units = list(self._numpy_units_generator())
    partition = partition_components(write_conflict_components(units), n_jobs)
    self._prepare_distributed(n_jobs)
    return [
        JobArgs(job_index=job_index, n_jobs=n_jobs)
        for job_index, indices in enumerate(partition)
        if indices
    ]

def reduce(
    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 `finalize` does not run.

    Args:
        func: The function to apply; under a parallel mapper it must be
            safe on worker threads (or processes).
        mapper: How the units are scheduled. `None` (the default) is
            serial; pass `ThreadedMapper()` or `ProcessMapper()` to fan
            out, sized by their own `max_workers` argument.

    Returns:
        One result per ROI, in ROI order: `results[i]` corresponds to
        `self.rois[i]`, regardless of the mapper's execution order.
    """
    self._require_no_halo_on_readonly_verbs()
    if mapper is None:
        mapper = BasicMapper[NumpyPipeType, R]()
    return mapper(func, list(self._numpy_units_generator(with_setters=False)))

def reduce_as_numpy(
    self,
    func: Callable[[NumpyPipeType], R],
    *,
    mapper: MapperProtocol[NumpyPipeType, R] | None = None,
) -> list[R]:
    """Alias for `reduce()`."""
    return self.reduce(func, mapper=mapper)

@deprecated(
    replacement="reduce() / reduce_as_numpy(mapper=...)",
    removed_in="1.2",
)
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.

    Deprecated: removed in ngio=1.2.

    Units are built read-only even on writable iterators: nothing is
    written and `finalize` does not run. Runs serially — the units hand
    `func` *lazy* dask arrays, so the per-unit work is graph
    construction, not IO. A parallel `mapper` raises (see `map_as_dask`).

    Returns:
        One result per ROI, in ROI order: `results[i]` corresponds to
        `self.rois[i]`, regardless of the mapper's execution order.
    """
    self._require_no_halo_on_readonly_verbs()
    _mapper = self._require_serial_dask_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.

    If two ROIs cover the same pixel, they are considered to overlap.
    This does not consider chunking or other storage details. Measured
    on the *read* regions — with a halo, neighbouring reads overlap by
    construction. For write safety, `check_if_write_units_overlap` is
    the right question.

    Returns:
        `True` if any ROIs overlap.
    """
    if len(self.rois) < 2:
        return False

    slicing_tuples = (
        g.slicing_ops.normalized_slicing_tuple
        for g in self._numpy_getters_generator()
    )
    return check_if_regions_overlap(slicing_tuples)

def require_no_regions_overlap(self) -> None:
    """Ensure that the Iterator's ROIs do not overlap."""
    if self.check_if_regions_overlap():
        raise NgioValueError("Some rois overlap.")

def map(
    self,
    func: Callable[[NumpyPipeType], NumpyPipeType],
    *,
    mapper: MapperProtocol[NumpyPipeType, NumpyPipeType] | None = None,
) -> None:
    """Apply a transformation function to each ROI's patch and write it back.

    Args:
        func: The transformation. Under a parallel mapper it runs on
            worker threads (or processes) and must be safe there.
        mapper: How the units are scheduled. `None` (the default) is
            serial — parallel writes stay explicit opt-in. Pass
            `ThreadedMapper()` or `ProcessMapper()` to fan out; each
            sizes its own pool from its `max_workers` argument, and both
            schedule the units into conflict-free waves (`plan_waves`).
    """
    if mapper is None:
        mapper = BasicMapper[NumpyPipeType, NumpyPipeType]()
    units = list(self._numpy_units_generator())
    self._validate_write_plan()
    mapper(func, units)
    # A partition slice runs only its share; the finalize (the one global
    # step) belongs to the unrestricted iterator, once every job is done.
    if self._partition is None:
        self.finalize()

def map_as_numpy(
    self,
    func: Callable[[NumpyPipeType], NumpyPipeType],
    *,
    mapper: MapperProtocol[NumpyPipeType, NumpyPipeType] | None = None,
) -> None:
    """Alias for `map()`."""
    return self.map(func, mapper=mapper)

@deprecated(
    replacement="map() / map_as_numpy(mapper=...)",
    removed_in="1.2",
)
def map_as_dask(
    self,
    func: Callable[[DaskPipeType], DaskPipeType],
    *,
    mapper: MapperProtocol[DaskPipeType, DaskPipeType] | None = None,
) -> None:
    """Apply a transformation function to each ROI's patch and write it back.

    Deprecated: removed in ngio=1.2.

    Runs serially: the dask pipes are already executed by dask's own
    scheduler, and the dask write path (`store_dask`) scopes a global
    dask config option, so concurrent callers are unsafe. A parallel
    `mapper` raises.
    """
    _mapper = self._require_serial_dask_mapper(mapper)
    units = list(self._dask_units_generator())
    self._validate_write_plan()
    _mapper(func, units)
    if self._partition is None:
        self.finalize()

def check_if_write_units_overlap(self) -> bool:
    """Check if any two ROIs write into the same write unit 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. The parallel
    mappers schedule such ROIs into separate waves, so this check answers
    "will my map run as a single fully-parallel wave?".

    This is O(n^2) in the number of ROIs; avoid calling it repeatedly in a
    loop.

    Returns:
        `True` if any two ROIs share a write unit.
    """
    if len(self.rois) < 2:
        return False

    footprints = (
        compute_write_footprint(setter)
        for setter in self._numpy_setters_generator()
        if setter is not None
    )
    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_write_units_overlap(self) -> None:
    """Ensure that the ROIs do not share write units on the output.

    The strict opt-in gate: the parallel mappers no longer refuse shared
    write units on their own (they wave-schedule around them), so call
    this to insist on a single-wave tiling instead.
    """
    if self.check_if_write_units_overlap():
        raise NgioValueError("Some ROIs share write units on the output.")