fix(tiling): 🐛 drop the halo on a one-voxel-thick axis

`tile_shape: "auto"` with `do_3D: false` produces z=1 tiles, so the shipped
`overlap: [4, 30, 30]` expanded every tile to 9 planes to keep 1 -- a 9x read
amplification. Worse, the 2-D guard added in v2.0.0 then refused the block,
because Cellpose without do_3D would read those planes as channels. So `auto`
plus a 2-D method could not run at all.

There is no context to gather along an axis the tile does not span, so
normalize_overlap now takes the tile shape and drops the halo where the tile
is a single voxel thick -- the behaviour the tiling guide already described
("z-tiles of size 1 get depth=0 in z even if you pass overlap=20") but that
stage_tile never implemented. Explicit tiles with room keep their halo, and
omitting tile_shape keeps the previous behaviour.

Also fixes the amplification prepare reports: it did not clip the halo to the
image the way stage_tile does, so a tile spanning the whole image claimed 5x
while reading exactly 1x. That figure is meant to be tuned from.

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

Author: lguerard

src/patchworks/_distributed.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import itertools
+import logging
 from pathlib import Path
 from typing import Callable, Sequence, Union
 
@@ -21,10 +22,16 @@
 from ._io import zarr_compressor_kwargs
 from ._relabel import relabel_sequential_array
 
+logger = logging.getLogger(__name__)
+
 Overlap = Union[int, Sequence[int]]
 
 
-def normalize_overlap(overlap: Overlap, ndim: int) -> tuple[int, ...]:
+def normalize_overlap(
+    overlap: Overlap,
+    ndim: int,
+    tile_shape: "Sequence[int] | None" = None,
+) -> tuple[int, ...]:
     """Expand an overlap spec to one halo width per axis.
 
     A scalar applies the same halo to every axis (the historical behaviour).
@@ -33,12 +40,21 @@ def normalize_overlap(overlap: Overlap, ndim: int) -> tuple[int, ...]:
     ``76 x 1084 x 1084`` to keep ``16 x 1024 x 1024`` -- 5.3x more voxels than
     it uses, nearly all of it in z.
 
+    With *tile_shape*, an axis only one voxel thick gets **no** halo. There is
+    no context to gather along an axis the tile does not span, and a 2-D
+    method handed the extra planes would read them as channels. This is the
+    ``tile_shape: "auto"`` + ``do_3D: false`` case, where tiles come out one
+    plane thick: a z-overlap of 4 would otherwise read 9 planes per tile to
+    keep 1.
+
     Parameters
     ----------
     overlap : int or sequence of int
         Halo width, shared or per-axis.
     ndim : int
         Number of axes the halo is applied to.
+    tile_shape : sequence of int, optional
+        Tile extent per axis. Used to drop halos an axis has no room for.
 
     Returns
     -------
@@ -55,6 +71,21 @@ def normalize_overlap(overlap: Overlap, ndim: int) -> tuple[int, ...]:
             )
     if any(o < 0 for o in values):
         raise ValueError(f"overlap must be non-negative, got {values}")
+
+    if tile_shape is not None:
+        clipped = tuple(
+            0 if int(t) <= 1 else o for o, t in zip(values, tile_shape)
+        )
+        if clipped != values:
+            dropped = [
+                i for i, (a, b) in enumerate(zip(values, clipped)) if a != b
+            ]
+            logger.info(
+                "dropping the halo on axis %s: the tile is 1 voxel thick "
+                "there, so there is no context to read.",
+                dropped,
+            )
+        values = clipped
     return values
 
 
@@ -173,7 +204,7 @@ def stage_tile(
     """
     shape = image.shape
     sl = spatial_tiles(shape, tile_shape)[index]
-    halo = normalize_overlap(overlap, len(sl))
+    halo = normalize_overlap(overlap, len(sl), tile_shape=tile_shape)
     expanded, trims = [], []
     for s, dim, ov in zip(sl, shape, halo):
         lo = max(0, s.start - ov)

tests/test_distributed.py

@@ -41,6 +41,53 @@ def test_normalize_overlap_scalar_and_per_axis():
         normalize_overlap([-1, 0, 0], 3)
 
 
+def test_no_halo_on_a_one_voxel_thick_axis():
+    """A 1-plane tile must get no halo on that axis.
+
+    `tile_shape: "auto"` with `do_3D: false` produces z=1 tiles, so a z-overlap
+    of 4 would read 9 planes to keep 1 (9x) -- and hand a 2-D method a stack it
+    would read as channels.
+    """
+    assert normalize_overlap([4, 30, 30], 3, tile_shape=(1, 8000, 8000)) == (
+        0,
+        30,
+        30,
+    )
+    # A tile with room keeps its halo.
+    assert normalize_overlap([4, 30, 30], 3, tile_shape=(16, 1024, 1024)) == (
+        4,
+        30,
+        30,
+    )
+    # Without a tile_shape nothing is clipped (the historical behaviour).
+    assert normalize_overlap([4, 30, 30], 3) == (4, 30, 30)
+
+
+def test_single_plane_tiles_segment_without_a_z_halo(tmp_path):
+    """End to end: z=1 tiles must not pull neighbouring planes into the tile."""
+    img = np.zeros((4, 32), "uint16")
+    img[1, 4:10] = 500  # only plane 1 has signal
+    tile = (1, 16)
+
+    stage = str(tmp_path / "stage.zarr")
+    create_stage(stage, img.shape, tile)
+    seen = []
+
+    def _recording(block):
+        seen.append(block.shape)
+        return _fn(block)
+
+    for i in range(len(spatial_tiles(img.shape, tile))):
+        stage_tile(img, _recording, stage, i, tile_shape=tile, overlap=[3, 2])
+
+    assert all(s[0] == 1 for s in seen), (
+        f"a z=1 tile was handed multiple planes: {seen}"
+    )
+    written = np.asarray(zarr.open_group(stage, mode="r")["staged"])
+    assert (written[0] == 0).all(), "plane 0 has no signal and must stay empty"
+    assert (written[1, 4:10] > 0).all()
+
+
 def test_auto_overlap_anisotropy_shrinks_z():
     """A 20x coarser z step needs 20x fewer planes for the same distance."""
     assert auto_overlap(30) == 30  # unchanged without voxel_size

workflow/scripts/prepare_tiles.py

@@ -85,17 +85,23 @@
 # A halo wider than the tile itself is not boundary context -- it means every
 # tile re-reads and re-segments its neighbours. Catch it here rather than
 # paying 5x the GPU time for results that get trimmed away.
-overlap = normalize_overlap(cfg.get("overlap", 0), len(tile_shape))
+overlap = normalize_overlap(
+    cfg.get("overlap", 0), len(tile_shape), tile_shape=tile_shape
+)
 for axis, (ov, extent) in enumerate(zip(overlap, tile_shape)):
     if ov >= extent:
         raise ValueError(
             f"overlap[{axis}]={ov} >= tile_shape[{axis}]={extent}: each tile "
             f"would read past its neighbours. Use a per-axis overlap, e.g. "
             f"overlap: {list(max(1, t // 4) for t in tile_shape)}"
         )
-amplification = np.prod(
-    [(t + 2 * o) for t, o in zip(tile_shape, overlap)]
-) / np.prod(tile_shape)
+# Clip the halo to the image the way stage_tile does, or a tile that already
+# spans an axis reports a halo it will never actually read -- a tile covering
+# the whole image would claim 5x while reading exactly 1x.
+read = np.prod(
+    [min(t + 2 * o, s) for t, o, s in zip(tile_shape, overlap, image.shape)]
+)
+amplification = read / np.prod(tile_shape)
 print(f"[patchworks] halo read amplification: {amplification:.2f}x")
 
 tiles = spatial_tiles(image.shape, tile_shape)