diff --git a/CLAUDE.md b/CLAUDE.md index 8b744a6..e7a9e24 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -43,9 +43,24 @@ The package lives in `bioimage_py/`: preserved on failure for `resume_from`) and `relabel_consecutive` (derives a consecutive mapping via a global `unique` reduction, then delegates the block-wise write to `relabel`; in place by default). `segmentation/multicut.py` ports the - bioimage-cpp-backed multicut cost transform + solvers (`compute_edge_costs`, - `multicut_decomposition` / `_gaec` / `_kernighan_lin`); meant to grow into multicut-based - segmentation. `segmentation/stitching.py` (`stitch_segmentation` / `stitch_tiled_segmentation`) + bioimage-cpp-backed multicut cost transform + whole-graph solvers (`compute_edge_costs`, + `multicut_decomposition` / `multicut_gaec` / `multicut_kernighan_lin`). + `segmentation/lifted_multicut.py` mirrors it for the lifted multicut + (`lifted_multicut_kernighan_lin` / `_gaec` / `_fusion_moves`, `get_lifted_multicut_solver`) plus + lifted-problem construction from a per-node label array (`lifted_edges_from_node_labels`, + `lifted_problem_from_node_labels`). `segmentation/blockwise_multicut.py` and + `segmentation/blockwise_lifted_multicut.py` add the hierarchical domain-decomposition solvers + (`blockwise_multicut` / `blockwise_lifted_multicut`): the per-block sub-problem solve is the **map** + step (distributed via the runner) and the union-find merge + edge contraction + final global solve + are the **in-process reduce**, with the block size doubling each level; both share + `segmentation/_blockwise_common.py` (the numpy reduce primitives, `reduce_problem`, and + `solve_subproblems_distributed`). They take `(graph, costs[, lifted_uv_ids, lifted_costs], + segmentation)` and return a node labeling — the segmentation defines the blocking domain (graph node + `i` == segment `i`); write the labeling to pixels with `relabel`. Unlike the elf reference, the + running node labeling is composed across levels (`node_ids = labels[unique(seg[block])]`) so + `n_levels > 1` is correct, and the graph/costs/labels are persisted to a temp zarr under the runner + temp root (dtypes canonicalized for zarr) for distributed backends. `segmentation/stitching.py` + (`stitch_segmentation` / `stitch_tiled_segmentation`) merges a tile-wise over-segmentation via a multicut over tile-interface object overlaps — the per-voxel phases (tile segmentation, overlap counting) run through the runner; RAG build + multicut solve are in-process private helpers (`_compute_rag` / `_project_node_labels_to_pixels`, @@ -148,9 +163,14 @@ the operations above (`stats`, `filters`, `segmentation.label` + `segmentation.w + `regionprops`, `copy`, and `downsample` — the latter built on the `ResizedSource` wrapper — the `evaluation` package: the parallel `contingency_table` primitive plus the metrics built on it — `segmentation.relabel` (apply a node labeling / relabeling map, in place by default; the canonical -node-label writer, see Conventions) + `segmentation.relabel_consecutive`, and -`segmentation.stitch_segmentation` / `stitch_tiled_segmentation` on the new `segmentation.multicut` -solvers). The slurm-only tests in `tests/test_slurm_runner.py` are skipped unless +node-label writer, see Conventions) + `segmentation.relabel_consecutive`, +`segmentation.stitch_segmentation` / `stitch_tiled_segmentation` on the `segmentation.multicut` +solvers, the whole-graph `segmentation.lifted_multicut` solvers + lifted-problem construction, and the +block-wise hierarchical `segmentation.blockwise_multicut` / `blockwise_lifted_multicut` — per-block +sub-problem solves distributed via the runner (`local` / `subprocess` / `slurm`), the reduce in-process +— verified by `tests/test_blockwise_multicut.py` (block-wise == whole-graph partition on a clean +problem, `local == subprocess` across workers and `n_levels`) and `tests/test_lifted_multicut.py`). The +slurm-only tests in `tests/test_slurm_runner.py` are skipped unless `sbatch` is on `PATH` and `BIOIMAGE_PY_SHARED_TMP` points at a shared filesystem; `subprocess` stays the CI proxy for the shared protocol. Note the slurm runner's key subtlety: per-task `.success` sentinels are written on compute nodes but can take up to the NFS attribute-cache timeout (~60 s) to become visible to diff --git a/bioimage_py/__init__.py b/bioimage_py/__init__.py index aa3ba7c..86f10fc 100644 --- a/bioimage_py/__init__.py +++ b/bioimage_py/__init__.py @@ -8,7 +8,7 @@ .. include:: ../docs/installation.md .. include:: ../docs/usage.md """ # noqa -from . import evaluation, filters, io, morphology, operations, segmentation, stats # noqa: F401 +from . import evaluation, filters, graph, io, morphology, operations, segmentation, stats # noqa: F401 from .copy import copy from .downsample import downsample from .runner import SlurmConfig, config_file_path, get_runner, write_slurm_config @@ -23,6 +23,7 @@ "segmentation", "morphology", "evaluation", + "graph", "operations", "io", "copy", diff --git a/bioimage_py/graph/__init__.py b/bioimage_py/graph/__init__.py new file mode 100644 index 0000000..0166d91 --- /dev/null +++ b/bioimage_py/graph/__init__.py @@ -0,0 +1,7 @@ +"""Graph operations: distributed region adjacency graphs and edge features.""" +from .distributed import distributed_edge_features, distributed_rag + +__all__ = [ + "distributed_rag", + "distributed_edge_features", +] diff --git a/bioimage_py/graph/distributed.py b/bioimage_py/graph/distributed.py new file mode 100644 index 0000000..701b55e --- /dev/null +++ b/bioimage_py/graph/distributed.py @@ -0,0 +1,336 @@ +"""Distributed region-adjacency graphs and edge features (block-wise, across backends). + +These operations build the region adjacency graph (RAG) of a labeled volume, and optionally its edge +features, block by block and merge the per-block results into one whole-volume graph. They wrap the +low-level primitives in :mod:`bioimage_cpp.graph.distributed` with the ``bioimage_py`` orchestration +(block iteration, halo sizing, I/O, and the in-process merge), so they scale identically across the +``local`` / ``subprocess`` / ``slurm`` backends -- mirroring the reduction ops in +:mod:`bioimage_py.evaluation` (per-block returns through ``runner.run(..., has_return_val=True)``, +merged in the orchestrating process). + +Two contracts the caller must honor: + +- **Globally consistent labels.** A segment must have the *same id in every block*. Run a + stitching / consistent-id step first (e.g. :func:`bioimage_py.segmentation.stitch_segmentation`, + or a distributed watershed with global ids) -- making per-block-local ids consistent is not part + of these operations. +- **Moment-only complex features.** Median / percentiles cannot be reconstructed from block + partials, so the distributed complex feature output is the moment subset ``[mean, std, min, max, + size]`` -- it equals the corresponding columns of the in-core complex features. + +The graph node ids are the label ids (node ``i`` is segment ``i``), so a multicut / clustering node +labeling is written back to pixels with the canonical node-label writer +:func:`bioimage_py.segmentation.relabel` using a dense array labeling. +""" +from __future__ import annotations + +from typing import List, Optional, Sequence, Tuple + +import bioimage_cpp as bic +import numpy as np + +from ..runner import get_runner +from ..runner.config import RunnerConfig +from ..sources import Source, SourceLike, as_source, capture_source, resolve_source +from ..util import BlockDescriptor, ComputeFn, check_direct, full_roi, to_roi + +__all__ = ["distributed_rag", "distributed_edge_features"] + +# Label dtypes the bioimage_cpp graph primitives accept directly; others are cast to uint64 per block. +_SUPPORTED_LABEL_DTYPES = ( + np.dtype("uint32"), np.dtype("uint64"), np.dtype("int32"), np.dtype("int64"), +) +_FEATURE_TYPES = ("edge_map", "affinity") + + +def _offset_halo(offsets: Sequence[Sequence[int]], ndim: int) -> List[int]: + """Per-axis halo for the affinity pass: ``max(1, max_offset)`` on every axis. + + The affinity pass extracts both the nearest-neighbor RAG (needs halo >= 1 on every axis) and the + affinity statistics (need halo >= ``max |offset component|`` per axis), so the halo is the + elementwise maximum of the two -- the ``1`` floor is what keeps NN edges on axes that only carry + long-range offsets. + """ + halo = [1] * ndim + for offset in offsets: + for axis in range(ndim): + halo[axis] = max(halo[axis], abs(int(offset[axis]))) + return halo + + +def _block_extract(labels: np.ndarray, data: Optional[np.ndarray], feature_type: str, + offsets: Optional[Sequence[Sequence[int]]], own_begin: Sequence[int], + own_shape: Sequence[int]) -> Tuple[np.ndarray, Optional[np.ndarray], + Optional[np.ndarray], int]: + """Extract one (haloed) block's owned edges, partial stats and label max. + + Shared by the direct fast path and the per-block compute function, so ``direct == blocked`` holds + by construction. Returns ``(graph_edges, stats_edges, stats, local_max)``: ``graph_edges`` are + the owned nearest-neighbor edges (feed the global graph); ``stats_edges`` / ``stats`` are the + edges / partial statistics for the feature fold (``None`` for the graph-only pass); ``local_max`` + is the maximum label id in the block (folds the node-count reduction into this pass). + """ + if labels.dtype not in _SUPPORTED_LABEL_DTYPES: + labels = labels.astype("uint64") + local_max = int(labels.max()) if labels.size else 0 + dist = bic.graph.distributed + if feature_type == "edge_map": + # The stats primitive's edges equal block_region_adjacency_edges for the block, so one call + # yields both the graph edges and the feature stats. + edges, stats = dist.block_edge_map_stats(labels, data, own_begin, own_shape) + return edges, edges, stats, local_max + if feature_type == "affinity": + # The affinity stats edges may include long-range-only pairs, so the graph must come from the + # nearest-neighbor extraction; long-range pairs are dropped by merge_block_edge_stats. + graph_edges = dist.block_region_adjacency_edges(labels, own_begin, own_shape) + stats_edges, stats = dist.block_affinity_stats(labels, data, offsets, own_begin, own_shape) + return graph_edges, stats_edges, stats, local_max + graph_edges = dist.block_region_adjacency_edges(labels, own_begin, own_shape) + return graph_edges, None, None, local_max + + +def _make_compute(feature_type: str, offsets: Optional[Sequence[Sequence[int]]], + data_handle: Optional[object]) -> ComputeFn: + """Build the per-block compute function (a cloudpicklable closure over the captured data handle). + + ``labels`` is the runner input (it defines the domain / blocking); ``data`` (the edge map or the + affinities) is captured and reopened per block, never passed as a runner input -- the affinities' + ``(channels, *spatial)`` shape would fail the runner's input shape check, and capturing is + uniform for the edge map too. + """ + + def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source], + mask: Optional[Source]) -> Tuple[np.ndarray, Optional[np.ndarray], + Optional[np.ndarray], int]: + labels_src = inputs[0] + outer = to_roi(block.outer_block) # halo included; symmetric + clamped at borders. + labels = np.ascontiguousarray(labels_src[outer]) + own_begin = [int(b) for b in block.inner_block_local.begin] + own_shape = [int(e) - int(b) + for b, e in zip(block.inner_block_local.begin, block.inner_block_local.end)] + data = None + if data_handle is not None: + data_src = resolve_source(data_handle) + if feature_type == "affinity": + data = np.ascontiguousarray(data_src[(slice(None),) + outer]) + else: + data = np.ascontiguousarray(data_src[outer]) + return _block_extract(labels, data, feature_type, offsets, own_begin, own_shape) + + return _compute + + +def _resolve_n_nodes(results: Sequence[tuple], n_nodes: Optional[int]) -> int: + """Return the node count: ``max(label) + 1`` (from the per-block maxes) unless given explicitly. + + Must be ``max(label) + 1`` (never ``edges.max() + 1``): isolated segments have no edges but still + need a node id for the labeling / writeback and for the stats fold's edge-endpoint validity. A + caller-supplied ``n_nodes`` must be at least this large. + """ + local_max = max((int(r[3]) for r in results), default=0) + required = local_max + 1 + if n_nodes is None: + return required + n_nodes = int(n_nodes) + if n_nodes < required: + raise ValueError( + f"n_nodes={n_nodes} is too small: the labels contain id {local_max}, so n_nodes must be " + f">= {required}." + ) + return n_nodes + + +def _build_graph(results: Sequence[tuple], n_nodes: int) -> "bic.graph.UndirectedGraph": + """Merge the per-block edges into the global :class:`bioimage_cpp.graph.UndirectedGraph`.""" + merged = bic.graph.distributed.merge_edges([r[0] for r in results]) + if len(merged) == 0: + # from_unique_edges must not be called with an empty edge array; build the node-only graph. + return bic.graph.UndirectedGraph(int(n_nodes)) + return bic.graph.UndirectedGraph.from_unique_edges(int(n_nodes), merged) + + +def _reduce_features(results: Sequence[tuple], graph: "bic.graph.UndirectedGraph", + complex_features: bool) -> np.ndarray: + """Fold the per-block partial stats onto the global edges and finalize the features.""" + dist = bic.graph.distributed + n_edges = int(graph.number_of_edges) + acc = dist.empty_edge_stats(n_edges) + if n_edges: + for _, stats_edges, stats, _ in results: + acc = dist.merge_block_edge_stats(graph, acc, stats_edges, stats) + return dist.finalize_edge_features(acc, compute_complex_features=complex_features) + + +def _validate_labels(labels_src: Source) -> None: + """Validate the label source is an integer 2D/3D image (the primitives' domain).""" + if not np.issubdtype(np.dtype(labels_src.dtype), np.integer): + raise ValueError(f"labels must be an integer label image, got dtype {labels_src.dtype}.") + if labels_src.ndim not in (2, 3): + raise ValueError(f"labels must be a 2D or 3D image, got ndim={labels_src.ndim}.") + + +def _validate_data(data_src: Source, labels_src: Source, feature_type: str, + offsets: Optional[Sequence[Sequence[int]]]) -> None: + """Validate the feature data shape against the labels (edge map same-shape / affinity channels).""" + label_shape = tuple(int(s) for s in labels_src.shape) + if feature_type == "edge_map": + if tuple(int(s) for s in data_src.shape) != label_shape: + raise ValueError( + f"edge_map data shape {tuple(data_src.shape)} must match the labels shape " + f"{label_shape}." + ) + else: # affinity + if data_src.ndim != labels_src.ndim + 1: + raise ValueError( + "affinity data must have shape (len(offsets), *labels.shape), got ndim=" + f"{data_src.ndim} for {labels_src.ndim}D labels." + ) + if tuple(int(s) for s in data_src.shape[1:]) != label_shape: + raise ValueError( + f"affinity data spatial shape {tuple(data_src.shape[1:])} must match the labels " + f"shape {label_shape}." + ) + if int(data_src.shape[0]) != len(offsets): + raise ValueError( + f"offsets length {len(offsets)} must match the affinity channel count " + f"{int(data_src.shape[0])}." + ) + + +def _run(labels: SourceLike, data: Optional[SourceLike], feature_type: str, + offsets: Optional[Sequence[Sequence[int]]], complex_features: bool, + block_shape: Optional[Tuple[int, ...]], n_nodes: Optional[int], num_workers: int, + job_type: str, job_config: Optional[RunnerConfig]): + """Shared driver for the graph-only ("rag") and feature passes (see the public functions).""" + labels_src = as_source(labels) + _validate_labels(labels_src) + ndim = labels_src.ndim + data_src = as_source(data) if data is not None else None + if data_src is not None: + _validate_data(data_src, labels_src, feature_type, offsets) + + halo = _offset_halo(offsets, ndim) if feature_type == "affinity" else [1] * ndim + + if check_direct(job_type, num_workers, block_shape, None, None): + # Direct fast path: one whole-array extraction feeding the same merge/fold helpers. + labels_arr = np.ascontiguousarray(labels_src[full_roi(ndim)]) + data_arr = None + if data_src is not None: + data_arr = np.ascontiguousarray( + data_src[(slice(None),) + full_roi(ndim)] if feature_type == "affinity" + else data_src[full_roi(ndim)]) + shape = [int(s) for s in labels_src.shape] + results = [_block_extract(labels_arr, data_arr, feature_type, offsets, [0] * ndim, shape)] + else: + runner = get_runner(job_type, job_config) + data_handle = capture_source(data_src, job_type) if data_src is not None else None + compute = _make_compute(feature_type, offsets, data_handle) + results = runner.run(compute, [labels], block_shape=block_shape, halo=halo, + num_workers=num_workers, has_return_val=True, + name=f"distributed-{feature_type}") + + n_nodes_ = _resolve_n_nodes(results, n_nodes) + graph = _build_graph(results, n_nodes_) + if feature_type == "rag": + return graph + features = _reduce_features(results, graph, complex_features) + return graph, features + + +def distributed_rag( + labels: SourceLike, + *, + block_shape: Optional[Tuple[int, ...]] = None, + n_nodes: Optional[int] = None, + num_workers: int = 1, + job_type: str = "local", + job_config: Optional[RunnerConfig] = None, +) -> "bic.graph.UndirectedGraph": + """Build the region adjacency graph of a labeled volume, block-wise. + + Each block is read with a one-pixel halo, its owned nearest-neighbor edges are extracted, and the + per-block edge sets are merged into one whole-volume graph. The result is exact regardless of how + objects straddle block boundaries: it equals the whole-volume + ``bioimage_cpp.graph.region_adjacency_graph(labels)``. Graph node ``i`` is segment ``i``. + + Args: + labels: The input segmentation (a numpy/zarr/n5 array or a `Source`); must be an integer 2D + or 3D label image with **globally consistent** ids (the same id in every block). + block_shape: Shape of the processing blocks. Defaults to the input chunk shape; required for + unchunked data. + n_nodes: The number of graph nodes. Defaults to ``max(label) + 1`` (derived in the block + pass). If given, must be at least ``max(label) + 1``. + num_workers: Number of parallel workers (threads for ``local``, tasks for distributed + backends). + job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``. + job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`). + + Returns: + The whole-volume region adjacency graph as a `bioimage_cpp.graph.UndirectedGraph`. + """ + return _run(labels, None, "rag", None, False, block_shape, n_nodes, num_workers, job_type, + job_config) + + +def distributed_edge_features( + labels: SourceLike, + data: SourceLike, + *, + feature_type: str = "edge_map", + offsets: Optional[Sequence[Sequence[int]]] = None, + complex_features: bool = False, + block_shape: Optional[Tuple[int, ...]] = None, + n_nodes: Optional[int] = None, + num_workers: int = 1, + job_type: str = "local", + job_config: Optional[RunnerConfig] = None, +) -> Tuple["bic.graph.UndirectedGraph", np.ndarray]: + """Build the region adjacency graph and its edge features in one block pass. + + Each block is read with a halo, its owned edges and partial edge statistics are extracted, and + both are merged into the whole-volume graph and finalized edge features. The features match the + moment columns of the in-core ``bioimage_cpp.graph.features`` on the whole volume: ``size``, + ``min`` and ``max`` are exact; ``mean`` and ``std`` match to floating-point tolerance (their + block-merge order is not fixed). + + Args: + labels: The input segmentation (a numpy/zarr/n5 array or a `Source`); must be an integer 2D + or 3D label image with **globally consistent** ids (the same id in every block). + data: The data the features are computed from. For ``feature_type="edge_map"`` an array the + same shape as ``labels`` (the edge value is the average of the two endpoint pixels). For + ``feature_type="affinity"`` an array of shape ``(len(offsets), *labels.shape)``. + feature_type: The feature source; one of ``"edge_map"`` or ``"affinity"``. + offsets: The affinity offsets (one per channel); required for ``feature_type="affinity"``, + ignored otherwise. The halo is sized from them (``max(1, max |offset|)`` per axis). + complex_features: If ``False`` the features are ``(E, 2)`` columns ``[mean, size]``; if + ``True`` they are ``(E, 5)`` columns ``[mean, std, min, max, size]``. + block_shape: Shape of the processing blocks. Defaults to the input chunk shape; required for + unchunked data. + n_nodes: The number of graph nodes. Defaults to ``max(label) + 1`` (derived in the block + pass). If given, must be at least ``max(label) + 1``. + num_workers: Number of parallel workers (threads for ``local``, tasks for distributed + backends). + job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``. + job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`). + + Returns: + A tuple ``(graph, features)`` of the whole-volume `bioimage_cpp.graph.UndirectedGraph` and + the ``(number_of_edges, k)`` feature array (rows aligned to the graph edge ids). + + Raises: + ValueError: If ``feature_type`` is unknown, or ``offsets`` is missing for the affinity + feature type, or the data shape is incompatible with the labels. + """ + if feature_type not in _FEATURE_TYPES: + types = ", ".join(_FEATURE_TYPES) + raise ValueError(f"feature_type must be one of ({types}), got {feature_type!r}.") + if feature_type == "affinity": + if offsets is None: + raise ValueError("offsets is required for feature_type='affinity'.") + offsets = [[int(v) for v in offset] for offset in offsets] + if not offsets: + raise ValueError("offsets must not be empty.") + if any(len(offset) != as_source(labels).ndim for offset in offsets): + raise ValueError("each offset must have length matching the labels ndim.") + return _run(labels, data, feature_type, offsets, complex_features, block_shape, n_nodes, + num_workers, job_type, job_config) diff --git a/bioimage_py/segmentation/__init__.py b/bioimage_py/segmentation/__init__.py index b3d0f6f..581801b 100644 --- a/bioimage_py/segmentation/__init__.py +++ b/bioimage_py/segmentation/__init__.py @@ -1,5 +1,10 @@ """Segmentation: connected-component labeling and related operations.""" +from .blockwise_lifted_multicut import blockwise_lifted_multicut +from .blockwise_multicut import blockwise_multicut from .label import label +from .lifted_multicut import (get_lifted_multicut_solver, lifted_edges_from_node_labels, + lifted_multicut_fusion_moves, lifted_multicut_gaec, + lifted_multicut_kernighan_lin, lifted_problem_from_node_labels) from .multicut import (compute_edge_costs, multicut_decomposition, multicut_gaec, multicut_kernighan_lin, transform_probabilities_to_costs) from .relabel import relabel, relabel_consecutive @@ -21,4 +26,12 @@ "multicut_decomposition", "multicut_gaec", "multicut_kernighan_lin", + "blockwise_multicut", + "blockwise_lifted_multicut", + "lifted_multicut_kernighan_lin", + "lifted_multicut_gaec", + "lifted_multicut_fusion_moves", + "get_lifted_multicut_solver", + "lifted_edges_from_node_labels", + "lifted_problem_from_node_labels", ] diff --git a/bioimage_py/segmentation/_blockwise_common.py b/bioimage_py/segmentation/_blockwise_common.py new file mode 100644 index 0000000..6f2f8d8 --- /dev/null +++ b/bioimage_py/segmentation/_blockwise_common.py @@ -0,0 +1,395 @@ +"""Shared machinery for the block-wise (lifted) multicut solvers. + +The block-wise multicut (Pape et al. 2017) solves a global multicut by hierarchical domain +decomposition: at each level the graph is partitioned into blocks, each block's induced sub-problem is +solved independently (the *map* step, distributed across the runner backends), and the per-block cut +decisions are merged into a coarser graph (the *reduce* step, in the orchestrating process). Repeating +this over successively larger blocks contracts the problem until a final global solve is cheap. + +This module holds the pieces shared by :mod:`bioimage_py.segmentation.blockwise_multicut` and +:mod:`bioimage_py.segmentation.blockwise_lifted_multicut`: + +- the numpy reduce primitives ported from ``elf`` (union-find merge + edge contraction with summed + costs; ``bioimage_cpp.utils.UnionFind`` replaces nifty's ufd, and :func:`_remap_edges_sum_costs` + replaces nifty's ``EdgeMapping``), +- :func:`reduce_problem`, the in-process reduction of one level, +- :func:`solve_subproblems_distributed`, the distributed map step (per-block sub-problem solve via + ``runner.run``), and +- the internal-solver dispatch-by-name (so the cloudpickled per-block closure captures a string, not a + solver object). + +The block's node membership is derived from the segmentation volume, mapped through the running +composed node labeling ``labels`` (original node id -> current supernode id). This keeps the *original* +segmentation as the only inter-level volume and, unlike the ``elf`` reference, makes ``n_levels > 1`` +correct (``elf`` indexes the contracted graph with original ids). +""" +from __future__ import annotations + +import os +import tempfile +from typing import Dict, List, Optional, Sequence, Tuple + +import bioimage_cpp as bic +import numpy as np + +from ..sources import Source, SourceSpec, as_source, from_spec +from ..util import BlockDescriptor, ComputeFn, full_roi, to_roi + + +def _undirected_graph_types() -> tuple: + """The concrete ``UndirectedGraph`` classes (the public wrapper and the raw ``_core`` class). + + ``bioimage_cpp`` exposes a Python wrapper ``bic.graph.UndirectedGraph`` whose instances (from + ``from_edges``) are distinct from the raw ``bic._core.UndirectedGraph`` returned by some builders + (e.g. ``from_unique_edges``, used by ``distributed_rag``), so a single ``isinstance`` check is not + enough. Both are accepted as an already-built graph. + """ + types = [bic.graph.UndirectedGraph] + core = getattr(bic, "_core", None) + core_type = getattr(core, "UndirectedGraph", None) + if core_type is not None and core_type not in types: + types.append(core_type) + return tuple(types) + + +_UNDIRECTED_GRAPH_TYPES = _undirected_graph_types() + + +def _relabel_from_zero(node_ids: np.ndarray) -> Tuple[np.ndarray, int, Dict[int, int]]: + """Map ``node_ids`` to consecutive integers from 0. Returns ``(relabeled, max_id, {old: new})``.""" + uniq, inverse = np.unique(node_ids, return_inverse=True) + relabeled = np.reshape(inverse, -1).astype(node_ids.dtype, copy=False) + mapping = {int(old): int(new) for new, old in enumerate(uniq)} + return relabeled, int(uniq.size - 1), mapping + + +def _relabel_keep_zero(labels: np.ndarray) -> np.ndarray: + """Map labels to consecutive integers from 1 while keeping 0 fixed at 0.""" + out = np.zeros_like(labels) + nonzero_mask = labels != 0 + if nonzero_mask.any(): + _, inverse = np.unique(labels[nonzero_mask], return_inverse=True) + out[nonzero_mask] = (np.reshape(inverse, -1) + 1).astype(labels.dtype, copy=False) + return out + + +def _remap_edges_sum_costs(uv_ids: np.ndarray, new_labels: np.ndarray, + costs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Contract ``uv_ids`` onto ``new_labels`` and sum the costs collapsing onto each supernode pair. + + Builds one row per unique pair of distinct supernodes (self-loops are dropped) and sums the costs + of all original edges mapping onto that pair. Replaces nifty's ``EdgeMapping``. + """ + uv_ids = np.asarray(uv_ids) + costs = np.asarray(costs) + if uv_ids.size == 0: + return uv_ids.reshape(0, 2).astype("uint64"), np.zeros(0, dtype=costs.dtype) + + remapped = new_labels[uv_ids] + keep = remapped[:, 0] != remapped[:, 1] + remapped = remapped[keep] + sub_costs = costs[keep] + if remapped.shape[0] == 0: + return np.zeros((0, 2), dtype="uint64"), np.zeros(0, dtype=sub_costs.dtype) + + # canonicalize each row so the same undirected edge always maps to the same key + u = np.minimum(remapped[:, 0], remapped[:, 1]).astype(np.int64) + v = np.maximum(remapped[:, 0], remapped[:, 1]).astype(np.int64) + n_new = int(new_labels.max()) + 1 + key = u * n_new + v + + unique_key, inverse = np.unique(key, return_inverse=True) + inverse = np.reshape(inverse, -1) + new_costs = np.zeros(unique_key.size, dtype=sub_costs.dtype) + np.add.at(new_costs, inverse, sub_costs) + + new_u = (unique_key // n_new).astype("uint64") + new_v = (unique_key % n_new).astype("uint64") + new_uv_ids = np.stack([new_u, new_v], axis=1) + return new_uv_ids, new_costs + + +def find_inner_lifted_edges(lifted_uv_ids: np.ndarray, node_list: np.ndarray) -> np.ndarray: + """Return the indices of the lifted edges with **both** endpoints in ``node_list``.""" + lifted_uv_ids = np.asarray(lifted_uv_ids) + if lifted_uv_ids.size == 0: + return np.zeros(0, dtype="int64") + lifted_indices = np.arange(len(lifted_uv_ids)) + inner_us = np.isin(lifted_uv_ids[:, 0], node_list) + inner_indices = lifted_indices[inner_us] + inner_uvs = lifted_uv_ids[inner_us] + inner_vs = np.isin(inner_uvs[:, 1], node_list) + return inner_indices[inner_vs] + + +def build_graph(n_nodes: int, uv_ids: np.ndarray) -> "bic.graph.UndirectedGraph": + """Build a ``bic.graph.UndirectedGraph`` from ``uv_ids`` (node-only when there are no edges).""" + uv_ids = np.ascontiguousarray(np.asarray(uv_ids), dtype="uint64") + if uv_ids.size == 0: + return bic.graph.UndirectedGraph(int(n_nodes)) + return bic.graph.UndirectedGraph.from_edges(int(n_nodes), uv_ids) + + +def as_undirected_graph(graph) -> "bic.graph.UndirectedGraph": + """Return ``graph`` as an ``UndirectedGraph`` (rebuild from edges if it is a RAG).""" + if isinstance(graph, _UNDIRECTED_GRAPH_TYPES): + return graph + return bic.graph.UndirectedGraph.from_edges( + graph.number_of_nodes, np.asarray(graph.uv_ids(), dtype="uint64"), + ) + + +def multicut_solver_by_name(name: str): + """Return a whole-graph multicut solver ``solver(graph, costs) -> node_labels`` for ``name``.""" + from .multicut import multicut_decomposition, multicut_gaec, multicut_kernighan_lin + + if name == "kernighan-lin": + return multicut_kernighan_lin + if name == "greedy-additive": + return multicut_gaec + if name == "decomposition": + return multicut_decomposition + raise ValueError( + f"Unknown internal multicut solver {name!r}; expected one of 'kernighan-lin', " + "'greedy-additive', 'decomposition'." + ) + + +def lifted_solver_by_name(name: str): + """Return a lifted solver ``solver(graph, costs, lifted_uv, lifted_costs) -> node_labels``.""" + from .lifted_multicut import lifted_multicut_gaec, lifted_multicut_kernighan_lin + + if name == "kernighan-lin": + return lifted_multicut_kernighan_lin + if name == "greedy-additive": + return lifted_multicut_gaec + raise ValueError( + f"Unknown internal lifted multicut solver {name!r}; expected one of 'kernighan-lin', " + "'greedy-additive'." + ) + + +def reduce_problem(graph, costs: np.ndarray, merge_mask: np.ndarray, + lifted_uv_ids: Optional[np.ndarray] = None, + lifted_costs: Optional[np.ndarray] = None): + """Reduce one level: union-find merge over the kept edges, then contract the graph. + + Args: + graph: The current level's graph. + costs: The current level's base edge costs (aligned to ``graph`` edge ids). + merge_mask: Boolean mask over the base edges; ``True`` edges are merged (kept), ``False`` cut. + lifted_uv_ids: Optional current-level lifted edges (contracted through the same labeling). + lifted_costs: Optional current-level lifted costs. + + Returns: + ``(new_graph, new_costs, new_labels)`` for the base problem, or + ``(new_graph, new_costs, new_labels, new_lifted_uv_ids, new_lifted_costs)`` when lifted arrays + are given. ``new_labels`` maps every current node id to its (consecutive) supernode id. + """ + n_nodes = int(graph.number_of_nodes) + nodes = np.arange(n_nodes, dtype="uint64") + uv_ids = np.asarray(graph.uv_ids(), dtype="uint64") + + ufd = bic.utils.UnionFind(n_nodes) + merge_uv = np.ascontiguousarray(uv_ids[merge_mask], dtype="uint64") + if merge_uv.size: + ufd.merge(merge_uv) + new_labels = _relabel_keep_zero(np.asarray(ufd.find(nodes))) + + new_uv_ids, new_costs = _remap_edges_sum_costs(uv_ids, new_labels, np.asarray(costs)) + n_new_nodes = int(new_labels.max()) + 1 if new_labels.size else 0 + new_graph = build_graph(n_new_nodes, new_uv_ids) + + if lifted_uv_ids is None: + return new_graph, new_costs, new_labels + + lifted_uv_ids = np.asarray(lifted_uv_ids, dtype="uint64") + if lifted_uv_ids.size: + new_lifted_uvs, new_lifted_costs = _remap_edges_sum_costs( + lifted_uv_ids, new_labels, np.asarray(lifted_costs), + ) + else: + new_lifted_uvs = lifted_uv_ids.reshape(0, 2) + new_lifted_costs = np.zeros(0, dtype="float64") + return new_graph, new_costs, new_labels, new_lifted_uvs, new_lifted_costs + + +# --- distributed map step --------------------------------------------------------------------------- + + +_CANONICAL_DTYPES = { + ("u", 1): np.uint8, ("u", 2): np.uint16, ("u", 4): np.uint32, ("u", 8): np.uint64, + ("i", 1): np.int8, ("i", 2): np.int16, ("i", 4): np.int32, ("i", 8): np.int64, + ("f", 2): np.float16, ("f", 4): np.float32, ("f", 8): np.float64, ("b", 1): np.bool_, +} + + +def _canonicalize(arr: np.ndarray) -> np.ndarray: + """Cast to the canonical numpy dtype for its (kind, itemsize). + + ``bioimage_cpp`` returns e.g. a ``ULongLong``-backed ``uint64`` (C ``unsigned long long``), which + is ``==``-equal to but a distinct dtype class from numpy's canonical ``UInt64`` -- and zarr v3's + dtype registry matches only the canonical class. Re-casting normalizes it so the array persists. + """ + arr = np.asarray(arr) + canon = _CANONICAL_DTYPES.get((arr.dtype.kind, arr.dtype.itemsize)) + if canon is None: + return np.ascontiguousarray(arr) + return np.ascontiguousarray(arr.astype(canon)) + + +def _persist_array(arr: np.ndarray, tmp_dir: str, name: str): + """Persist an array to a temp zarr and return its reopen spec (for distributed workers).""" + import zarr + + arr = _canonicalize(arr) + path = os.path.join(tmp_dir, f"{name}.zarr") + chunks = tuple(max(1, int(s)) for s in arr.shape) if arr.ndim else None + z = zarr.open_array(path, mode="w", shape=arr.shape, dtype=arr.dtype, chunks=chunks) + if arr.size: + z[:] = arr + return as_source(z).to_spec() + + +def _capture_graph(graph, job_type: str, tmp_dir: Optional[str]): + """Capture the graph for the per-block closure: the live object for local, its edges' spec else.""" + if job_type == "local": + return graph + return _persist_array(np.asarray(graph.uv_ids(), dtype="uint64"), tmp_dir, "uv") + + +def _capture_array(arr: np.ndarray, job_type: str, tmp_dir: Optional[str], name: str): + """Capture an array for the per-block closure: the array for local, a reopen spec otherwise.""" + if job_type == "local": + return np.ascontiguousarray(arr) + return _persist_array(arr, tmp_dir, name) + + +def _read_array(handle) -> np.ndarray: + """Materialize a captured array handle (a numpy array for local, a spec to reopen otherwise).""" + if isinstance(handle, np.ndarray): + return handle + src = from_spec(handle) + return np.asarray(src[full_roi(src.ndim)]) + + +def _make_solve_block(graph_handle, n_nodes: int, costs_handle, labels_handle, + solver_name: str, has_halo: bool, lifted_handles) -> ComputeFn: + """Build the per-block sub-problem solver (a cloudpicklable closure over the captured handles). + + ``inputs[0]`` is the segmentation (it defines the blocking domain). The graph, costs, composed + labels and (optional) lifted arrays are captured as handles and materialized + cached once per + worker process. The block returns its cut edge ids (inner edges cut by the local solve, plus all + block-boundary edges) through the return-value channel. + """ + cache: Dict[str, object] = {} + + def _graph() -> "bic.graph.UndirectedGraph": + g = cache.get("graph") + if g is None: + # Distributed: graph_handle is a SourceSpec of the persisted edges -> rebuild. + # Local: graph_handle is the live graph object -> use directly. + g = build_graph(int(n_nodes), _read_array(graph_handle)) \ + if isinstance(graph_handle, SourceSpec) else graph_handle + cache["graph"] = g + cache["uv_ids"] = np.asarray(g.uv_ids(), dtype="uint64") + return g # type: ignore[return-value] + + def _cached(key: str, handle) -> np.ndarray: + val = cache.get(key) + if val is None: + val = _read_array(handle) + cache[key] = val + return val # type: ignore[return-value] + + def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source], + mask: Optional[Source]) -> np.ndarray: + seg_src = inputs[0] + roi = to_roi(block.outer_block) if has_halo else to_roi(block) + node_ids = np.unique(seg_src[roi]).astype("uint64") + if labels_handle is not None: + labels = _cached("labels", labels_handle) + node_ids = np.unique(labels[node_ids]).astype("uint64") + + graph = _graph() + uv_ids = cache["uv_ids"] # type: ignore[assignment] + inner_edges, outer_edges = graph.extract_subgraph_from_nodes(node_ids) + inner_edges = np.asarray(inner_edges) + outer_edges = np.asarray(outer_edges).astype("uint64") + if inner_edges.size == 0: + return outer_edges + + sub_uvs = np.asarray(uv_ids)[inner_edges] + _, max_id, mapping = _relabel_from_zero(node_ids) + sub_uvs = bic.utils.take_dict(mapping, np.ascontiguousarray(sub_uvs, dtype="uint64")) + sub_graph = build_graph(max_id + 1, sub_uvs) + sub_costs = _cached("costs", costs_handle)[inner_edges] + + if lifted_handles is None: + sub_result = multicut_solver_by_name(solver_name)(sub_graph, sub_costs) + else: + luv_handle, lcosts_handle = lifted_handles + lifted_uv = _cached("luv", luv_handle) + inner_lifted = find_inner_lifted_edges(lifted_uv, node_ids) + if inner_lifted.size: + sub_lifted_uvs = bic.utils.take_dict( + mapping, np.ascontiguousarray(lifted_uv[inner_lifted], dtype="uint64"), + ) + sub_lifted_costs = _cached("lcosts", lcosts_handle)[inner_lifted] + sub_result = lifted_solver_by_name(solver_name)( + sub_graph, sub_costs, sub_lifted_uvs, sub_lifted_costs, + ) + else: + sub_result = multicut_solver_by_name(solver_name)(sub_graph, sub_costs) + + sub_result = np.asarray(sub_result) + cut = sub_result[sub_uvs[:, 0]] != sub_result[sub_uvs[:, 1]] + return np.concatenate([inner_edges[cut].astype("uint64"), outer_edges]) + + return _compute + + +def solve_subproblems_distributed(runner, graph, costs: np.ndarray, segmentation_src: Source, + labels: Optional[np.ndarray], *, block_shape: Tuple[int, ...], + halo: Optional[Sequence[int]], internal_solver: str, + num_workers: int, job_type: str, tmp_dir: Optional[str], + name: str, lifted_uv_ids: Optional[np.ndarray] = None, + lifted_costs: Optional[np.ndarray] = None) -> np.ndarray: + """Run the per-block sub-problem solves (map) and reduce them to a base-edge merge mask. + + Returns a boolean mask over ``graph`` edges: ``True`` where the edge is kept (merged), ``False`` + where it is cut by at least one block (block-boundary edges are always cut at the current level). + """ + n_edges = int(graph.number_of_edges) + has_halo = halo is not None + level_dir: Optional[str] = None + if job_type != "local": + level_dir = tempfile.mkdtemp(dir=tmp_dir) + + graph_handle = _capture_graph(graph, job_type, level_dir) + costs_handle = _capture_array(np.asarray(costs), job_type, level_dir, "costs") + labels_handle = None if labels is None \ + else _capture_array(np.asarray(labels), job_type, level_dir, "labels") + lifted_handles = None + if lifted_uv_ids is not None: + luv = np.asarray(lifted_uv_ids, dtype="uint64").reshape(-1, 2) + lifted_handles = ( + _capture_array(luv, job_type, level_dir, "luv"), + _capture_array(np.asarray(lifted_costs, dtype="float64"), job_type, level_dir, "lcosts"), + ) + + compute = _make_solve_block(graph_handle, int(graph.number_of_nodes), costs_handle, + labels_handle, internal_solver, has_halo, lifted_handles) + results: List[Optional[np.ndarray]] = runner.run( + compute, [segmentation_src], outputs=[], block_shape=block_shape, halo=halo, + num_workers=num_workers, has_return_val=True, name=name, + ) + + cut = np.zeros(n_edges, dtype=bool) + for res in results: + if res is None: + continue + res = np.asarray(res) + if res.size: + cut[res.astype("int64")] = True + return ~cut diff --git a/bioimage_py/segmentation/blockwise_lifted_multicut.py b/bioimage_py/segmentation/blockwise_lifted_multicut.py new file mode 100644 index 0000000..2d5ed0d --- /dev/null +++ b/bioimage_py/segmentation/blockwise_lifted_multicut.py @@ -0,0 +1,121 @@ +"""Block-wise (hierarchical) lifted multicut with a distributed per-block solve. + +:func:`blockwise_lifted_multicut` is the lifted counterpart of +:func:`bioimage_py.segmentation.blockwise_multicut`: it solves a global lifted multicut by hierarchical +domain decomposition, distributing the per-block sub-problem solves through the runner and reducing them +in the orchestrating process. The lifted edges/costs are carried alongside the base problem -- each +block additionally solves its interior lifted edges (falling back to the plain multicut solve for a +block with none), and the reduce step contracts the lifted edges through the same node labeling as the +base edges. It returns a node labeling; project it back to pixels with +:func:`bioimage_py.segmentation.relabel`. +""" +from __future__ import annotations + +import shutil +import tempfile +from typing import Optional, Sequence, Tuple + +import numpy as np + +from ..runner import get_runner +from ..runner.config import RunnerConfig +from ..sources import SourceLike, as_source +from ._blockwise_common import (as_undirected_graph, lifted_solver_by_name, reduce_problem, + solve_subproblems_distributed) + +__all__ = ["blockwise_lifted_multicut"] + + +def blockwise_lifted_multicut( + graph, + costs: np.ndarray, + lifted_uv_ids: np.ndarray, + lifted_costs: np.ndarray, + segmentation: SourceLike, + *, + block_shape: Tuple[int, ...], + n_levels: int = 1, + internal_solver: str = "kernighan-lin", + halo: Optional[Sequence[int]] = None, + num_workers: int = 1, + job_type: str = "local", + job_config: Optional[RunnerConfig] = None, +) -> np.ndarray: + """Solve the lifted multicut problem block-wise (hierarchical domain decomposition). + + The per-block sub-problem solves are the distributed map step (via the runner); the union-find + merge, edge (and lifted-edge) contraction and final global lifted solve run in the orchestrating + process. + + Args: + graph: The base graph (or region adjacency graph). Its node ids must equal the segment ids of + ``segmentation`` (node ``i`` is segment ``i``). + costs: The base edge costs (aligned to the base edge ids). + lifted_uv_ids: The lifted edges as a ``(n_lifted, 2)`` array of node-id pairs. + lifted_costs: The lifted edge costs, aligned to ``lifted_uv_ids``. + segmentation: The segmentation whose ids are the graph nodes (a numpy/zarr/n5 array or a + `Source`). It defines the blocking domain; for distributed backends it must be file-backed + (reopenable). It is only read. + block_shape: The block shape of the finest (level 0) decomposition. Required. The block size is + doubled each level and clamped to the volume shape. + n_levels: The number of hierarchy levels (>= 1). + internal_solver: The solver used for the per-block sub-problems and the final global solve; one + of ``"kernighan-lin"`` or ``"greedy-additive"``. + halo: Optional per-axis halo added to each block when reading the segmentation. ``None`` reads + the block region only. + num_workers: Number of parallel workers (threads for ``local``, tasks for distributed + backends). + job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``. + job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`). + + Returns: + The node label solution (a dense 1D array indexed by node id == segment id). + + Raises: + ValueError: If ``n_levels < 1`` or ``block_shape`` does not match the segmentation ndim. + """ + if n_levels < 1: + raise ValueError(f"n_levels must be >= 1, got {n_levels}.") + + graph = as_undirected_graph(graph) + costs = np.asarray(costs) + lifted_uv_ids = np.asarray(lifted_uv_ids, dtype="uint64").reshape(-1, 2) + lifted_costs = np.asarray(lifted_costs, dtype="float64") + seg_src = as_source(segmentation) + shape = tuple(int(s) for s in seg_src.shape) + block_shape = tuple(int(b) for b in block_shape) + if len(block_shape) != len(shape): + raise ValueError( + f"block_shape {block_shape} must match the segmentation ndim {len(shape)}." + ) + + runner = get_runner(job_type, job_config) + tmp_dir: Optional[str] = None + if job_type != "local": + tmp_dir = tempfile.mkdtemp(prefix="bioimage_py_blockwise_lmc_", dir=runner.config.tmp_root) + + graph_, costs_ = graph, costs + lifted_uv_, lifted_costs_ = lifted_uv_ids, lifted_costs + labels: Optional[np.ndarray] = None + for level in range(n_levels): + level_block_shape = tuple(min(b * (2 ** level), s) for b, s in zip(block_shape, shape)) + merge_mask = solve_subproblems_distributed( + runner, graph_, costs_, seg_src, labels, + block_shape=level_block_shape, halo=halo, internal_solver=internal_solver, + num_workers=num_workers, job_type=job_type, tmp_dir=tmp_dir, + name=f"blockwise-lifted-multicut-l{level}", + lifted_uv_ids=lifted_uv_, lifted_costs=lifted_costs_, + ) + graph_, costs_, new_labels, lifted_uv_, lifted_costs_ = reduce_problem( + graph_, costs_, merge_mask, lifted_uv_ids=lifted_uv_, lifted_costs=lifted_costs_, + ) + labels = new_labels if labels is None else new_labels[labels] + + final_labels = np.asarray( + lifted_solver_by_name(internal_solver)(graph_, costs_, lifted_uv_, lifted_costs_) + ) + node_labels = final_labels[labels] + + if tmp_dir is not None: # preserved on failure (an exception above skips this) for debugging. + shutil.rmtree(tmp_dir, ignore_errors=True) + return node_labels diff --git a/bioimage_py/segmentation/blockwise_multicut.py b/bioimage_py/segmentation/blockwise_multicut.py new file mode 100644 index 0000000..1b511c9 --- /dev/null +++ b/bioimage_py/segmentation/blockwise_multicut.py @@ -0,0 +1,117 @@ +"""Block-wise (hierarchical) multicut with a distributed per-block solve. + +:func:`blockwise_multicut` solves a global multicut by hierarchical domain decomposition: at each level +every block solves its induced sub-problem independently (the *map* step, distributed across the +``local`` / ``subprocess`` / ``slurm`` backends via the runner), and the per-block cut decisions are +merged into a coarser graph in the orchestrating process (the *reduce* step). The block size doubles +each level, so block-boundary edges that could not be decided locally become interior edges of the next +level; a final global solve on the fully contracted graph completes the partition. + +This scales the whole-graph solvers in :mod:`bioimage_py.segmentation.multicut` to problems whose +per-block solves are too many for a single process, at the cost of an approximation (the decomposition +fixes block-boundary edges as cut at each level). It returns a node labeling; project it back to pixels +with :func:`bioimage_py.segmentation.relabel` (``relabel(segmentation, node_labels, output=...)``), the +canonical distributed node-label writer. +""" +from __future__ import annotations + +import shutil +import tempfile +from typing import Optional, Sequence, Tuple + +import numpy as np + +from ..runner import get_runner +from ..runner.config import RunnerConfig +from ..sources import SourceLike, as_source +from ._blockwise_common import (as_undirected_graph, multicut_solver_by_name, reduce_problem, + solve_subproblems_distributed) + +__all__ = ["blockwise_multicut"] + + +def blockwise_multicut( + graph, + costs: np.ndarray, + segmentation: SourceLike, + *, + block_shape: Tuple[int, ...], + n_levels: int = 1, + internal_solver: str = "kernighan-lin", + halo: Optional[Sequence[int]] = None, + num_workers: int = 1, + job_type: str = "local", + job_config: Optional[RunnerConfig] = None, +) -> np.ndarray: + """Solve the multicut problem block-wise (hierarchical domain decomposition). + + The per-block sub-problem solves are the distributed map step (via the runner); the union-find + merge, edge contraction and final global solve run in the orchestrating process. + + Args: + graph: The graph (or region adjacency graph) of the multicut problem. Its node ids must equal + the segment ids of ``segmentation`` (node ``i`` is segment ``i``) -- the contract of + :func:`bioimage_py.graph.distributed_rag`. + costs: The edge costs of the multicut problem (aligned to the graph edge ids; positive costs + attractive, negative repulsive). + segmentation: The segmentation whose ids are the graph nodes (a numpy/zarr/n5 array or a + `Source`). It defines the blocking domain; for distributed backends it must be file-backed + (reopenable). It is only read. + block_shape: The block shape of the finest (level 0) decomposition. Required. The block size is + doubled each level and clamped to the volume shape. + n_levels: The number of hierarchy levels (>= 1). More levels contract more block-boundary edges + before the final global solve. + internal_solver: The solver used for the per-block sub-problems and the final global solve; one + of ``"kernighan-lin"``, ``"greedy-additive"`` or ``"decomposition"``. + halo: Optional per-axis halo added to each block when reading the segmentation (block node + membership then includes the neighborhood). ``None`` reads the block region only. + num_workers: Number of parallel workers (threads for ``local``, tasks for distributed + backends). + job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``. + job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`). + + Returns: + The node label solution (a dense 1D array indexed by node id == segment id). + + Raises: + ValueError: If ``n_levels < 1`` or ``block_shape`` does not match the segmentation ndim. + """ + if n_levels < 1: + raise ValueError(f"n_levels must be >= 1, got {n_levels}.") + + graph = as_undirected_graph(graph) + costs = np.asarray(costs) + seg_src = as_source(segmentation) + shape = tuple(int(s) for s in seg_src.shape) + block_shape = tuple(int(b) for b in block_shape) + if len(block_shape) != len(shape): + raise ValueError( + f"block_shape {block_shape} must match the segmentation ndim {len(shape)}." + ) + + runner = get_runner(job_type, job_config) + tmp_dir: Optional[str] = None + if job_type != "local": + tmp_dir = tempfile.mkdtemp(prefix="bioimage_py_blockwise_mc_", dir=runner.config.tmp_root) + + graph_, costs_ = graph, costs + labels: Optional[np.ndarray] = None + for level in range(n_levels): + # The block size doubles each level and is clamped to the volume (a block never exceeds it). + level_block_shape = tuple(min(b * (2 ** level), s) for b, s in zip(block_shape, shape)) + merge_mask = solve_subproblems_distributed( + runner, graph_, costs_, seg_src, labels, + block_shape=level_block_shape, halo=halo, internal_solver=internal_solver, + num_workers=num_workers, job_type=job_type, tmp_dir=tmp_dir, + name=f"blockwise-multicut-l{level}", + ) + graph_, costs_, new_labels = reduce_problem(graph_, costs_, merge_mask) + labels = new_labels if labels is None else new_labels[labels] + + # Final global solve on the fully contracted graph, then compose back to the original node ids. + final_labels = np.asarray(multicut_solver_by_name(internal_solver)(graph_, costs_)) + node_labels = final_labels[labels] + + if tmp_dir is not None: # preserved on failure (an exception above skips this) for debugging. + shutil.rmtree(tmp_dir, ignore_errors=True) + return node_labels diff --git a/bioimage_py/segmentation/lifted_multicut.py b/bioimage_py/segmentation/lifted_multicut.py new file mode 100644 index 0000000..06cf9a5 --- /dev/null +++ b/bioimage_py/segmentation/lifted_multicut.py @@ -0,0 +1,269 @@ +"""Lifted multicut solver functionality (backed by :mod:`bioimage_cpp.graph.lifted_multicut`). + +The lifted multicut extends the multicut (see :mod:`bioimage_py.segmentation.multicut`) with an +additional set of *lifted* edges. A lifted edge connects two nodes that need not be adjacent in the +base graph and contributes its cost to the objective iff its endpoints land in different partition +elements -- exactly like a regular edge -- but it is not itself part of the graph whose cut is being +optimized. This lets long-range attraction / repulsion (e.g. from an external segmentation or class +probabilities) bias the partition without adding spurious short-range connectivity. + +The problem is described by the base graph, the base ``costs`` (one per base edge), and a separate pair +of arrays for the lifted problem: ``lifted_uv_ids`` (``(n_lifted, 2)`` node-id pairs) and +``lifted_costs`` (one per lifted edge). All solvers return a node labeling (a dense 1D array indexed by +node id); positive costs are attractive, negative costs repulsive, as for the plain multicut. +""" +from __future__ import annotations + +import functools +from typing import Callable, Optional, Tuple + +import bioimage_cpp as bic +import numpy as np + +__all__ = [ + "lifted_multicut_kernighan_lin", + "lifted_multicut_gaec", + "lifted_multicut_fusion_moves", + "get_lifted_multicut_solver", + "lifted_edges_from_node_labels", + "lifted_problem_from_node_labels", +] + + +def _to_lifted_objective(graph, costs: np.ndarray, lifted_uv_ids: np.ndarray, + lifted_costs: np.ndarray): + """Build a ``bic.graph.lifted_multicut.LiftedMulticutObjective`` from a graph, costs and lifted arrays.""" + if isinstance(graph, bic.graph.UndirectedGraph): + graph_ = graph + else: + graph_ = bic.graph.UndirectedGraph.from_edges( + graph.number_of_nodes, np.asarray(graph.uv_ids(), dtype="uint64"), + ) + return bic.graph.lifted_multicut.LiftedMulticutObjective( + graph_, costs, + lifted_uvs=np.asarray(lifted_uv_ids, dtype="uint64"), + lifted_costs=np.asarray(lifted_costs), + ) + + +def _get_lifted_solver(internal_solver: str): + """Return a fresh ``bic.graph.lifted_multicut`` solver instance for the given name.""" + if internal_solver == "kernighan-lin": + return bic.graph.lifted_multicut.LiftedKernighanLinMulticut() + elif internal_solver == "greedy-additive": + return bic.graph.lifted_multicut.LiftedGreedyAdditiveMulticut() + else: + raise ValueError(f"{internal_solver} cannot be used as internal lifted solver.") + + +def lifted_multicut_kernighan_lin( + graph, + costs: np.ndarray, + lifted_uv_ids: np.ndarray, + lifted_costs: np.ndarray, + warmstart: bool = True, +) -> np.ndarray: + """Solve the lifted multicut problem with the Kernighan-Lin solver. + + Args: + graph: The base graph (or region adjacency graph) of the lifted multicut problem. + costs: The base edge costs, aligned to the base edge ids. + lifted_uv_ids: The lifted edges as a ``(n_lifted, 2)`` array of node-id pairs. + lifted_costs: The lifted edge costs, aligned to ``lifted_uv_ids``. + warmstart: Whether to warmstart with the lifted greedy-additive solution. + + Returns: + The node label solution to the lifted multicut problem. + """ + objective = _to_lifted_objective(graph, costs, lifted_uv_ids, lifted_costs) + if warmstart: + solver = bic.graph.lifted_multicut.LiftedChainedSolvers([ + bic.graph.lifted_multicut.LiftedGreedyAdditiveMulticut(), + bic.graph.lifted_multicut.LiftedKernighanLinMulticut(), + ]) + else: + solver = bic.graph.lifted_multicut.LiftedKernighanLinMulticut() + return solver.optimize(objective) + + +def lifted_multicut_gaec( + graph, + costs: np.ndarray, + lifted_uv_ids: np.ndarray, + lifted_costs: np.ndarray, +) -> np.ndarray: + """Solve the lifted multicut problem with the greedy-additive edge contraction solver. + + Args: + graph: The base graph (or region adjacency graph) of the lifted multicut problem. + costs: The base edge costs, aligned to the base edge ids. + lifted_uv_ids: The lifted edges as a ``(n_lifted, 2)`` array of node-id pairs. + lifted_costs: The lifted edge costs, aligned to ``lifted_uv_ids``. + + Returns: + The node label solution to the lifted multicut problem. + """ + objective = _to_lifted_objective(graph, costs, lifted_uv_ids, lifted_costs) + return bic.graph.lifted_multicut.LiftedGreedyAdditiveMulticut().optimize(objective) + + +def lifted_multicut_fusion_moves( + graph, + costs: np.ndarray, + lifted_uv_ids: np.ndarray, + lifted_costs: np.ndarray, + n_threads: int = 1, + warmstart: bool = True, + sigma: float = 2.0, + seed_fraction: float = 0.05, + num_it: int = 10, + num_it_stop: int = 4, +) -> np.ndarray: + """Solve the lifted multicut problem with the fusion-moves solver. + + Fusion moves generate proposal solutions (via a watershed proposal generator) and fuse them with + the current best solution, escaping the local optima that the greedy / KL solvers can get stuck in. + + Args: + graph: The base graph (or region adjacency graph) of the lifted multicut problem. + costs: The base edge costs, aligned to the base edge ids. + lifted_uv_ids: The lifted edges as a ``(n_lifted, 2)`` array of node-id pairs. + lifted_costs: The lifted edge costs, aligned to ``lifted_uv_ids``. + n_threads: The number of threads used to generate and fuse proposals in parallel. + warmstart: Whether to warmstart with the lifted greedy-additive and Kernighan-Lin solutions. + sigma: The noise sigma of the watershed proposal generator. + seed_fraction: The fraction of nodes used as seeds by the watershed proposal generator. + num_it: The number of fusion-move iterations. + num_it_stop: Stop after this many iterations without improvement. + + Returns: + The node label solution to the lifted multicut problem. + """ + objective = _to_lifted_objective(graph, costs, lifted_uv_ids, lifted_costs) + proposal_generator = bic.graph.lifted_multicut.WatershedProposalGenerator( + sigma=sigma, n_seeds_fraction=seed_fraction, + ) + solver = bic.graph.lifted_multicut.FusionMoveLiftedMulticut( + proposal_generator=proposal_generator, + sub_solver=bic.graph.lifted_multicut.LiftedKernighanLinMulticut(), + number_of_iterations=num_it, + stop_if_no_improvement=num_it_stop, + number_of_threads=n_threads, + ) + if warmstart: + solver = bic.graph.lifted_multicut.LiftedChainedSolvers([ + bic.graph.lifted_multicut.LiftedGreedyAdditiveMulticut(), + bic.graph.lifted_multicut.LiftedKernighanLinMulticut(), + solver, + ]) + return solver.optimize(objective) + + +# Registry of the lifted multicut solvers (name -> function). +_solvers = { + "kernighan-lin": lifted_multicut_kernighan_lin, + "greedy-additive": lifted_multicut_gaec, + "fusion-moves": lifted_multicut_fusion_moves, +} + + +def get_lifted_multicut_solver(name: str, **kwargs) -> Callable[..., np.ndarray]: + """Get a lifted multicut solver by name, with solver keyword arguments bound. + + Args: + name: The name of the solver; one of ``"kernighan-lin"``, ``"greedy-additive"`` or + ``"fusion-moves"``. + kwargs: Keyword arguments bound to the solver (e.g. ``warmstart`` or ``n_threads``). + + Returns: + A callable ``solver(graph, costs, lifted_uv_ids, lifted_costs)`` returning the node labeling. + + Raises: + ValueError: If ``name`` is not a known lifted multicut solver. + """ + if name not in _solvers: + names = ", ".join(_solvers.keys()) + raise ValueError(f"Lifted multicut solver must be one of ({names}), got {name!r}.") + return functools.partial(_solvers[name], **kwargs) + + +def lifted_edges_from_node_labels( + graph, + node_labels: np.ndarray, + graph_depth: int, + mode: str = "different", + ignore_label: Optional[int] = 0, + n_threads: Optional[int] = None, +) -> np.ndarray: + """Discover lifted edges from per-node labels within a graph neighborhood. + + A lifted edge is inserted between every pair of nodes that are within ``graph_depth`` hops of each + other in the base graph (excluding the base edges themselves) and whose labels satisfy ``mode``. + + Args: + graph: The base graph (or region adjacency graph). + node_labels: A per-node label array (length ``graph.number_of_nodes``). + graph_depth: The maximum number of graph hops between the endpoints of a lifted edge. + mode: Which node-label pairs to connect; one of ``"all"``, ``"same"`` or ``"different"``. + ignore_label: A node label to ignore; nodes with this label are not connected. Pass ``None`` + to disable. + n_threads: The number of threads to use. Defaults to all available cores. + + Returns: + The lifted edges as a ``(n_lifted, 2)`` array of node-id pairs. + """ + if isinstance(graph, bic.graph.UndirectedGraph): + graph_ = graph + else: + graph_ = bic.graph.UndirectedGraph.from_edges( + graph.number_of_nodes, np.asarray(graph.uv_ids(), dtype="uint64"), + ) + n_threads_ = 0 if n_threads is None else int(n_threads) + return bic.graph.lifted_multicut.lifted_edges_from_node_labels( + graph_, np.asarray(node_labels), int(graph_depth), + mode=mode, ignore_label=ignore_label, number_of_threads=n_threads_, + ) + + +def lifted_problem_from_node_labels( + graph, + node_labels: np.ndarray, + graph_depth: int, + same_segment_cost: float, + different_segment_cost: float, + mode: str = "all", + ignore_label: Optional[int] = 0, + n_threads: Optional[int] = None, +) -> Tuple[np.ndarray, np.ndarray]: + """Build a lifted problem (edges and costs) from per-node labels. + + Lifted edges are discovered with :func:`lifted_edges_from_node_labels`, then each edge is assigned a + flat cost: ``same_segment_cost`` when its endpoints share a node label (attractive if positive), + ``different_segment_cost`` when they differ (repulsive if negative). This is the node-label variant + of the lifted-from-segmentation construction: derive ``node_labels`` externally (e.g. the dominant + overlap of each superpixel with an external segmentation) and pass them here. + + Args: + graph: The base graph (or region adjacency graph). + node_labels: A per-node label array (length ``graph.number_of_nodes``). + graph_depth: The maximum number of graph hops between the endpoints of a lifted edge. + same_segment_cost: The cost assigned to lifted edges whose endpoints share a label. + different_segment_cost: The cost assigned to lifted edges whose endpoints differ in label. + mode: Which node-label pairs to connect; one of ``"all"``, ``"same"`` or ``"different"``. + ignore_label: A node label to ignore; nodes with this label are not connected. Pass ``None`` + to disable. + n_threads: The number of threads to use. Defaults to all available cores. + + Returns: + A ``(lifted_uv_ids, lifted_costs)`` tuple: the ``(n_lifted, 2)`` lifted edges and their costs. + """ + node_labels = np.asarray(node_labels) + lifted_uv_ids = lifted_edges_from_node_labels( + graph, node_labels, graph_depth, mode=mode, ignore_label=ignore_label, n_threads=n_threads, + ) + lifted_uv_ids = np.asarray(lifted_uv_ids) + if len(lifted_uv_ids) == 0: + return lifted_uv_ids.reshape(0, 2).astype("uint64"), np.zeros(0, dtype="float64") + same = node_labels[lifted_uv_ids[:, 0]] == node_labels[lifted_uv_ids[:, 1]] + lifted_costs = np.where(same, float(same_segment_cost), float(different_segment_cost)) + return lifted_uv_ids, lifted_costs.astype("float64") diff --git a/bioimage_py/segmentation/stitching.py b/bioimage_py/segmentation/stitching.py index 9b8d129..b0508cf 100644 --- a/bioimage_py/segmentation/stitching.py +++ b/bioimage_py/segmentation/stitching.py @@ -29,7 +29,7 @@ from ..runner import get_runner from ..runner.config import RunnerConfig -from ..sources import Source, SourceLike, as_source, from_spec +from ..sources import Source, SourceLike, as_source, capture_source, resolve_source from ..util import BlockDescriptor, ComputeFn, full_roi, get_blocking, to_roi from .multicut import compute_edge_costs, multicut_decomposition from .relabel import relabel @@ -140,7 +140,7 @@ def _make_seg_overlap(shape: Tuple[int, ...], tile_shape: Tuple[int, ...], def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source], mask: Optional[Source]) -> Optional[np.ndarray]: - store = _resolve_source(store_handle) + store = resolve_source(store_handle) blocking = get_blocking(shape, tile_shape) block_id = blocking.coordinates_to_block_id([int(c) for c in block.begin]) rows: List[np.ndarray] = [] @@ -181,8 +181,8 @@ def _make_segment(shape: Tuple[int, ...], tile_shape: Tuple[int, ...], tile_over def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence[Source], mask: Optional[Source]) -> Optional[np.ndarray]: - input_ = _resolve_source(input_handle) - store = _resolve_source(store_handle) + input_ = resolve_source(input_handle) + store = resolve_source(store_handle) output_ = outputs[0] blocking = get_blocking(shape, tile_shape) block_id = blocking.coordinates_to_block_id([int(c) for c in block.inner_block.begin]) @@ -215,20 +215,6 @@ def _compute(block: BlockDescriptor, inputs: Sequence[Source], outputs: Sequence # Orchestration helpers. # --------------------------------------------------------------------------------------------- -def _resolve_source(handle) -> Source: - """Resolve a closure-captured handle to a live `Source` (already-live for local, spec otherwise).""" - return handle if isinstance(handle, Source) else from_spec(handle) - - -def _capture(source: Source, job_type: str): - """Capture a source for a per-block closure: the live object for local, its spec otherwise. - - For distributed backends ``to_spec()`` also validates the source is reopenable (raising a clear - error for in-memory numpy arrays). - """ - return source if job_type == "local" else source.to_spec() - - def _prepare_output(output: Optional[SourceLike], shape: Tuple[int, ...], job_type: str) -> SourceLike: """Resolve the output array: allocate a numpy array for local, require a file-backed one otherwise.""" if output is not None: @@ -468,7 +454,7 @@ def stitch_segmentation( "uint64. Reduce the tile shape or the volume size.") runner = get_runner(job_type, job_config) - input_handle = _capture(src, job_type) + input_handle = capture_source(src, job_type) store, store_cleanup, store_handle = _make_tile_store(n_blocks, max_halo, job_type, job_config) try: # Stage 1: segment each haloed tile, offset its ids, write the inner block + the store slot. diff --git a/bioimage_py/sources/__init__.py b/bioimage_py/sources/__init__.py index b9b04ed..9dec7c1 100644 --- a/bioimage_py/sources/__init__.py +++ b/bioimage_py/sources/__init__.py @@ -2,7 +2,8 @@ from .array_source import ArraySource from .base import Source, SourceSpec from .cloudvolume_source import CloudVolumeSource, open_cloudvolume -from .dispatch import SourceLike, as_source, from_spec, register_source +from .dispatch import (SourceLike, as_source, capture_source, from_spec, register_source, + resolve_source) from .file_source import FileSource, open_source from .webknossos_source import WebKnossosSource, open_webknossos @@ -12,8 +13,10 @@ "SourceSpec", "SourceLike", "as_source", + "capture_source", "from_spec", "register_source", + "resolve_source", "FileSource", "open_source", "CloudVolumeSource", diff --git a/bioimage_py/sources/dispatch.py b/bioimage_py/sources/dispatch.py index 22d569b..0c6111b 100644 --- a/bioimage_py/sources/dispatch.py +++ b/bioimage_py/sources/dispatch.py @@ -57,6 +57,40 @@ def as_source(obj: "SourceLike") -> Source: raise TypeError(f"Cannot convert object of type {type(obj)!r} to a Source.") +def capture_source(source: Source, job_type: str) -> Union[Source, SourceSpec]: + """Capture a source for a per-block closure: the live object for local, its spec otherwise. + + Per-block compute functions are cloudpickled for distributed backends, so a captured source + must be picklable and reopenable on the worker. For the ``local`` backend the live `Source` is + captured directly; for distributed backends ``to_spec()`` is captured, which also validates that + the source is reopenable (raising a clear error for in-memory numpy arrays). Reconstruct the + captured handle on the worker with :func:`resolve_source`. + + Args: + source: The source to capture. + job_type: The execution backend (``"local"`` / ``"subprocess"`` / ``"slurm"``). + + Returns: + The live `Source` for ``local``, otherwise its `SourceSpec`. + """ + return source if job_type == "local" else source.to_spec() + + +def resolve_source(handle: Union[Source, SourceSpec]) -> Source: + """Resolve a closure-captured handle to a live `Source`. + + The inverse of :func:`capture_source`: an already-live `Source` (the ``local`` case) is returned + unchanged, while a `SourceSpec` (the distributed case) is reopened with :func:`from_spec`. + + Args: + handle: The captured handle (a live `Source` or a `SourceSpec`). + + Returns: + A live `Source`. + """ + return handle if isinstance(handle, Source) else from_spec(handle) + + def from_spec(spec: SourceSpec) -> Source: """Reconstruct a :class:`Source` from its :class:`SourceSpec`.""" if spec.kind in ("zarr", "z5py"): diff --git a/tests/graph/test_distributed_rag.py b/tests/graph/test_distributed_rag.py new file mode 100644 index 0000000..f464bca --- /dev/null +++ b/tests/graph/test_distributed_rag.py @@ -0,0 +1,227 @@ +"""Parity tests for the distributed RAG / edge-feature orchestration. + +The core guarantee mirrors ``tests/test_runner_parity.py`` and the bioimage-cpp reference +(``tests/graph/test_distributed_rag.py``): computing the graph / features block-wise and merging +reproduces the whole-volume result -- exactly for the region graph and the ``size/min/max`` feature +columns, and to floating-point tolerance for ``mean/std`` -- and ``direct == LocalRunner == +SubprocessRunner``. +""" +import numpy as np +import pytest + +import bioimage_cpp as bic +import bioimage_py as bp + +LABEL_DTYPES = ["uint32", "uint64", "int32", "int64"] + +# In-core complex feature columns are (mean, median, std, min, max, p5, p10, p25, p75, p90, p95, +# size); the distributed complex output is the moment subset (mean, std, min, max, size). +_COMPLEX_MOMENT_COLUMNS = [0, 2, 3, 4, 11] +# Distributed complex columns that are exact across backends / block order (min, max, size) and +# those that are only float-order stable (mean, std). +_EXACT_COMPLEX = [2, 3, 4] +_ALLCLOSE_COMPLEX = [0, 1] + + +def _assert_complex_matches_whole(feats, whole_complex): + """Assert distributed complex features equal the in-core moment columns.""" + for dist_col, whole_col in zip(range(5), _COMPLEX_MOMENT_COLUMNS): + if dist_col in _EXACT_COMPLEX: + np.testing.assert_array_equal(feats[:, dist_col], whole_complex[:, whole_col]) + else: + np.testing.assert_allclose(feats[:, dist_col], whole_complex[:, whole_col], + rtol=1e-6, atol=1e-9) + + +def _assert_features_backend_equal(a, b, complex_features): + """Assert two distributed feature arrays agree (exact on size/min/max, allclose on mean/std).""" + if complex_features: + for col in _EXACT_COMPLEX: + np.testing.assert_array_equal(a[:, col], b[:, col]) + for col in _ALLCLOSE_COMPLEX: + np.testing.assert_allclose(a[:, col], b[:, col], rtol=1e-6, atol=1e-9) + else: # (mean, size) + np.testing.assert_array_equal(a[:, 1], b[:, 1]) + np.testing.assert_allclose(a[:, 0], b[:, 0], rtol=1e-6, atol=1e-9) + + +# -------------------------------------------------------------------------------------------------- +# Region adjacency graph +# -------------------------------------------------------------------------------------------------- + +@pytest.mark.parametrize("dtype", LABEL_DTYPES) +def test_rag_equality_2d(dtype, rng): + labels = rng.integers(0, 7, size=(18, 21)).astype(dtype) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + graph = bp.graph.distributed_rag(labels) # direct + np.testing.assert_array_equal(graph.uv_ids(), whole) + assert graph.number_of_nodes == int(labels.max()) + 1 + + +def test_rag_equality_3d(rng): + labels = rng.integers(0, 5, size=(7, 8, 9)).astype("uint64") + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + graph = bp.graph.distributed_rag(labels) + np.testing.assert_array_equal(graph.uv_ids(), whole) + + +@pytest.mark.parametrize("block_shape", [(8, 8), (7, 5)]) # a divisor and a non-divisor of the shape +def test_rag_backend_parity(block_shape, zarr_factory, rng): + labels = rng.integers(0, 6, size=(24, 20)).astype("uint64") + z = zarr_factory(labels, chunks=(8, 8)) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + + direct = bp.graph.distributed_rag(labels) # local, 1 worker, no block shape + np.testing.assert_array_equal(direct.uv_ids(), whole) + for nw, job in [(4, "local"), (3, "subprocess")]: + graph = bp.graph.distributed_rag(z, block_shape=block_shape, num_workers=nw, job_type=job) + np.testing.assert_array_equal(graph.uv_ids(), whole, + err_msg=f"rag mismatch nw={nw} job={job} bs={block_shape}") + + +# -------------------------------------------------------------------------------------------------- +# Edge-map features +# -------------------------------------------------------------------------------------------------- + +@pytest.mark.parametrize("dtype", ["uint32", "int64"]) +def test_edge_map_features_2d(dtype, rng): + labels = rng.integers(0, 6, size=(12, 15)).astype(dtype) + edge_map = rng.standard_normal((12, 15)).astype("float64") + rag = bic.graph.region_adjacency_graph(labels) + + graph, simple = bp.graph.distributed_edge_features(labels, edge_map, complex_features=False) + np.testing.assert_array_equal(graph.uv_ids(), rag.uv_ids()) + whole_simple = bic.graph.features.edge_map_features(rag, labels, edge_map) + np.testing.assert_array_equal(simple[:, 1], whole_simple[:, 1]) # size exact + np.testing.assert_allclose(simple[:, 0], whole_simple[:, 0], rtol=1e-6, atol=1e-9) # mean + + _, complex_ = bp.graph.distributed_edge_features(labels, edge_map, complex_features=True) + _assert_complex_matches_whole(complex_, bic.graph.features.edge_map_features_complex( + rag, labels, edge_map)) + + +def test_edge_map_features_3d(rng): + labels = rng.integers(0, 5, size=(6, 7, 8)).astype("uint64") + edge_map = rng.standard_normal((6, 7, 8)).astype("float64") + rag = bic.graph.region_adjacency_graph(labels) + _, complex_ = bp.graph.distributed_edge_features(labels, edge_map, complex_features=True) + _assert_complex_matches_whole(complex_, bic.graph.features.edge_map_features_complex( + rag, labels, edge_map)) + + +@pytest.mark.parametrize("complex_features", [False, True]) +def test_edge_map_backend_parity(complex_features, zarr_factory, rng): + labels = rng.integers(0, 6, size=(24, 20)).astype("uint64") + edge_map = rng.standard_normal((24, 20)).astype("float64") + zl = zarr_factory(labels, chunks=(8, 8)) + ze = zarr_factory(edge_map, chunks=(8, 8)) + + _, direct = bp.graph.distributed_edge_features(labels, edge_map, complex_features=complex_features) + for nw, job in [(4, "local"), (3, "subprocess")]: + _, feats = bp.graph.distributed_edge_features( + zl, ze, complex_features=complex_features, block_shape=(8, 8), num_workers=nw, + job_type=job) + _assert_features_backend_equal(direct, feats, complex_features) + + +# -------------------------------------------------------------------------------------------------- +# Affinity features +# -------------------------------------------------------------------------------------------------- + +@pytest.mark.parametrize("offsets", [[[1, 0], [0, 1]], [[1, 0], [0, 1], [3, 0], [0, 3]]]) +def test_affinity_features_2d(offsets, rng): + labels = rng.integers(0, 6, size=(13, 14)).astype("uint32") + aff = rng.standard_normal((len(offsets),) + labels.shape).astype("float64") + rag = bic.graph.region_adjacency_graph(labels) + + graph, complex_ = bp.graph.distributed_edge_features( + labels, aff, feature_type="affinity", offsets=offsets, complex_features=True) + # The graph stays the nearest-neighbor RAG even with long-range offsets. + np.testing.assert_array_equal(graph.uv_ids(), rag.uv_ids()) + _assert_complex_matches_whole(complex_, bic.graph.features.affinity_features_complex( + rag, labels, aff, offsets)) + + +def test_affinity_features_3d(rng): + labels = rng.integers(0, 5, size=(5, 6, 7)).astype("uint64") + offsets = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 2, 0]] + aff = rng.standard_normal((len(offsets),) + labels.shape).astype("float64") + rag = bic.graph.region_adjacency_graph(labels) + graph, complex_ = bp.graph.distributed_edge_features( + labels, aff, feature_type="affinity", offsets=offsets, complex_features=True) + np.testing.assert_array_equal(graph.uv_ids(), rag.uv_ids()) + _assert_complex_matches_whole(complex_, bic.graph.features.affinity_features_complex( + rag, labels, aff, offsets)) + + +def test_affinity_backend_parity(zarr_factory, rng): + offsets = [[1, 0], [0, 1], [3, 0], [0, 3]] # includes long-range offsets (halo > block-face) + labels = rng.integers(0, 6, size=(24, 20)).astype("uint64") + aff = rng.standard_normal((len(offsets),) + labels.shape).astype("float64") + zl = zarr_factory(labels, chunks=(8, 8)) + za = zarr_factory(aff, chunks=(len(offsets), 8, 8)) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + + _, direct = bp.graph.distributed_edge_features( + labels, aff, feature_type="affinity", offsets=offsets, complex_features=True) + for nw, job in [(4, "local"), (3, "subprocess")]: + graph, feats = bp.graph.distributed_edge_features( + zl, za, feature_type="affinity", offsets=offsets, complex_features=True, + block_shape=(8, 8), num_workers=nw, job_type=job) + np.testing.assert_array_equal(graph.uv_ids(), whole, err_msg=f"nw={nw} job={job}") + _assert_features_backend_equal(direct, feats, True) + + +# -------------------------------------------------------------------------------------------------- +# Edge cases +# -------------------------------------------------------------------------------------------------- + +def test_single_label_empty_graph(): + labels = np.zeros((10, 12), dtype="uint32") + graph = bp.graph.distributed_rag(labels) + assert graph.number_of_nodes == 1 + assert graph.number_of_edges == 0 + + edge_map = np.ones((10, 12), dtype="float64") + graph2, simple = bp.graph.distributed_edge_features(labels, edge_map) + assert graph2.number_of_edges == 0 + assert simple.shape == (0, 2) + _, complex_ = bp.graph.distributed_edge_features(labels, edge_map, complex_features=True) + assert complex_.shape == (0, 5) + + +def test_isolated_segments_kept_as_nodes(): + # Only ids 0 and 5 are present (adjacent), so nodes 1..4 are isolated; n_nodes must be max+1=6. + labels = np.zeros((6, 6), dtype="uint32") + labels[:, 3:] = 5 + graph = bp.graph.distributed_rag(labels) + assert graph.number_of_nodes == 6 + np.testing.assert_array_equal(graph.uv_ids(), np.array([[0, 5]], dtype="uint64")) + + +def test_unsupported_dtype_is_cast(rng): + labels = rng.integers(0, 6, size=(12, 12)).astype("uint16") + whole = bic.graph.region_adjacency_graph(labels.astype("uint32")).uv_ids() + graph = bp.graph.distributed_rag(labels) + np.testing.assert_array_equal(graph.uv_ids(), whole) + + +def test_n_nodes_explicit_and_too_small(rng): + labels = rng.integers(0, 6, size=(12, 12)).astype("uint32") + graph = bp.graph.distributed_rag(labels, n_nodes=100) + assert graph.number_of_nodes == 100 + with pytest.raises(ValueError, match="n_nodes"): + bp.graph.distributed_rag(labels, n_nodes=3) + + +def test_invalid_arguments_raise(rng): + labels = rng.integers(0, 4, size=(8, 8)).astype("uint32") + data = rng.standard_normal((8, 8)).astype("float64") + with pytest.raises(ValueError, match="feature_type"): + bp.graph.distributed_edge_features(labels, data, feature_type="bogus") + with pytest.raises(ValueError, match="offsets is required"): + bp.graph.distributed_edge_features(labels, data, feature_type="affinity") + with pytest.raises(ValueError, match="must match the labels shape"): + bp.graph.distributed_edge_features(labels, rng.standard_normal((8, 7)).astype("float64")) + with pytest.raises(ValueError, match="integer"): + bp.graph.distributed_rag(data) diff --git a/tests/graph/test_distributed_rag_slurm.py b/tests/graph/test_distributed_rag_slurm.py new file mode 100644 index 0000000..c0ffb84 --- /dev/null +++ b/tests/graph/test_distributed_rag_slurm.py @@ -0,0 +1,102 @@ +"""Cluster-gated slurm parity tests for the distributed RAG / edge-feature orchestration. + +These submit real sbatch array jobs, so they are skipped unless ``sbatch`` is on ``PATH`` and +``BIOIMAGE_PY_SHARED_TMP`` points at a shared filesystem (same gate as ``tests/test_slurm_runner.py``). +They assert that the slurm backend reproduces the whole-volume graph / edge features exactly for the +region graph and the ``size/min/max`` columns, and to floating-point tolerance for ``mean/std`` -- +extending the ``local`` / ``subprocess`` coverage in ``tests/graph/test_distributed_rag.py`` to slurm. + +Partition/account default to this cluster's CPU test queue and can be overridden via the +``BIOIMAGE_PY_SLURM_PARTITION`` / ``BIOIMAGE_PY_SLURM_ACCOUNT`` env vars. +""" +import os +import shutil + +import numpy as np +import pytest + +import bioimage_cpp as bic +import bioimage_py as bp +from bioimage_py.runner import SlurmConfig + +# Skip unless the slurm CLI and a shared filesystem (for the job temp + test arrays) are both +# available; computed here rather than imported from conftest to avoid a fragile module import. +pytestmark = pytest.mark.skipif( + shutil.which("sbatch") is None or not os.environ.get("BIOIMAGE_PY_SHARED_TMP"), + reason="needs sbatch on PATH and BIOIMAGE_PY_SHARED_TMP set to a shared-filesystem dir", +) + +# Distributed complex feature columns are (mean, std, min, max, size); min/max/size are exact across +# backends / block order, mean/std are only float-order stable. +_EXACT_COMPLEX = [2, 3, 4] +_ALLCLOSE_COMPLEX = [0, 1] + + +def _cfg(tmp_root, **overrides): + """Build a small SlurmConfig for quick CPU test jobs (duplicated from test_slurm_runner.py).""" + params = dict( + tmp_root=tmp_root, + partition=os.environ.get("BIOIMAGE_PY_SLURM_PARTITION", "standard96:test"), + account=os.environ.get("BIOIMAGE_PY_SLURM_ACCOUNT", "nim00007"), + time="00:10:00", + cpus_per_task=1, + mem="2G", + poll_interval=5.0, + ) + params.update(overrides) + return SlurmConfig(**params) + + +def _assert_features_backend_equal(a, b, complex_features): + """Assert two distributed feature arrays agree (exact on size/min/max, allclose on mean/std).""" + if complex_features: + for col in _EXACT_COMPLEX: + np.testing.assert_array_equal(a[:, col], b[:, col]) + for col in _ALLCLOSE_COMPLEX: + np.testing.assert_allclose(a[:, col], b[:, col], rtol=1e-6, atol=1e-9) + else: # (mean, size) + np.testing.assert_array_equal(a[:, 1], b[:, 1]) + np.testing.assert_allclose(a[:, 0], b[:, 0], rtol=1e-6, atol=1e-9) + + +def test_slurm_rag_parity(shared_zarr_factory, rng, shared_tmp_path): + labels = rng.integers(0, 6, size=(24, 20)).astype("uint64") + z = shared_zarr_factory(labels, chunks=(8, 8)) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + + graph = bp.graph.distributed_rag(z, block_shape=(8, 8), num_workers=4, job_type="slurm", + job_config=_cfg(shared_tmp_path)) + np.testing.assert_array_equal(graph.uv_ids(), whole) + + +@pytest.mark.parametrize("complex_features", [False, True]) +def test_slurm_edge_map_features_parity(complex_features, shared_zarr_factory, rng, shared_tmp_path): + labels = rng.integers(0, 6, size=(24, 20)).astype("uint64") + edge_map = rng.standard_normal((24, 20)).astype("float64") + zl = shared_zarr_factory(labels, chunks=(8, 8)) + ze = shared_zarr_factory(edge_map, chunks=(8, 8)) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + + _, direct = bp.graph.distributed_edge_features(labels, edge_map, complex_features=complex_features) + graph, feats = bp.graph.distributed_edge_features( + zl, ze, complex_features=complex_features, block_shape=(8, 8), num_workers=4, + job_type="slurm", job_config=_cfg(shared_tmp_path)) + np.testing.assert_array_equal(graph.uv_ids(), whole) + _assert_features_backend_equal(direct, feats, complex_features) + + +def test_slurm_affinity_features_parity(shared_zarr_factory, rng, shared_tmp_path): + offsets = [[1, 0], [0, 1], [3, 0], [0, 3]] # includes long-range offsets (halo > block-face) + labels = rng.integers(0, 6, size=(24, 20)).astype("uint64") + aff = rng.standard_normal((len(offsets),) + labels.shape).astype("float64") + zl = shared_zarr_factory(labels, chunks=(8, 8)) + za = shared_zarr_factory(aff, chunks=(len(offsets), 8, 8)) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + + _, direct = bp.graph.distributed_edge_features( + labels, aff, feature_type="affinity", offsets=offsets, complex_features=True) + graph, feats = bp.graph.distributed_edge_features( + zl, za, feature_type="affinity", offsets=offsets, complex_features=True, + block_shape=(8, 8), num_workers=4, job_type="slurm", job_config=_cfg(shared_tmp_path)) + np.testing.assert_array_equal(graph.uv_ids(), whole) + _assert_features_backend_equal(direct, feats, True) diff --git a/tests/test_blockwise_multicut.py b/tests/test_blockwise_multicut.py new file mode 100644 index 0000000..891c583 --- /dev/null +++ b/tests/test_blockwise_multicut.py @@ -0,0 +1,161 @@ +"""Tests for block-wise (lifted) multicut. + +The core guarantees, mirroring ``tests/test_runner_parity.py``: +- the block-wise (hierarchical) solve reproduces the whole-graph multicut partition on a clean + problem, and +- ``local == subprocess`` for the distributed per-block map step (exact, across workers and levels). +""" +import numpy as np +import pytest + +import bioimage_py as bp +from bioimage_py.graph import distributed_rag +from bioimage_py.segmentation import (blockwise_lifted_multicut, blockwise_multicut, + lifted_multicut_kernighan_lin, multicut_kernighan_lin, relabel) + +# 4x4 superpixels of 8x8 (a 32x32 volume); macro clusters are 2x2 blocks of superpixels (4 macros). +_SP_SIZE = 8 +_GRID = 4 +_BLOCK_SHAPE = (16, 16) + + +def _same_partition(a, b): + """Whether two node labelings define the same partition (ids may differ).""" + a = np.asarray(a) + b = np.asarray(b) + return np.array_equal(a[:, None] == a[None, :], b[:, None] == b[None, :]) + + +def _build_problem(): + """Build ``(segmentation, macro, graph, uv_ids)`` for the macro-cluster test problem.""" + size = _SP_SIZE * _GRID + sp = np.zeros((size, size), dtype="uint32") + sid = 0 + for i in range(0, size, _SP_SIZE): + for j in range(0, size, _SP_SIZE): + sp[i:i + _SP_SIZE, j:j + _SP_SIZE] = sid + sid += 1 + ids = np.arange(sid) + macro = (ids // _GRID // 2) * (_GRID // 2) + (ids % _GRID) // 2 + graph = distributed_rag(sp, block_shape=_BLOCK_SHAPE) + return sp, macro, graph, np.asarray(graph.uv_ids()) + + +def _clean_costs(uv, macro, attractive=10.0, repulsive=-10.0): + """Attractive within a macro cluster, repulsive across -> the multicut is exactly ``macro``.""" + return np.where(macro[uv[:, 0]] == macro[uv[:, 1]], attractive, repulsive) + + +@pytest.mark.parametrize("n_levels", [1, 2, 3]) +def test_blockwise_multicut_matches_whole_graph(n_levels): + sp, macro, graph, uv = _build_problem() + costs = _clean_costs(uv, macro) + + whole = np.asarray(multicut_kernighan_lin(graph, costs)) + assert _same_partition(whole, macro) # the clean problem recovers the macro clusters + + node_labels = np.asarray(blockwise_multicut(graph, costs, sp, block_shape=_BLOCK_SHAPE, + n_levels=n_levels)) + assert len(node_labels) == graph.number_of_nodes + assert _same_partition(node_labels, whole) + assert _same_partition(node_labels, macro) + + +@pytest.mark.parametrize("n_levels", [1, 2]) +@pytest.mark.parametrize("num_workers", [1, 3]) +def test_blockwise_multicut_backend_parity(n_levels, num_workers, zarr_factory): + sp, macro, graph, uv = _build_problem() + costs = _clean_costs(uv, macro) + seg = zarr_factory(sp, chunks=_BLOCK_SHAPE) + + results = {} + for job in ("local", "subprocess"): + results[job] = np.asarray(blockwise_multicut( + graph, costs, seg, block_shape=_BLOCK_SHAPE, n_levels=n_levels, + num_workers=num_workers, job_type=job)) + # The distributed map step is deterministic -> the node labelings are bit-identical. + np.testing.assert_array_equal(results["local"], results["subprocess"]) + + +def test_blockwise_multicut_halo(): + sp, macro, graph, uv = _build_problem() + costs = _clean_costs(uv, macro) + node_labels = np.asarray(blockwise_multicut(graph, costs, sp, block_shape=_BLOCK_SHAPE, + n_levels=2, halo=(2, 2))) + assert _same_partition(node_labels, macro) + + +def test_blockwise_multicut_then_relabel(): + # End to end: solve block-wise, then project the node labeling to pixels with relabel. + sp, macro, graph, uv = _build_problem() + costs = _clean_costs(uv, macro) + node_labels = np.asarray(blockwise_multicut(graph, costs, sp, block_shape=_BLOCK_SHAPE, + n_levels=2)).astype("uint32") + seg = np.asarray(relabel(sp, node_labels)) + assert seg.shape == sp.shape + # Projecting the node labeling to pixels is exactly node_labels indexed by the superpixel id. + np.testing.assert_array_equal(seg, node_labels[sp]) + # The resulting pixel segmentation has exactly one label per macro cluster. + assert len(np.unique(seg)) == len(np.unique(macro)) + + +# -------------------------------------------------------------------------------------------------- +# Lifted +# -------------------------------------------------------------------------------------------------- + +def _lifted_matters_problem(): + """Base weakly repulsive (plain MC = singletons); within-macro lifted edges (strongly attractive).""" + sp, macro, graph, uv = _build_problem() + costs = np.full(len(uv), -1.0) + pairs = [] + for m in np.unique(macro): + members = np.where(macro == m)[0] + for x in members[1:]: + pairs.append([members[0], x]) + lifted_uv = np.array(pairs, dtype="uint64") + lifted_costs = np.full(len(lifted_uv), 50.0) + return sp, macro, graph, costs, lifted_uv, lifted_costs + + +@pytest.mark.parametrize("n_levels", [1, 2, 3]) +def test_blockwise_lifted_matches_whole_graph(n_levels): + sp, macro, graph, costs, lifted_uv, lifted_costs = _lifted_matters_problem() + + whole = np.asarray(lifted_multicut_kernighan_lin(graph, costs, lifted_uv, lifted_costs)) + assert _same_partition(whole, macro) # the lifted edges recover the macro clusters + + node_labels = np.asarray(blockwise_lifted_multicut( + graph, costs, lifted_uv, lifted_costs, sp, block_shape=_BLOCK_SHAPE, n_levels=n_levels)) + assert _same_partition(node_labels, whole) + assert _same_partition(node_labels, macro) + + +@pytest.mark.parametrize("num_workers", [1, 3]) +def test_blockwise_lifted_backend_parity(num_workers, zarr_factory): + sp, macro, graph, costs, lifted_uv, lifted_costs = _lifted_matters_problem() + seg = zarr_factory(sp, chunks=_BLOCK_SHAPE) + + results = {} + for job in ("local", "subprocess"): + results[job] = np.asarray(blockwise_lifted_multicut( + graph, costs, lifted_uv, lifted_costs, seg, block_shape=_BLOCK_SHAPE, n_levels=2, + num_workers=num_workers, job_type=job)) + np.testing.assert_array_equal(results["local"], results["subprocess"]) + + +# -------------------------------------------------------------------------------------------------- +# Errors +# -------------------------------------------------------------------------------------------------- + +def test_blockwise_multicut_invalid_args(): + sp, macro, graph, uv = _build_problem() + costs = _clean_costs(uv, macro) + with pytest.raises(ValueError): + blockwise_multicut(graph, costs, sp, block_shape=_BLOCK_SHAPE, n_levels=0) + with pytest.raises(ValueError): + blockwise_multicut(graph, costs, sp, block_shape=(16, 16, 16), n_levels=1) + + +def test_blockwise_multicut_exported(): + assert hasattr(bp.segmentation, "blockwise_multicut") + assert hasattr(bp.segmentation, "blockwise_lifted_multicut") diff --git a/tests/test_lifted_multicut.py b/tests/test_lifted_multicut.py new file mode 100644 index 0000000..9f18899 --- /dev/null +++ b/tests/test_lifted_multicut.py @@ -0,0 +1,112 @@ +"""Tests for the whole-graph lifted multicut solvers and lifted-problem construction.""" +import numpy as np +import pytest + +import bioimage_cpp as bic +import bioimage_py as bp +from bioimage_py.segmentation import (get_lifted_multicut_solver, lifted_edges_from_node_labels, + lifted_multicut_gaec, lifted_multicut_fusion_moves, + lifted_multicut_kernighan_lin, lifted_problem_from_node_labels, + multicut_kernighan_lin) + + +def _line_graph(n=4): + """A simple path graph 0-1-...-(n-1) as a ``bic.graph.UndirectedGraph``.""" + uv = np.array([[i, i + 1] for i in range(n - 1)], dtype="uint64") + return bic.graph.UndirectedGraph.from_edges(n, uv), uv + + +def test_lifted_edge_splits_attractive_blob(): + # Base edges are attractive -> plain multicut merges the whole path into one segment. A strong + # repulsive lifted edge between the two ends (the canonical lifted use) must split them apart. + graph, _ = _line_graph(4) + costs = np.full(graph.number_of_edges, 5.0) + lifted_uv = np.array([[0, 3]], dtype="uint64") + lifted_costs = np.array([-50.0]) + + plain = np.asarray(multicut_kernighan_lin(graph, costs)) + assert len(np.unique(plain)) == 1 # all merged + + for solver in (lifted_multicut_kernighan_lin, lifted_multicut_gaec, lifted_multicut_fusion_moves): + node_labels = np.asarray(solver(graph, costs, lifted_uv, lifted_costs)) + assert len(node_labels) == graph.number_of_nodes + assert node_labels[0] != node_labels[3], solver.__name__ + + +def test_lifted_multicut_matches_plain_without_lifted_edges(): + # With no lifted edges the lifted solver must reproduce the plain multicut partition. + graph, _ = _line_graph(5) + costs = np.array([5.0, -5.0, 5.0, -5.0]) + empty_uv = np.zeros((0, 2), dtype="uint64") + empty_costs = np.zeros(0) + + plain = np.asarray(multicut_kernighan_lin(graph, costs)) + lifted = np.asarray(lifted_multicut_kernighan_lin(graph, costs, empty_uv, empty_costs)) + # Compare as partitions (ids may differ). + assert np.array_equal(plain[:, None] == plain[None, :], lifted[:, None] == lifted[None, :]) + + +def test_get_lifted_multicut_solver_dispatch(): + graph, _ = _line_graph(4) + costs = np.full(graph.number_of_edges, 5.0) + lifted_uv = np.array([[0, 3]], dtype="uint64") + lifted_costs = np.array([-50.0]) + + solver = get_lifted_multicut_solver("kernighan-lin", warmstart=False) + node_labels = np.asarray(solver(graph, costs, lifted_uv, lifted_costs)) + assert node_labels[0] != node_labels[3] + + with pytest.raises(ValueError): + get_lifted_multicut_solver("not-a-solver") + + +def test_lifted_edges_from_node_labels(): + # Path 0-1-2-3-4; labels split ends (1 vs 2), middle ignored (0). + graph, _ = _line_graph(5) + node_labels = np.array([1, 0, 0, 0, 2], dtype="uint64") + + diff = np.asarray(lifted_edges_from_node_labels(graph, node_labels, graph_depth=4, + mode="different", ignore_label=0)) + # Only nodes 0 and 4 carry a (different) label -> exactly the (0, 4) lifted edge. + assert diff.shape == (1, 2) + assert set(diff[0].tolist()) == {0, 4} + + # ignore_label=0 excludes the labelled-0 nodes; "same" with no shared non-zero labels -> empty. + same = np.asarray(lifted_edges_from_node_labels(graph, node_labels, graph_depth=4, + mode="same", ignore_label=0)) + assert same.shape[0] == 0 + + +def test_lifted_problem_from_node_labels_costs(): + graph, _ = _line_graph(5) + node_labels = np.array([1, 1, 2, 2, 2], dtype="uint64") + lifted_uv, lifted_costs = lifted_problem_from_node_labels( + graph, node_labels, graph_depth=4, same_segment_cost=7.0, different_segment_cost=-3.0, + mode="all", ignore_label=None, + ) + lifted_uv = np.asarray(lifted_uv) + lifted_costs = np.asarray(lifted_costs) + assert len(lifted_uv) == len(lifted_costs) > 0 + # Every cost is the same/different constant matching its endpoints' labels. + for (u, v), c in zip(lifted_uv, lifted_costs): + expected = 7.0 if node_labels[u] == node_labels[v] else -3.0 + assert c == expected + + +def test_lifted_problem_from_node_labels_empty(): + graph, _ = _line_graph(4) + node_labels = np.zeros(4, dtype="uint64") # all ignored + lifted_uv, lifted_costs = lifted_problem_from_node_labels( + graph, node_labels, graph_depth=3, same_segment_cost=1.0, different_segment_cost=-1.0, + mode="different", ignore_label=0, + ) + assert np.asarray(lifted_uv).shape == (0, 2) + assert np.asarray(lifted_costs).shape == (0,) + + +def test_lifted_multicut_exported(): + for name in ("lifted_multicut_kernighan_lin", "lifted_multicut_gaec", + "lifted_multicut_fusion_moves", "get_lifted_multicut_solver", + "lifted_edges_from_node_labels", "lifted_problem_from_node_labels", + "blockwise_lifted_multicut"): + assert hasattr(bp.segmentation, name)