Merge dev: auto-tile halo fix and calibration-derived decon sizes (v2.1.0)

Fixes a v2.0.0 bug and removes a duplicated calibration.

tile_shape: "auto" with do_3D: false produces z=1 tiles, so the shipped
overlap [4, 30, 30] expanded each tile to 9 planes to keep 1 -- a 9x read
amplification -- and then hit the 2-D guard added in v2.0.0, because Cellpose
without do_3D would read those planes as channels. That combination could not
run at all. An axis the tile does not span has no context to gather, so the
halo is dropped there: the behaviour the tiling guide already described but
stage_tile never implemented. Explicit tiles with room keep their halo.

The amplification figure prepare logs now clips to the image the way
stage_tile does; a tile spanning the whole image used to claim 5x while
reading exactly 1x, and that number exists to be tuned from.

Deconvolution voxel sizes are now derived from the image's own NGFF
calibration -- lateral from X/Y, axial from Z -- instead of being retyped
into the config where the two could drift apart. A deconvolution given the
wrong sampling does not fail, it just returns a subtly wrong result. Values
set explicitly still win, so a PSF sampled differently from the data keeps
its own. The wiring is generic: any custom function declaring a voxel_size
parameter receives the calibration, so a method needing physical units can
ask for them rather than have them configured.

Author: lguerard

docs/api/plugins/ome_zarr.md

@@ -25,6 +25,10 @@ resolution, matching anisotropic microscopy stacks.
 
 ::: patchworks.plugins.ome_zarr.register_labels
 
+## read_pixel_size
+
+::: patchworks.plugins.ome_zarr.read_pixel_size
+
 ## NGFF metadata layout
 
 NGFF 0.4 is defined over zarr v2 and puts its keys at the top level; 0.5 is

docs/examples/dog.md

@@ -65,18 +65,36 @@ Requires `cupy` (matching your CUDA version, e.g. `pip install cupy-cuda12x`)
 
 ## With deconvolution first
 
+Let the voxel sizes come from the image rather than retyping them — the
+lateral ones from X/Y, the axial ones from Z:
+
 ```python
+from patchworks.plugins.ome_zarr import read_pixel_size
+
 fn = dog_label_fn(
     low_sigma=1.0, high_sigma=3.0, threshold=0.02,
-    decon_kwargs=dict(
-        psf=psf, dxpsf=xy_scale, dxdata=xy_scale,
-        dzpsf=z_scale, dzdata=z_scale,
-        wavelength=525, na=1.4, nimm=1.515,
-    ),
+    decon_kwargs=dict(psf=psf, wavelength=525, na=1.4, nimm=1.515),
+    voxel_size=read_pixel_size(IMAGE),   # -> dxdata/dzdata/dxpsf/dzpsf
 )
 result = tile_process(IMAGE, fn, tile_shape=(1, 1024, 1024), overlap=32)
 ```
 
+Anything you set in `decon_kwargs` yourself wins, so a PSF sampled
+differently from the data keeps its own sizes:
+
+```python
+decon_kwargs=dict(psf=psf, dxpsf=0.05, dzpsf=0.1, wavelength=525, ...)
+```
+
+!!! warning "A wrong voxel size does not fail loudly"
+    Deconvolution given the wrong sampling still runs and still returns an
+    image — just a subtly wrong one. That is the reason to derive these from
+    the store's own calibration instead of keeping a second copy in a config
+    that can drift.
+
+    `read_pixel_size` returns `{}` for an uncalibrated store; then nothing is
+    filled in and you must supply the sizes yourself.
+
 !!! note "Deconvolution always needs a GPU"
     `pycudadecon` is CUDA-only, independent of `dog_label_fn`'s own `use_gpu`
     flag (which only picks the backend for the blur/label steps). A SLURM job
@@ -146,15 +164,27 @@ custom:
     threshold: 0.02
     decon_kwargs:
       psf: "/path/to/psf.tif"
-      dxpsf: 0.1
-      dxdata: 0.1
-      dzpsf: 0.2
-      dzdata: 0.2
       wavelength: 525
       na: 1.4
       nimm: 1.515
 ```
 
+No voxel sizes: the workflow reads `image.zarr`'s own calibration and fills
+in `dxdata`/`dxpsf` from X/Y and `dzdata`/`dzpsf` from Z. Set any of them in
+`decon_kwargs` to override — for instance a PSF sampled finer than the data:
+
+```yaml
+      dxpsf: 0.05
+      dzpsf: 0.1
+```
+
+!!! tip "This works for your own methods too"
+    The injection is not DoG-specific. Any `custom` function that declares a
+    `voxel_size` parameter receives `{"z": .., "y": .., "x": ..}` from the
+    store, so a method needing physical units never has to keep a second copy
+    of the calibration in its config. If the store is uncalibrated the
+    workflow says so and passes nothing.
+
 Run it exactly like a Cellpose config:
 
 ```bash

docs/guide/custom_segmentation.md

@@ -118,6 +118,22 @@ def segment(tile: np.ndarray, prob_thresh: float = 0.5) -> np.ndarray:
 
 Using a GPU? Just let your framework see it — nothing extra needed.
 
+!!! tip "Ask for the voxel size instead of configuring it"
+    Declare a `voxel_size` parameter and the workflow passes the image's own
+    NGFF calibration, `{"z": .., "y": .., "x": ..}` in micrometers:
+
+    ```python
+    def segment(tile, *, voxel_size=None, min_diameter_um=5.0):
+        min_px = min_diameter_um / voxel_size["x"]   # µm -> pixels
+        ...
+    ```
+
+    That keeps physical parameters tied to the image rather than duplicated
+    in a config that can drift away from it. Setting `voxel_size` in
+    `custom.kwargs` overrides it; an uncalibrated store passes nothing and
+    logs a warning. This is how the DoG plugin gets its deconvolution voxel
+    sizes.
+
 !!! tip "You do not need to modify patchworks"
     `method: "custom"` imports any `(tile) -> labels` callable, so a new
     method is a module of your own and a config block — nothing in the

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)

src/patchworks/plugins/dog.py

@@ -34,13 +34,16 @@
 
 from __future__ import annotations
 
+import logging
 from functools import partial
 from typing import Any, Callable
 
 import numpy as np
 
 from .._gpu import free_gpu_caches, retry_on_oom
 
+logger = logging.getLogger(__name__)
+
 
 def _require_cupy():
     """Raise an actionable ImportError if cupy is not installed.
@@ -77,13 +80,72 @@ def _require_pycudadecon():
         ) from exc
 
 
+def decon_voxel_kwargs(
+    voxel_size: dict[str, float], psf_voxel_size: dict[str, float] | None = None
+) -> dict[str, float]:
+    """pycudadecon's voxel-size arguments, derived from a calibration.
+
+    ``dxdata``/``dzdata`` describe the **image**, so they come straight from
+    its calibration. ``dxpsf``/``dzpsf`` describe the **PSF file**, which is
+    only the same when the PSF was acquired on the same system at the same
+    sampling -- pass *psf_voxel_size* when it was not.
+
+    Getting these wrong does not fail loudly: deconvolution just returns a
+    subtly wrong result, which is why deriving them beats retyping them.
+
+    Parameters
+    ----------
+    voxel_size : dict
+        Image calibration, e.g. from
+        :func:`patchworks.plugins.ome_zarr.read_pixel_size`. Needs ``x`` (or
+        ``y``) for the lateral size and ``z`` for the axial one.
+    psf_voxel_size : dict, optional
+        The PSF's own calibration. Defaults to *voxel_size*.
+
+    Returns
+    -------
+    dict
+        ``dxdata``/``dzdata``/``dxpsf``/``dzpsf``, omitting any the
+        calibration cannot supply.
+
+    Examples
+    --------
+    >>> decon_voxel_kwargs({"z": 0.2, "y": 0.1, "x": 0.1})
+    {'dxdata': 0.1, 'dzdata': 0.2, 'dxpsf': 0.1, 'dzpsf': 0.2}
+    """
+
+    def _lateral(cal):
+        x, y = cal.get("x"), cal.get("y")
+        if x and y and abs(x - y) > 1e-9 * max(x, y):
+            logger.warning(
+                "anisotropic lateral calibration (x=%s, y=%s); pycudadecon "
+                "takes a single lateral size, using x.",
+                x,
+                y,
+            )
+        return x or y
+
+    psf_cal = psf_voxel_size if psf_voxel_size is not None else voxel_size
+    out: dict[str, float] = {}
+    for key, value in (
+        ("dxdata", _lateral(voxel_size)),
+        ("dzdata", voxel_size.get("z")),
+        ("dxpsf", _lateral(psf_cal)),
+        ("dzpsf", psf_cal.get("z")),
+    ):
+        if value:
+            out[key] = float(value)
+    return out
+
+
 def dog_label_fn(
     low_sigma: float | tuple[float, ...],
     high_sigma: float | tuple[float, ...],
     threshold: float,
     *,
     use_gpu: bool = False,
     decon_kwargs: dict[str, Any] | None = None,
+    voxel_size: dict[str, float] | None = None,
 ) -> Callable[[np.ndarray], np.ndarray]:
     """Return a ready-to-use DoG labeler for ``tile_process``.
 
@@ -108,6 +170,15 @@ def dog_label_fn(
         PSF support when deconvolving, so edge tiles don't get truncated
         context.
 
+        ``dxdata``/``dzdata``/``dxpsf``/``dzpsf`` are filled in from
+        *voxel_size* when you leave them out, so the voxel sizes cannot drift
+        away from the image they describe. Anything you set explicitly wins.
+    voxel_size:
+        Physical voxel size as ``{"z": .., "y": .., "x": ..}``. The Snakemake
+        workflow passes the image's own calibration automatically; from the
+        API, :func:`patchworks.plugins.ome_zarr.read_pixel_size` reads it from
+        a store.
+
     Returns
     -------
     Callable[[ndarray], ndarray]
@@ -117,6 +188,19 @@ def dog_label_fn(
         _require_cupy()
     if decon_kwargs is not None:
         _require_pycudadecon()
+        if voxel_size:
+            # Derived values fill gaps only -- an explicit dxpsf for a PSF
+            # sampled differently from the data must not be overwritten.
+            derived = decon_voxel_kwargs(voxel_size)
+            missing = {
+                k: v for k, v in derived.items() if k not in decon_kwargs
+            }
+            if missing:
+                logger.info(
+                    "decon voxel sizes taken from the image calibration: %s",
+                    missing,
+                )
+                decon_kwargs = {**decon_kwargs, **missing}
     cfg = {
         "low_sigma": low_sigma,
         "high_sigma": high_sigma,

src/patchworks/plugins/ome_zarr.py

@@ -743,6 +743,35 @@ def _write_multiscales(
     write_ngff_attrs(zarr.open_group(group_path, mode="a"), multiscales=[entry])
 
 
+def read_pixel_size(store: Union[str, Path]) -> PixelSize:
+    """Physical voxel size recorded in an OME-ZARR's level-0 metadata.
+
+    The calibration the conversion carried over from the source file, as
+    ``{"z": .., "y": .., "x": ..}`` in micrometers. Axes left at scale 1.0
+    (uncalibrated) are omitted, so an empty dict means the store carries no
+    usable calibration.
+
+    Use this instead of retyping voxel sizes into a config: a deconvolution
+    told the wrong voxel size produces a plausible-looking but wrong result.
+
+    Parameters
+    ----------
+    store : str or Path
+        Path of the OME-ZARR group.
+
+    Returns
+    -------
+    dict
+        ``{axis: size}`` for calibrated spatial axes.
+
+    Examples
+    --------
+    >>> read_pixel_size("scan.zarr")  # doctest: +SKIP
+    {'z': 0.2, 'y': 0.1, 'x': 0.1}
+    """
+    return _read_zarr_calibration(store, "")
+
+
 def _read_zarr_calibration(store: Union[str, Path], axes: str) -> PixelSize:
     """Read level-0 spatial scale from an existing OME-ZARR, if any.