perf(merge): ⚡ skip chunks that hold no labels, and size occupancy blocks

The per-tile counts already say which chunks are background, so the merge
need not touch them: an empty chunk is now neither read nor written. Zarr
leaves an unwritten chunk unmaterialised and reads it back as the fill value,
so on a sparse image this saves I/O *and* disk. Likewise a boundary column
next to an empty chunk cannot produce a touching pair, so the scan skips it.
On a one-blob test image that is 1 of 16 chunks relabelled and 0 of 24
boundary columns read.

Doing that exposed why skip_empty was doing nothing on smaller images: the
occupancy block defaulted to (1, 128, 128), which on a 128x128 image is one
block spanning the whole plane, so every tile over-covered it and tested
occupied. The map was right; the resolution was useless. block_for_tile sizes
the block to a quarter of the tile, and prepare/run_multi pass it. The same
test image goes from 16/16 tiles segmented to 1/16.

Two bugs in build_occupancy_map surfaced with it. overwrite=True built the
new map and then threw it away: the guard for "a concurrent run finished
first" fired unconditionally, so the stale map was returned. And a map built
with a different block was reused forever, so changing tile_shape silently
kept the old resolution. It now rebuilds when the block differs, and an
explicit overwrite clears the old directory first.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: lguerard

src/patchworks/__init__.py

@@ -46,7 +46,12 @@
 )
 from ._io import auto_empty_threshold, estimate_empty_tiles, load_ome_zarr
 from ._merge import capped_output_chunks, merge_tile_labels
-from ._occupancy import build_occupancy_map, occupancy_path, tile_occupancy
+from ._occupancy import (
+    block_for_tile,
+    build_occupancy_map,
+    occupancy_path,
+    tile_occupancy,
+)
 from ._postprocess import dilate_labels
 from ._relabel import relabel_sequential_array, relabel_sequential_zarr
 from ._relations import label_relations
@@ -67,6 +72,7 @@
     "load_ome_zarr",
     "estimate_empty_tiles",
     "auto_empty_threshold",
+    "block_for_tile",
     "build_occupancy_map",
     "occupancy_path",
     "tile_occupancy",

src/patchworks/_merge.py

@@ -157,6 +157,7 @@ def _scan_touching_pairs(
     chunk_shape: tuple[int, ...],
     label_offsets: "np.ndarray | None" = None,
     n_workers: int = 1,
+    has_labels: "np.ndarray | None" = None,
 ) -> np.ndarray:
     """Scan chunk-boundary slabs; return (N, 2) int64 array of touching pairs.
 
@@ -183,6 +184,11 @@ def _scan_touching_pairs(
     n_workers : int
         Threads to scan with. Each task reads its own slab, so this is
         embarrassingly parallel and was needlessly serial.
+    has_labels : np.ndarray, optional
+        Per-chunk boolean: does this chunk hold any label? A boundary next to
+        an empty chunk can never produce a touching pair, so those columns are
+        not read at all. On a sparse image (``skip_empty``) that is most of
+        them.
 
     Returns
     -------
@@ -197,15 +203,36 @@ def _scan_touching_pairs(
 
     # One task per chunk-column of one boundary face. Each reads its own slab
     # and produces its own pairs, so they are independent.
-    tasks: list[tuple[int, int, tuple[int, ...]]] = []
+    tasks: list[tuple[int, int, tuple[int, ...], int, int]] = []
+    skipped = 0
     for ax, pos in _boundary_face_specs(shape, chunk_shape):
         face_axes = [a for a in range(arr.ndim) if a != ax]
         face_ranges = [range(0, shape[a], chunk_shape[a]) for a in face_axes]
         for offsets in _iproduct(*face_ranges):
-            tasks.append((ax, pos, offsets))
+            grid = [0] * arr.ndim
+            for a, off in zip(face_axes, offsets):
+                grid[a] = off // chunk_shape[a]
+            grid[ax] = pos // chunk_shape[ax]
+            b_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
+            grid[ax] -= 1  # the chunk on the near side of the boundary
+            a_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
+            # A pair needs a non-zero label on *both* sides, so a boundary
+            # touching an empty chunk cannot produce one -- don't read it.
+            if has_labels is not None and not (
+                has_labels[a_idx] and has_labels[b_idx]
+            ):
+                skipped += 1
+                continue
+            tasks.append((ax, pos, offsets, a_idx, b_idx))
+    if skipped:
+        logger.info(
+            "boundary scan: skipping %d/%d columns that border an empty chunk",
+            skipped,
+            skipped + len(tasks),
+        )
 
-    def _one(task: tuple[int, int, tuple[int, ...]]) -> "np.ndarray | None":
-        ax, pos, offsets = task
+    def _one(task: tuple) -> "np.ndarray | None":
+        ax, pos, offsets, a_idx, b_idx = task
         face_axes = [a for a in range(arr.ndim) if a != ax]
         sl: list = [slice(None)] * arr.ndim
         sl[ax] = slice(pos - 1, pos + 1)
@@ -215,13 +242,6 @@ def _one(task: tuple[int, int, tuple[int, ...]]) -> "np.ndarray | None":
         a_vals = slab[0].ravel().astype(np.int64)
         b_vals = slab[1].ravel().astype(np.int64)
         if label_offsets is not None:
-            grid = [0] * arr.ndim
-            for a, off in zip(face_axes, offsets):
-                grid[a] = off // chunk_shape[a]
-            grid[ax] = pos // chunk_shape[ax]
-            b_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
-            grid[ax] -= 1  # the chunk on the near side of the boundary
-            a_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
             a_vals[a_vals > 0] += label_offsets[a_idx]
             b_vals[b_vals > 0] += label_offsets[b_idx]
         mask = (a_vals > 0) & (b_vals > 0) & (a_vals != b_vals)
@@ -514,6 +534,11 @@ def zarr_native_merge(
         with ``offset`` an exclusive cumulative sum -- O(n_chunks) arithmetic
         instead of reading and rewriting the entire store to renumber it.
         ``None`` falls back to the streaming renumber pass.
+
+        A count of 0 also means that chunk holds no labels, so the merge can
+        skip it outright -- neither reading it nor writing zeros over it. Zarr
+        leaves an unwritten chunk unmaterialised and reads it back as the
+        fill value, so on a sparse image this saves both I/O and disk.
     sequential : bool
         Renumber the merged labels to a contiguous ``1..N``. Free here: the id
         domain is dense by construction, so the compaction is a ``np.unique``
@@ -554,7 +579,9 @@ def zarr_native_merge(
         max_label = int(offsets[-1] + counts_arr[-1]) if n_chunks else 0
     else:
         offsets = None
+        counts_arr = None
         max_label = _make_globally_unique(arr, shape, chunk_shape)
+    has_labels = None if counts_arr is None else counts_arr > 0
     logger.info(
         "zarr_native_merge: shape=%s chunks=%s max_label=%d (%s)",
         shape,
@@ -571,6 +598,7 @@ def zarr_native_merge(
         chunk_shape,
         label_offsets=offsets,
         n_workers=n_workers,
+        has_labels=has_labels,
     )
     logger.info(
         "zarr_native_merge: %d touching pairs → building LUT", len(pairs)
@@ -623,11 +651,22 @@ def zarr_native_merge(
         )
         for idx in _iproduct(*[range(n) for n in n_per_dim])
     ]
+    # A chunk that wrote no labels is all background. Skipping it leaves the
+    # output chunk unwritten, which zarr reads back as the fill value (0) and
+    # never stores -- so an empty region costs neither I/O nor disk.
     tasks = [
         (sl, int(offsets[i]) if offsets is not None else 0)
         for i, sl in enumerate(chunk_slices)
+        if has_labels is None or has_labels[i]
     ]
-    n_w = max(1, min(n_workers, n_chunks))
+    if has_labels is not None and len(tasks) < n_chunks:
+        logger.info(
+            "zarr_native_merge: relabeling %d of %d chunks (%d hold no labels)",
+            len(tasks),
+            n_chunks,
+            n_chunks - len(tasks),
+        )
+    n_w = max(1, min(n_workers, max(1, len(tasks))))
     logger.info(
         "zarr_native_merge: relabeling %d chunks with %d worker(s)…",
         n_chunks,

src/patchworks/_occupancy.py

@@ -30,7 +30,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from itertools import product as _iproduct
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, Sequence, Union
 
 import numpy as np
 import zarr
@@ -40,6 +40,47 @@
 logger = logging.getLogger(__name__)
 
 DEFAULT_BLOCK = (1, 128, 128)
+# Blocks per tile edge. A block is the finest thing the map can resolve, so it
+# has to be well under the tile or every tile over-covers the same block and
+# comes back "occupied" -- correct, but useless as a skip list.
+_BLOCKS_PER_TILE = 4
+
+
+def block_for_tile(
+    tile_shape: Sequence[int], cap: Sequence[int] = DEFAULT_BLOCK
+) -> tuple[int, ...]:
+    """Occupancy block sized so a tile spans several blocks.
+
+    The map can only resolve whole blocks, so a block as large as the tile
+    makes every tile hit the same block and test occupied. Sizing it to a
+    fraction of the tile keeps the answer discriminating, while the *cap*
+    keeps the map small on big tiles.
+
+    Parameters
+    ----------
+    tile_shape : sequence of int
+        The tile shape the map will be queried with.
+    cap : sequence of int, optional
+        Largest block per axis.
+
+    Returns
+    -------
+    tuple of int
+        Block shape, at least 1 per axis.
+
+    Examples
+    --------
+    >>> block_for_tile((16, 1024, 1024))
+    (1, 128, 128)
+    >>> block_for_tile((8, 32, 32))
+    (1, 8, 8)
+    """
+    return tuple(
+        max(1, min(int(c), int(t) // _BLOCKS_PER_TILE or 1))
+        for t, c in zip(tile_shape, cap)
+    )
+
+
 # Bytes to pull per read while pooling; keeps peak memory flat regardless of
 # how large the image is.
 _READ_TARGET_BYTES = 128 * 1024**2
@@ -153,8 +194,25 @@ def build_occupancy_map(
     store = str(image_store)
     out_path = occupancy_path(store, level)
     if not overwrite and Path(out_path).exists():
-        logger.info("occupancy map already present at %s", out_path)
-        return out_path
+        # Reuse only if it was built at the block we want. A map left over
+        # from a run with a different tile_shape is coarser (or finer) than
+        # this run needs, and silently reusing it would degrade every
+        # occupancy answer that follows.
+        try:
+            existing = tuple(zarr.open_array(out_path, mode="r").attrs["block"])
+        except Exception:
+            existing = None
+        if existing == tuple(block):
+            logger.info("occupancy map already present at %s", out_path)
+            return out_path
+        logger.info(
+            "rebuilding occupancy map at %s: it was built with block %s, "
+            "this run needs %s",
+            out_path,
+            existing,
+            tuple(block),
+        )
+        overwrite = True
 
     root = zarr.open_group(store, mode="r")
     src = _level_array(root, level, store)
@@ -229,9 +287,13 @@ def _one(starts: tuple[int, ...]) -> None:
         raise
 
     if Path(out_path).exists():
-        # Another run finished first; theirs is as good as ours.
-        shutil.rmtree(tmp_path, ignore_errors=True)
-        return out_path
+        if not overwrite:
+            # A concurrent run finished first; theirs is as good as ours.
+            shutil.rmtree(tmp_path, ignore_errors=True)
+            return out_path
+        # Explicit rebuild: clear the old one, or os.replace would fail on a
+        # non-empty directory and we would keep serving the stale map.
+        shutil.rmtree(out_path, ignore_errors=True)
     try:
         os.replace(tmp_path, out_path)
     except OSError:
@@ -294,6 +356,16 @@ def tile_occupancy(
             f"tile_shape is {len(tile_shape)}-D but the occupancy map is "
             f"{len(block)}-D"
         )
+    coarse = [
+        (ax, b, t) for ax, (b, t) in enumerate(zip(block, tile_shape)) if b >= t
+    ]
+    if coarse:
+        logger.warning(
+            "occupancy blocks are as large as the tile on axes %s; every tile "
+            "will over-cover the same block and test occupied. Rebuild the "
+            "map with block=block_for_tile(tile_shape).",
+            [ax for ax, _, _ in coarse],
+        )
 
     tile_grid = tuple(-(-s // t) for s, t in zip(sp_shape, tile_shape))
     occupancy = np.zeros(tile_grid, dtype=bool)

tests/test_distributed.py

@@ -1,5 +1,7 @@
 """Tests for the per-tile distributed building blocks."""
 
+from pathlib import Path
+
 import numpy as np
 import pytest
 import zarr
@@ -151,6 +153,45 @@ def test_offsets_match_the_renumber_pass(tmp_path):
     assert set(ids.tolist()) == {1, 2, 3}, f"expected 3 objects, got {ids}"
 
 
+def test_empty_chunks_are_skipped_not_just_short_circuited(tmp_path):
+    """Chunks with no labels must cost neither I/O nor disk.
+
+    A count of 0 already tells the merge the chunk is background, so it need
+    not be read at all -- and leaving the output chunk unwritten means zarr
+    never materialises it.
+    """
+    img = np.zeros((8, 64), "uint16")
+    img[1:4, 1:7] = 500  # wholly inside tile 0; the other seven are empty
+    tile = (8, 8)
+
+    stage = str(tmp_path / "stage.zarr")
+    create_stage(stage, img.shape, tile)
+    counts = {
+        i: stage_tile(img, _fn, stage, i, tile_shape=tile, overlap=0)
+        for i in range(len(spatial_tiles(img.shape, tile)))
+    }
+    assert sum(1 for n in counts.values() if n == 0) == 7, "7 empty tiles"
+
+    out = str(tmp_path / "out.zarr")
+    merged = merge_tile_labels(
+        stage,
+        write_to=out,
+        input_component="staged",
+        sequential_labels=True,
+        label_counts=counts,
+    ).compute()
+
+    # Correctness first: the one real object survives, everything else is 0.
+    assert set(np.unique(merged).tolist()) == {0, 1}
+    assert merged[1:4, 1:7].min() == 1
+
+    # And the empty chunks were never written: zarr stores only the chunks it
+    # was given, so the skipped ones leave no file behind.
+    chunk_dir = Path(out) / "labels" / "c" / "0"
+    written = sorted(p.name for p in chunk_dir.iterdir())
+    assert written == ["0"], f"only the occupied chunk should exist: {written}"
+
+
 def test_stage_tile_returns_dense_label_count(tmp_path):
     """stage_tile reports how many labels it wrote, and writes exactly 1..n."""
     img = np.zeros((8, 16), "uint16")

tests/test_occupancy.py

@@ -5,6 +5,7 @@
 import zarr
 
 from patchworks import (
+    block_for_tile,
     build_occupancy_map,
     estimate_empty_tiles,
     occupancy_path,
@@ -123,6 +124,16 @@ def test_build_is_idempotent_and_covers_ragged_edges(tmp_path):
 
     assert build_occupancy_map(store, block=(1, 8, 8)) == path
 
+    # ...but a *different* block must rebuild rather than serve a stale map,
+    # or changing tile_shape would silently keep the old resolution forever.
+    build_occupancy_map(store, block=(1, 4, 4))
+    assert tuple(zarr.open_array(path, mode="r").attrs["block"]) == (1, 4, 4)
+
+    # And an explicit overwrite must actually take effect: the guard against
+    # a concurrent build used to discard the freshly built map.
+    build_occupancy_map(store, block=(1, 8, 8), overwrite=True)
+    assert tuple(zarr.open_array(path, mode="r").attrs["block"]) == (1, 8, 8)
+
 
 def test_map_lives_outside_the_image_store(tmp_path):
     """The map must not make the image store an invalid zarr hierarchy.
@@ -188,6 +199,33 @@ def _counting(self, key):
     )
 
 
+def test_block_must_be_finer_than_the_tile(tmp_path):
+    """A block as coarse as the tile makes every tile test occupied.
+
+    The map can only resolve whole blocks, so one block spanning the tile is
+    over-covered by every tile -- correct, but useless as a skip list. This is
+    what silently disabled skip_empty on smaller images.
+    """
+    assert block_for_tile((16, 1024, 1024)) == (1, 128, 128)
+    assert block_for_tile((8, 32, 32)) == (1, 8, 8)
+    assert block_for_tile((1, 4, 4)) == (1, 1, 1)  # never below 1
+
+    # A tile-sized block marks everything occupied...
+    data = np.zeros((1, 8, 128, 128), "uint16")
+    data[0, 2:5, 4:20, 4:20] = 900  # one blob in one corner
+    store = _write_store(tmp_path, data)
+    tile = (8, 32, 32)
+
+    build_occupancy_map(store, block=(1, 128, 128))
+    coarse = tile_occupancy(store, tile, channel=0, threshold=100)
+    assert coarse["n_occupied"] == coarse["n_tiles"], "coarse blocks over-cover"
+
+    # ...while a tile-derived block finds only the tile that has signal.
+    build_occupancy_map(store, block=block_for_tile(tile), overwrite=True)
+    fine = tile_occupancy(store, tile, channel=0, threshold=100)
+    assert fine["n_occupied"] == 1, f"expected 1 occupied, got {fine}"
+
+
 def test_tile_shape_dimensionality_is_checked(tmp_path):
     data = np.zeros((1, 2, 16, 16), "uint16")
     store = _write_store(tmp_path, data)