fix(dog): 🐛 restore the field of view cudaDecon crops away

An edge tile failed with "could not broadcast input array from shape
(13,1020,1020) into shape (14,1024,1024)" raised six frames deep inside
zarr, which named neither the tile nor the function responsible.

pycudadecon returns a slightly smaller volume for some input sizes. patchworks
needs one label per input voxel: both the halo trim and the destination slice
come from the tile geometry, so a shrunken result has nowhere to go. The decon
output is now re-centred to the input shape, logged at WARNING with both
shapes so a large or asymmetric crop -- which would shift labels rather than
just cost a voxel of halo -- cannot pass silently.

stage_tile also checks the returned shape itself and names the offending
function, so any custom method that changes shape fails with a readable
message instead of a zarr broadcast error.

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

Author: lguerard

src/patchworks/_distributed.py

@@ -213,6 +213,22 @@ def stage_tile(
         trims.append((s.start - lo, hi - s.stop))
     block = np.asarray(image[tuple(expanded)])
     out = np.asarray(fn(block))
+    if out.shape != block.shape:
+        # Caught here rather than 6 frames deep in zarr's codec pipeline as
+        # "could not broadcast input array from shape (13,1020,1020) into
+        # shape (14,1024,1024)", which says nothing about which function is
+        # at fault. A segmentation function must label the voxels it was
+        # given: the halo trim and the destination slice are both computed
+        # from the tile's geometry, so a different shape has no defined
+        # placement.
+        name = getattr(fn, "__name__", type(fn).__name__)
+        raise ValueError(
+            f"segmentation function {name!r} returned shape {out.shape} for "
+            f"a tile of shape {block.shape} (tile {index}). It must return "
+            "one label per input voxel. Some deconvolution backends crop "
+            "their output -- pad or centre it back to the input shape before "
+            "returning."
+        )
     sel = tuple(
         slice(left, out.shape[i] - right)
         for i, (left, right) in enumerate(trims)

src/patchworks/plugins/dog.py

@@ -238,6 +238,61 @@ def _run(block: np.ndarray, dog_dict: dict[str, Any]) -> np.ndarray:
     )
 
 
+def _restore_shape(arr: np.ndarray, shape: tuple[int, ...]) -> np.ndarray:
+    """Centre *arr* back into an array of *shape*, cropping or edge-padding.
+
+    Deconvolution must not change the field of view: patchworks writes the
+    result into a destination slice derived from the tile's geometry, so one
+    label per input voxel is required.
+
+    Centring is the right correction for a symmetric crop, which is what
+    apodisation produces. The discrepancies observed are small (a voxel in z,
+    a few in x/y) and land inside the halo, which is discarded anyway -- so
+    the labels that survive the trim are unaffected. It is logged at WARNING
+    with the exact shapes so a larger, non-symmetric crop cannot pass
+    silently.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        The deconvolved volume.
+    shape : tuple of int
+        The shape it must be returned at (the input tile's).
+
+    Returns
+    -------
+    np.ndarray
+        An array of exactly *shape*.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> _restore_shape(np.ones((13, 1020)), (14, 1024)).shape
+    (14, 1024)
+    """
+    logger.warning(
+        "deconvolution returned %s for a %s input; re-centring to the input "
+        "shape. patchworks needs one label per input voxel. A large or "
+        "asymmetric difference here would shift labels -- check the PSF and "
+        "voxel sizes if this is more than a few voxels.",
+        arr.shape,
+        shape,
+    )
+    # Crop first, so an axis that grew is handled before padding the rest.
+    crop = tuple(
+        slice((a - s) // 2, (a - s) // 2 + s) if a > s else slice(None)
+        for a, s in zip(arr.shape, shape)
+    )
+    arr = arr[crop]
+    pad = tuple(
+        ((s - a) // 2, s - a - (s - a) // 2) if a < s else (0, 0)
+        for a, s in zip(arr.shape, shape)
+    )
+    if any(lo or hi for lo, hi in pad):
+        arr = np.pad(arr, pad, mode="edge")
+    return arr
+
+
 def _segment_once(
     block: np.ndarray, dog_dict: dict[str, Any], use_gpu: bool
 ) -> np.ndarray:
@@ -253,7 +308,16 @@ def _segment_once(
         # in the per-tile timing that tile_process logs.
         from pycudadecon import decon
 
+        before = img.shape
         img = decon(images=img, **decon_kwargs)
+        if img.shape != before:
+            # cudaDecon returns a slightly smaller volume for some input
+            # sizes (e.g. (14,1024,1024) -> (13,1020,1020) on an edge tile).
+            # patchworks needs one label per input voxel: the halo trim and
+            # the destination slice are both derived from the tile geometry,
+            # so a shrunken result has nowhere to go and used to surface as
+            # an unreadable broadcast error from inside zarr.
+            img = _restore_shape(img, before)
 
     if use_gpu:
         import cupy as cp

tests/test_dog.py

@@ -107,3 +107,56 @@ def _decon(images, **kwargs):
     assert captured["dxdata"] == 0.1  # filled from the calibration
     assert captured["dzdata"] == 0.2
     assert captured["dzpsf"] == 0.2
+
+
+def test_restore_shape_recentres_a_cropped_decon():
+    """cudaDecon can hand back a smaller volume than it was given.
+
+    Observed on a real edge tile: (14, 1024, 1024) in, (13, 1020, 1020) out.
+    patchworks needs one label per input voxel, so the field of view has to be
+    restored before the DoG step.
+    """
+    from patchworks.plugins.dog import _restore_shape
+
+    arr = np.arange(13 * 1020 * 1020, dtype="float32").reshape(13, 1020, 1020)
+    out = _restore_shape(arr, (14, 1024, 1024))
+    assert out.shape == (14, 1024, 1024)
+    # The original content is preserved, centred, not resampled.
+    assert np.array_equal(out[0:13, 2:1022, 2:1022], arr)
+
+
+def test_restore_shape_handles_growth_and_exact_fit():
+    """It must be a no-op when shapes already match, and crop when larger."""
+    from patchworks.plugins.dog import _restore_shape
+
+    same = np.ones((4, 8, 8), dtype="float32")
+    assert _restore_shape(same, (4, 8, 8)).shape == (4, 8, 8)
+    bigger = np.ones((6, 12, 12), dtype="float32")
+    assert _restore_shape(bigger, (4, 8, 8)).shape == (4, 8, 8)
+    # And a mix: one axis short, one long.
+    mixed = np.ones((2, 12), dtype="float32")
+    assert _restore_shape(mixed, (4, 8)).shape == (4, 8)
+
+
+def test_stage_tile_rejects_a_shape_changing_function(tmp_path):
+    """A wrong-shaped return must name the culprit, not blow up inside zarr.
+
+    This used to surface as "could not broadcast input array from shape
+    (13,1020,1020) into shape (14,1024,1024)" six frames deep in zarr's codec
+    pipeline, which says nothing about which function misbehaved.
+    """
+    import dask.array as da
+    import pytest
+
+    from patchworks import create_stage, stage_tile
+
+    image = da.zeros((8, 32, 32), chunks=(4, 16, 16), dtype="uint16")
+    stage = str(tmp_path / "stage.zarr")
+    create_stage(stage, image.shape, (4, 16, 16))
+
+    def crops(block):
+        """Stand-in for a deconvolution backend that trims its output."""
+        return np.zeros(tuple(s - 1 for s in block.shape), dtype="int32")
+
+    with pytest.raises(ValueError, match="one label per input voxel"):
+        stage_tile(image, crops, stage, 0, tile_shape=(4, 16, 16), overlap=2)