patchworks 2.6.4__tar.gz → 2.6.6__tar.gz

{patchworks-2.6.4 → patchworks-2.6.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: patchworks
-Version: 2.6.4
+Version: 2.6.6
 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.4 → patchworks-2.6.6}/docs/guide/label_relations.md

@@ -6,7 +6,21 @@ the other — by streaming both arrays chunk by chunk, so it scales to
 hundreds of thousands of objects without loading anything fully into RAM.
 
 Both label arrays must share the exact same chunk layout — same
-`tile_shape`/pyramid `level` when they were produced.
+`tile_shape`/pyramid `level` when they were produced. `label_relations()`
+raises rather than silently doing something slower if they don't; rechunk
+one side to match first (same shape, different chunking is a normal dask
+op — some extra I/O reading across misaligned source chunks, not a
+correctness issue):
+
+```python
+if nuclei.chunks != cells.chunks:
+    cells = cells.rechunk(nuclei.chunks)
+```
+
+(The cluster workflow's `run_multi.py`/`relate.py` does this automatically
+— see below — so two configs segmented at different `tile_shape` still
+relate correctly; you only need to do it by hand when calling
+`label_relations()` directly.)
 
 ```python
 import dask.array as da

{patchworks-2.6.4 → patchworks-2.6.6}/docs/guide/snakemake.md

@@ -398,11 +398,22 @@ results/image.zarr/labels/nuclei_labels/
 results/image.zarr/labels/cyto_labels/
 ```
 
-!!! tip "Keep `tile_shape` (and `level`) identical across configs"
-    Different segmentations of the same image can use different `channel` and
-    `cellpose:` settings freely, but keep `tile_shape`/`level` the same across
-    configs — the label arrays then share the exact same chunk layout, which
-    [`label_relations()`](label_relations.md) requires.
+!!! tip "Keep `level` identical across configs; `tile_shape` doesn't have to match anymore"
+    Different segmentations of the same image can use different `channel`
+    and `cellpose:` settings freely, but keep `level` the same across
+    configs so the label arrays cover the same voxel grid.
+
+    `tile_shape` itself no longer has to match. Two configs naturally
+    produce different on-disk chunk layouts when `tile_shape: "auto"`
+    charges a `nuclei_channel` config for a smaller tile than a
+    single-channel one, or when one config was segmented before the other's
+    settings changed, or a cheaper method (e.g. a DoG detector) sized its
+    own tile differently — `relate.py` detects a chunk mismatch per pair and
+    rechunks the finer side to match before relating, so this is no longer
+    something you need to plan around. (`run_multi.py` *also* auto-pins
+    every config under `"auto"` to one shared, smallest-computed tile
+    up front — mainly useful to keep segmentation itself GPU-efficient
+    across configs, not required for the relate step to work.)
 
 See [Relating labels across segmentations](label_relations.md) for what
 `label_relations()` returns and how to save it yourself — the cluster

{patchworks-2.6.4 → patchworks-2.6.6}/tests/test_run_multi.py

@@ -8,6 +8,8 @@ sys.path.insert(
     0, str(Path(__file__).resolve().parents[1] / "workflow" / "scripts")
 )
 
+import numpy as np  # noqa: E402
+import openpyxl  # noqa: E402
 import pytest  # noqa: E402
 import yaml  # noqa: E402
 
@@ -186,28 +188,146 @@ def test_relate_script_has_the_real_bookkeeping():
     assert "openpyxl" in src
 
 
-def test_auto_tile_shape_with_a_lone_nuclei_channel_is_refused():
-    """Matching `tile_shape` *values* are not enough when one config is 2-ch.
+def test_relate_rechunks_mismatched_label_arrays(tmp_path):
+    """A chunk-layout mismatch must be rechunked away, not require a re-run.
 
-    "auto" == "auto" passes the plain equality check, but the sizer charges
-    per channel, so the nuclei_channel config gets a smaller tile. The label
-    groups then disagree on chunk layout and label_relations raises -- after
-    every segmentation has already run, which is the expensive way to find out.
+    Two configs are free to have segmented at different tile_shape (one
+    published before the other's config changed, or a cheaper method sized
+    its own tile differently) -- label_relations() itself refuses mismatched
+    chunks by design, but that only means the caller has to rechunk one side
+    first, not that the whole segmentation needs redoing.
+    """
+    import zarr
+
+    from relate import run_relations
+
+    image_store = str(tmp_path / "image.zarr")
+
+    # a: labels 1 and 2, split at x=5. b: a single label 10 covering all of
+    # a's label 1 and none of label 2 -- built with a *different* chunking.
+    a_data = np.zeros((1, 10), dtype=np.int32)
+    a_data[0, :5] = 1
+    a_data[0, 5:] = 2
+    b_data = np.zeros((1, 10), dtype=np.int32)
+    b_data[0, :5] = 10
+
+    root = zarr.open_group(image_store, mode="w")
+    labels = root.require_group("labels")
+    a_grp = labels.require_group("nuclei_labels")
+    a_arr = a_grp.create_array(
+        name="0", shape=a_data.shape, chunks=(1, 2), dtype=np.int32
+    )
+    a_arr[:] = a_data
+    a_grp.attrs["sequential_labels"] = True
+    a_grp.attrs["n_objects"] = 2
+
+    b_grp = labels.require_group("cyto_labels")
+    b_arr = b_grp.create_array(
+        name="0", shape=b_data.shape, chunks=(1, 5), dtype=np.int32
+    )
+    b_arr[:] = b_data
+    b_grp.attrs["sequential_labels"] = True
+    b_grp.attrs["n_objects"] = 1
+
+    out_dir = tmp_path / "work"
+    out_dir.mkdir()
+    run_relations(
+        str(out_dir),
+        image_store,
+        [{"a": "nuclei_labels", "b": "cyto_labels", "output": "rel.xlsx"}],
+    )
+
+    wb = openpyxl.load_workbook(out_dir / "rel.xlsx")
+    rows = {
+        row[0]: (row[1], row[2], row[3])
+        for row in wb["nuclei_labels"].iter_rows(min_row=2, values_only=True)
+    }
+    assert rows[1] == (10, 5, 1.0)  # label 1 fully inside b's label 10
+    assert rows[2] == (None, 0, 0)  # label 2 touches nothing in b
+
+
+def test_mixed_nuclei_channel_auto_passes_validation():
+    """A channel-count mismatch under `tile_shape: "auto"` is no longer
+
+    refused at validation time -- it's resolved automatically instead (see
+    `_resolve_shared_tile_shape`), which needs the converted image's real
+    shape/dtype and so can only run after phase A, not from
+    `_validate_configs()`. This used to `sys.exit` here; asserting that
+    would now be testing the wrong layer.
     """
     paths = [Path("a.yaml"), Path("b.yaml")]
     base = {"work_dir": "/w", "tile_shape": "auto", "level": 0}
 
-    bad = [
+    mixed = [
         {**base, "label_name": "a", "channel": 0, "nuclei_channel": 1},
         {**base, "label_name": "b", "channel": 2},
     ]
-    with pytest.raises(SystemExit):
-        _validate_configs(paths, bad)
+    assert _validate_configs(paths, mixed) == "/w"
 
     # Same pair with one explicit shape is fine: both get that tile.
-    pinned = [{**c, "tile_shape": [16, 512, 512]} for c in bad]
+    pinned = [{**c, "tile_shape": [16, 512, 512]} for c in mixed]
     assert _validate_configs(paths, pinned) == "/w"
 
     # And "auto" is fine when every config carries the same channel count.
-    both = [{**bad[0]}, {**bad[1], "nuclei_channel": 3}]
+    both = [{**mixed[0]}, {**mixed[1], "nuclei_channel": 3}]
     assert _validate_configs(paths, both) == "/w"
+
+
+def test_resolve_shared_tile_shape_pins_the_smallest_candidate(
+    tmp_path, monkeypatch
+):
+    """The shared tile must be the tightest of every config's own budget.
+
+    A larger tile than some config's own sizer output would ask that config
+    for more memory than its settings were judged to need -- only the
+    smallest candidate is safe for every config at once.
+    """
+    import run_multi
+
+    class _FakeImage:
+        shape = (10, 100, 100)
+        dtype = "uint16"
+
+    calls = []
+
+    def _fake_load_ome_zarr(store, *, channel, level):
+        calls.append((store, channel, level))
+        return _FakeImage()
+
+    # One tile per config, matched up by call order (channel 0 then 1).
+    fake_tiles = [(8, 64, 64), (4, 32, 32)]
+
+    def _fake_sizer_cellpose(shape, dtype, **kwargs):
+        return fake_tiles[len(calls) - 1]
+
+    monkeypatch.setattr(
+        "patchworks.load_ome_zarr", _fake_load_ome_zarr, raising=False
+    )
+    monkeypatch.setattr(
+        "patchworks.auto_tile_shape_cellpose",
+        _fake_sizer_cellpose,
+        raising=False,
+    )
+
+    seg_cfgs = [
+        {
+            "channel": 0,
+            "level": 0,
+            "method": "cellpose",
+            "cellpose": {"do_3D": True, "gpu": True},
+        },
+        {
+            "channel": 1,
+            "level": 0,
+            "nuclei_channel": 2,
+            "method": "cellpose",
+            "cellpose": {"do_3D": True, "gpu": True},
+        },
+    ]
+    out = run_multi._resolve_shared_tile_shape(
+        seg_cfgs, "/w/image.zarr", str(tmp_path)
+    )
+    assert out == tmp_path / ".multi_tile_shape.generated.yaml"
+    written = yaml.safe_load(out.read_text())
+    assert written == {"tile_shape": [4, 32, 32]}  # the smaller candidate
+    assert len(calls) == 2

{patchworks-2.6.4 → patchworks-2.6.6}/workflow/scripts/relate.py

@@ -16,9 +16,15 @@ from __future__ import annotations
 
 import argparse
 import json
+import math
 from pathlib import Path
 
 
+def _num_chunks(arr: "da.Array") -> int:  # noqa: F821 - dask imported lazily by callers
+    """Total chunk count of a dask array, for picking the coarser side to rechunk to."""
+    return math.prod(len(c) for c in arr.chunks)
+
+
 def _label_ids(image_store: str, name: str) -> list[int]:
     """Ids present in a label image, without scanning the volume.
 
@@ -72,6 +78,35 @@ def run_relations(
         print(f"[relate] relating {a_name} -> {b_name} …", flush=True)
         a = da.from_zarr(image_store, component=f"labels/{a_name}/0")
         b = da.from_zarr(image_store, component=f"labels/{b_name}/0")
+
+        # label_relations() requires matching chunk layouts (it walks both
+        # arrays block-by-block at the same index) but two configs are free
+        # to have segmented at different tile_shape -- e.g. one already
+        # published before the other's config changed, or a cheaper method
+        # naturally sized its tile differently. Same shape, different
+        # chunking is a normal dask op (extra I/O reading across misaligned
+        # source chunks, not a correctness issue), so rechunk the finer side
+        # to the coarser one here rather than require identical tile_shape
+        # across every config up front.
+        if a.chunks != b.chunks:
+            a_n, b_n = _num_chunks(a), _num_chunks(b)
+            if a_n <= b_n:
+                print(
+                    f"[relate] {a_name} chunks {a.chunks} != {b_name} "
+                    f"chunks {b.chunks}; rechunking {b_name} to match "
+                    f"{a_name} (fewer chunks)",
+                    flush=True,
+                )
+                b = b.rechunk(a.chunks)
+            else:
+                print(
+                    f"[relate] {a_name} chunks {a.chunks} != {b_name} "
+                    f"chunks {b.chunks}; rechunking {a_name} to match "
+                    f"{b_name} (fewer chunks)",
+                    flush=True,
+                )
+                a = a.rechunk(b.chunks)
+
         table = label_relations(a, b)
 
         # label_relations() only returns a-objects that touch a b-object.

{patchworks-2.6.4 → patchworks-2.6.6}/workflow/scripts/run_multi.py

@@ -54,6 +54,7 @@ def _snakemake_cmd(
     extra: list[str] | None = None,
     jobname_prefix: str | None = None,
     common: Path | None = None,
+    extra_configfiles: list[Path] | None = None,
 ) -> list[str]:
     """Build one snakemake invocation.
 
@@ -66,10 +67,17 @@ def _snakemake_cmd(
     settings every config shares -- the input, the work_dir, everything
     ``convert`` reads -- live in one file and the per-config file carries only
     what actually differs.
+
+    *extra_configfiles*, when given, are appended after *configfile* and so
+    win over both it and *common* -- used to pin a driver-computed value
+    (e.g. a resolved ``tile_shape``) across every config without editing any
+    config file on disk.
     """
     configfiles = [str(configfile.resolve())]
     if common is not None:
         configfiles.insert(0, str(common.resolve()))
+    if extra_configfiles:
+        configfiles += [str(p.resolve()) for p in extra_configfiles]
     cmd = [
         "snakemake",
         "-s",
@@ -234,21 +242,12 @@ def _validate_configs(paths: list[Path], cfgs: list[dict]) -> str:
 
     # `tile_shape: "auto"` is identical as a *value* across configs while
     # producing different tiles: the sizer charges per channel, so a config
-    # with nuclei_channel gets a smaller one. The label groups then disagree on
-    # chunk layout and label_relations raises -- after every segmentation has
-    # run. Matching values are not enough here, so check the inputs that feed
-    # the sizer instead.
-    if {repr(cfg.get("tile_shape", "auto")) for cfg in cfgs} == {repr("auto")}:
-        if len({cfg.get("nuclei_channel") is not None for cfg in cfgs}) != 1:
-            problems.append(
-                'tile_shape: "auto" sizes a nuclei_channel config smaller '
-                "(a tile carries two channels), so the label groups would "
-                "not share a chunk layout and label_relations would fail "
-                f"after every segmentation had run; got "
-                f"{_spread('nuclei_channel')}. Set one explicit tile_shape in "
-                "the file `common:` points at, sized for the two-channel "
-                "config."
-            )
+    # with nuclei_channel gets a smaller one. The label groups would then
+    # disagree on chunk layout and label_relations would raise -- after every
+    # segmentation had run. That used to be a hard error asking for a manual
+    # explicit tile_shape; main() now resolves and pins one automatically
+    # (see _resolve_shared_tile_shape), once the converted image exists to
+    # size against, so there is nothing to check here anymore.
 
     # Phase A converts once, from the first config. Anything `convert` reads
     # out of a later config is therefore silently ignored -- someone setting
@@ -292,6 +291,92 @@ def _validate_configs(paths: list[Path], cfgs: list[dict]) -> str:
     return work_dirs.pop()
 
 
+def _resolve_shared_tile_shape(
+    seg_cfgs: list[dict], image_store: str, work_dir: str
+) -> Path:
+    """Auto-size ``tile_shape`` once, shared across every config.
+
+    ``tile_shape: "auto"`` resolves differently per config when
+    ``nuclei_channel`` differs -- the sizer charges per channel, so a
+    two-channel config gets a smaller tile. Left alone, the label groups
+    would end up with different chunk layouts and ``label_relations`` would
+    raise, after every segmentation had already run.
+
+    Computes what each config's own settings (channel count, ``do_3D``,
+    ``diameter``, GPU budget) would actually produce, then pins every config
+    to the *smallest* of them by voxel count -- the most memory-constrained
+    case, and safe for every config since a smaller tile only ever asks for
+    less memory than that config's own budget allows, never more. Configs
+    can therefore use different ``do_3D``/``diameter``/channel-count
+    settings and still end up with one shared, valid ``tile_shape``.
+
+    Only called once the converted image exists (the sizer needs its real
+    shape/dtype), so this runs from ``main()`` after phase A, not from
+    ``_validate_configs()``.
+
+    Returns
+    -------
+    Path
+        A generated one-key YAML file (``tile_shape: [...]``), meant to be
+        passed as an ``extra_configfiles`` entry to ``_snakemake_cmd`` so it
+        overrides every config's own (or common's) ``tile_shape`` value.
+    """
+    from functools import partial
+
+    import numpy as np
+    from patchworks import (
+        auto_tile_shape,
+        auto_tile_shape_cellpose,
+        load_ome_zarr,
+    )
+
+    candidates = []
+    for cfg in seg_cfgs:
+        image = load_ome_zarr(
+            image_store, channel=cfg["channel"], level=int(cfg.get("level", 0))
+        )
+        gpu_gb = cfg.get("gpu_memory_gb")
+        gpu_bytes = int(gpu_gb * 1024**3) if gpu_gb else None
+        n_channels = 2 if cfg.get("nuclei_channel") is not None else 1
+        method = cfg.get("method", "cellpose")
+        if method == "cellpose":
+            cp = cfg["cellpose"]
+            sizer = partial(
+                auto_tile_shape_cellpose,
+                do_3D=cp.get("do_3D", False),
+                use_gpu=cp.get("gpu", True),
+                diameter=cp.get("diameter"),
+                gpu_memory=gpu_bytes,
+                n_channels=n_channels,
+            )
+        else:
+            sizer = partial(
+                auto_tile_shape,
+                use_gpu=gpu_bytes is not None,
+                gpu_memory=gpu_bytes,
+                n_channels=n_channels,
+            )
+        candidates.append(tuple(int(x) for x in sizer(image.shape, image.dtype)))
+
+    tile_shape = min(candidates, key=lambda t: int(np.prod(t)))
+    print(
+        f'[run_multi] tile_shape: "auto" resolves differently across '
+        f"configs (nuclei_channel differs); pinning every config to the "
+        f"smallest computed tile {list(tile_shape)} so the label arrays "
+        f"share a chunk layout. Candidates were {[list(c) for c in candidates]}.",
+        flush=True,
+    )
+
+    override_path = Path(work_dir) / ".multi_tile_shape.generated.yaml"
+    override_path.write_text(
+        "# Generated by run_multi.py -- pins tile_shape across configs so\n"
+        "# label_relations sees matching chunk layouts. Safe to delete; it\n"
+        "# is regenerated on the next multi-config run.\n"
+        f"tile_shape: {list(tile_shape)}\n"
+    )
+    return override_path
+
+
 def _resolve(workflow_dir: Path, path_str: str) -> Path:
     path = Path(path_str)
     return path if path.is_absolute() else workflow_dir / path
@@ -453,6 +538,24 @@ def main() -> None:
         )
         sys.exit(rc)
 
+    # tile_shape: "auto" resolves differently per config when nuclei_channel
+    # differs (the sizer charges per channel) -- pin every config to one
+    # shared, computed value so the label arrays end up with matching chunk
+    # layouts. Needs the just-converted image, so this can only happen here,
+    # not in _validate_configs(). Dry runs never reach a real image.zarr.
+    tile_override = None
+    if not args.dry_run and {
+        cfg.get("tile_shape", "auto") for cfg in seg_cfgs
+    } == {"auto"}:
+        channel_counts = {
+            2 if cfg.get("nuclei_channel") is not None else 1
+            for cfg in seg_cfgs
+        }
+        if len(channel_counts) > 1:
+            tile_override = _resolve_shared_tile_shape(
+                seg_cfgs, image_store, work_dir
+            )
+
     # Phase B: the segmentations touch disjoint files under
     # work_dir/<label_name>/, so run them together and let the GPU partition
     # stay busy instead of idling through each config's prepare and merge.
@@ -470,6 +573,7 @@ def main() -> None:
             # Names the config in squeue, so concurrent runs are tellable apart.
             jobname_prefix=slurm_jobname_prefix(cfg["label_name"]),
             common=common_path,
+            extra_configfiles=[tile_override] if tile_override else None,
         )
         print(f"[run_multi] $ {' '.join(cmd)}", flush=True)
         procs.append((cfg_path.name, subprocess.Popen(cmd, cwd=workflow_dir)))