Merge branch dev: fix decon output shape on edge tiles

pycudadecon returns a smaller volume for some input sizes, which failed edge
tiles with an unreadable broadcast error from inside zarr. The output is
re-centred to the input shape, and stage_tile now rejects any segmentation
function that changes shape by name.

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)