patchworks 2.6.0__tar.gz → 2.6.1__tar.gz

{patchworks-2.6.0 → patchworks-2.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 2.6.0
+Version: 2.6.1
 Summary: Tiled processing of arbitrarily large images with globally consistent labels
 Project-URL: Homepage, https://github.com/imcf/patchworks
 Project-URL: Issues, https://github.com/imcf/patchworks/issues

{patchworks-2.6.0 → patchworks-2.6.1}/docs/guide/snakemake.md

@@ -321,9 +321,10 @@ occupancy map and the staged labels are byte-for-byte what a single-channel
 run produces, and `merge` and `label_relations` need no changes. Two things
 follow from that:
 
-- A tile holds twice the bytes, so a hand-set `tile_shape` sized to fill a GPU
-  may need halving. `tile_shape: "auto"` sizes from the single-channel array
-  and does not yet know about the pair.
+- A tile holds twice the bytes. `tile_shape: "auto"` accounts for this — it is
+  told the tile carries two channels and shrinks each spatial side by ~1/√2,
+  so the tile still fits the same VRAM budget. A **hand-set** `tile_shape`
+  sized to fill a GPU has no such protection and needs halving yourself.
 - The translation is version-specific. Cellpose 3 gets `channels: [1, 2]`
   (1-based into the channel axis, `0` = grayscale); Cellpose 4 (cpsam) dropped
   `channels` entirely and simply reads both. Either is overridable by setting

{patchworks-2.6.0 → patchworks-2.6.1}/src/patchworks/_chunks.py

@@ -264,6 +264,7 @@ def auto_tile_shape(
     gpu_memory: int | None = None,
     available_memory: int | None = None,
     n_workers: int | None = None,
+    n_channels: int = 1,
     verbose: bool = False,
 ) -> tuple[int, ...]:
     """Balanced tile shape for general-purpose 3-D processing.
@@ -287,6 +288,10 @@ def auto_tile_shape(
         Available host RAM in bytes; auto-queried when None.
     n_workers:
         Number of parallel workers (divides the RAM budget).
+    n_channels:
+        Channels each tile carries (default 1). Above 1 the per-voxel cost
+        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
+        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
     verbose:
         Log the chosen shape and estimated tile size.
 
@@ -302,7 +307,12 @@ def auto_tile_shape(
     (128, 512, 512)
     """
     n_workers = n_workers or cpu_allocation()
-    itemsize = np.dtype(dtype).itemsize
+    # A tile holds n_channels planes per voxel (e.g. Cellpose's
+    # cyto+nuclei pair), so the per-voxel cost -- and every budget
+    # derived from it below -- scales with them.
+    if n_channels < 1:
+        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
+    itemsize = np.dtype(dtype).itemsize * n_channels
     n_spatial = min(3, len(shape))
 
     if use_gpu:
@@ -359,6 +369,7 @@ def auto_tile_shape_cellpose(
     n_workers: int | None = None,
     model_memory_bytes: int = 2 * 1024**3,
     cellpose_memory_factor: int = 20,
+    n_channels: int = 1,
     verbose: bool = False,
 ) -> tuple[int, ...]:
     """Cellpose-optimised tile shape.
@@ -391,6 +402,10 @@ def auto_tile_shape_cellpose(
         Memory consumed by the Cellpose model weights (default 2 GiB).
     cellpose_memory_factor:
         Cellpose allocates roughly this multiple of raw input bytes (default 20×).
+    n_channels:
+        Channels each tile carries (default 1). Above 1 the per-voxel cost
+        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
+        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
     verbose:
         Log the chosen shape and memory estimates.
 
@@ -406,7 +421,12 @@ def auto_tile_shape_cellpose(
     (1, 2048, 2048)
     """
     n_workers = n_workers or cpu_allocation()
-    itemsize = np.dtype(dtype).itemsize
+    # A tile holds n_channels planes per voxel (e.g. Cellpose's
+    # cyto+nuclei pair), so the per-voxel cost -- and every budget
+    # derived from it below -- scales with them.
+    if n_channels < 1:
+        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
+    itemsize = np.dtype(dtype).itemsize * n_channels
 
     if use_gpu:
         total_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()

{patchworks-2.6.0 → patchworks-2.6.1}/tests/test_core.py

@@ -389,3 +389,32 @@ def test_channel_selection_respects_the_stores_axes(tmp_path):
     store4 = str(tmp_path / "czyx.zarr")
     to_ome_zarr(vol4, store4, axes="czyx", n_levels=1, progress=False)
     assert load_ome_zarr(store4, channel=1, level=0).shape == (8, 64, 64)
+
+
+def test_auto_tile_shape_charges_for_extra_channels():
+    """A 2-channel tile must fit the same byte budget, not twice it.
+
+    `nuclei_channel` doubles what a tile holds while the tile geometry stays
+    single-channel, so a sizer blind to it hands the GPU a tile needing twice
+    the VRAM it budgeted for.
+    """
+    import pytest
+
+    from patchworks import auto_tile_shape, auto_tile_shape_cellpose
+
+    shape, dtype = (128, 2048, 2048), "uint16"
+
+    one = auto_tile_shape(shape, dtype)
+    two = auto_tile_shape(shape, dtype, n_channels=2)
+    # Same bytes overall: 2 channels of roughly half the area each.
+    assert np.prod(two) * 2 <= np.prod(one)
+    assert np.prod(two) * 2 >= np.prod(one) * 0.9
+
+    kw = dict(diameter=30, do_3D=True, use_gpu=True, gpu_memory=24 * 1024**3)
+    cp_one = auto_tile_shape_cellpose(shape, dtype, **kw)
+    cp_two = auto_tile_shape_cellpose(shape, dtype, n_channels=2, **kw)
+    assert np.prod(cp_two) * 2 <= np.prod(cp_one)
+    assert cp_two[0] == cp_one[0]  # do_3D still pins z to the full extent
+
+    with pytest.raises(ValueError, match="n_channels"):
+        auto_tile_shape(shape, dtype, n_channels=0)

{patchworks-2.6.0 → patchworks-2.6.1}/workflow/config/config_cyto.yaml

@@ -14,9 +14,11 @@ channel: 0 # cytoplasm/membrane channel
 #
 # Only the segment step reads it: the tile geometry, the occupancy map and the
 # staged labels all stay exactly as they are for a single channel -- the pair
-# is stacked on a leading axis that is carried into each tile, not tiled. It
-# does double the bytes a tile holds, so a hand-set `tile_shape` sized to the
-# GPU may need halving.
+# is stacked on a leading axis that is carried into each tile, not tiled.
+#
+# It does double the bytes a tile holds. `tile_shape: "auto"` is told about
+# that and shrinks each spatial side by ~1/sqrt(2) to fit the same VRAM
+# budget; the hand-set tile_shape in common.yaml is not, so halve it yourself.
 #
 # Cellpose 3 turns this into channels: [1, 2]; Cellpose 4 (cpsam) dropped
 # `channels` and just reads both. Override either by setting `channels:` or

{patchworks-2.6.0 → patchworks-2.6.1}/workflow/scripts/prepare_tiles.py

@@ -44,6 +44,10 @@ if ts == "auto":
     # "GPU memory query failed" fallback). None => the built-in 8 GiB default.
     gpu_gb = cfg.get("gpu_memory_gb")
     gpu_bytes = int(gpu_gb * 1024**3) if gpu_gb else None
+    # `image` is single-channel here (the geometry is spatial), but segment
+    # reads nuclei_channel alongside it, so a tile costs twice the bytes.
+    # Without this the sizer would hand the GPU a tile it cannot hold.
+    n_channels = 2 if cfg.get("nuclei_channel") is not None else 1
     if method == "cellpose":
         cp = cfg["cellpose"]
         sizer = partial(
@@ -52,6 +56,7 @@ if ts == "auto":
             use_gpu=cp.get("gpu", True),
             diameter=cp.get("diameter"),
             gpu_memory=gpu_bytes,
+            n_channels=n_channels,
         )
     else:
         # cfg["cellpose"] used to be read unconditionally here, so a DoG or
@@ -59,7 +64,10 @@ if ts == "auto":
         # 'cellpose'. The Cellpose estimator's memory model wouldn't apply to
         # them anyway.
         sizer = partial(
-            auto_tile_shape, use_gpu=gpu_bytes is not None, gpu_memory=gpu_bytes
+            auto_tile_shape,
+            use_gpu=gpu_bytes is not None,
+            gpu_memory=gpu_bytes,
+            n_channels=n_channels,
         )
     tile_shape = tuple(sizer(image.shape, image.dtype))
 else: