diff --git a/.gitignore b/.gitignore index d495756..804c095 100644 --- a/.gitignore +++ b/.gitignore @@ -27,5 +27,6 @@ venv/ # Data *.h5 +*.npz CLAUDE.md diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 904f637..4b09083 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -486,6 +486,209 @@ Notes: `ProposalGenerator` and provide your own `_build` returning a C++ proposal-generator object if you need to extend the set. +## Lifted Multicut + +Nifty exposes lifted multicut through a separate objective + solver hierarchy. +`bioimage-cpp` mirrors the structure with `LiftedMulticutObjective` and a +`LiftedMulticutSolver` class hierarchy. + +Nifty: + +```python +import nifty.graph.opt.lifted_multicut as nlmc + +objective = nlmc.liftedMulticutObjective(graph) +objective.insertLiftedEdgesBfs(max_distance=3) +for u, v, w in lifted_weights: + objective.setCost(u, v, w) +solver = objective.liftedMulticutGreedyAdditiveFactory().create(objective) +labels = solver.optimize() +``` + +bioimage-cpp: + +```python +import bioimage_cpp as bic + +objective = bic.graph.LiftedMulticutObjective( + graph, + edge_costs, + lifted_uvs=lifted_uvs, + lifted_costs=lifted_costs, + bfs_distance=3, # optional: also insert zero-weight lifted edges within k hops +) +labels = bic.graph.LiftedGreedyAdditiveMulticut().optimize(objective) +energy = objective.energy(labels) +``` + +`LiftedMulticutObjective` accepts: + +- `graph` — an `UndirectedGraph` or `RegionAdjacencyGraph`. The constructor + copies the topology, so further mutations on the input graph do not affect + the objective. +- `edge_costs` — 1D `float64` array of length `graph.number_of_edges`. +- `lifted_uvs` / `lifted_costs` — optional `(n_lifted, 2)` uint64 array and 1D + float64 array of equal length, listing the additional lifted edges and + their weights. +- `bfs_distance` — optional positive integer. Adds a zero-weight lifted edge + for every pair of nodes within this many base-graph hops of each other + (excluding nodes already connected by a base edge). Pairs with both + `lifted_uvs` and `bfs_distance` to seed the topology and then update + specific weights. +- `overwrite_existing` — when `True`, lifted entries that coincide with an + existing edge replace its weight; the default accumulates. + +Available solvers (no fusion-move / ILP solvers yet): + +| nifty factory | bioimage-cpp solver | +| --- | --- | +| `liftedMulticutGreedyAdditiveFactory()` | `LiftedGreedyAdditiveMulticut()` | +| `liftedMulticutKernighanLinFactory(...)` | `LiftedKernighanLinMulticut(...)` | +| `chainedSolversFactory([...])` | `LiftedChainedSolvers([...])` | + +A typical warm-started solve combines greedy and KL: + +```python +solver = bic.graph.LiftedChainedSolvers([ + bic.graph.LiftedGreedyAdditiveMulticut(), + bic.graph.LiftedKernighanLinMulticut(number_of_outer_iterations=10), +]) +labels = solver.optimize(objective) +``` + +Notes: + +- Output labels are dense `uint64` ids in `0 .. number_of_clusters - 1`. +- Every output cluster is *base-graph connected* — both solvers enforce this + invariant. A strongly attractive lifted edge between two nodes that have no + base-graph path between them will not merge their clusters. +- `LiftedKernighanLinMulticut` warm-starts from the lifted greedy-additive + solution when the objective's current labels are the trivial singleton + labeling. +- `objective.set_cost(u, v, weight, overwrite=False)` updates or inserts a + single lifted edge. +- The lifted graph is exposed via `objective.lifted_graph`; the first + `objective.number_of_base_edges` edges are exactly the base edges in the + same order as in `graph`. + +### Building a lifted multicut problem from affinities + +For the common case of lifted multicut on a watershed over-segmentation, +nifty offers `nifty.graph.rag.computeLiftedEdgesFromRagAndOffsets` (lifted +edge discovery) and per-channel affinity accumulators. bioimage-cpp exposes +two focused helpers that cover the same workflow: + +```python +# Discover lifted edges implied by long-range affinity offsets. 1-hop offsets +# are skipped automatically, so the full offset list can be passed in. +lifted_uvs = bic.graph.lifted_edges_from_affinities( + rag, oversegmentation, offsets, number_of_threads=0, +) + +# Accumulate (mean, size) statistics per lifted edge. Pixel pairs whose +# (u, v) does not appear in `lifted_uvs` are skipped, so local edges are +# never contaminated with long-range affinities. +lifted_features = bic.graph.lifted_affinity_features( + oversegmentation, affinities, offsets, lifted_uvs, + number_of_threads=0, +) +# For the 12-column feature set (mean, median, std, min, max, percentiles, size): +lifted_features = bic.graph.lifted_affinity_features_complex(...) +``` + +The output column conventions match the local-edge variants +(`SIMPLE_EDGE_FEATURE_NAMES`, `COMPLEX_EDGE_FEATURE_NAMES`). + +End-to-end pipeline (also in `examples/segmentation/lifted_multicut_from_affinities.py`): + +```python +rag = bic.graph.region_adjacency_graph(oversegmentation) +local_costs = local_threshold - bic.graph.affinity_features( + rag, oversegmentation, direct_affinities, direct_offsets, +)[:, 0] +lifted_uvs = bic.graph.lifted_edges_from_affinities( + rag, oversegmentation, long_range_offsets, +) +lifted_costs = lifted_threshold - bic.graph.lifted_affinity_features( + oversegmentation, long_range_affinities, long_range_offsets, lifted_uvs, +)[:, 0] +objective = bic.graph.LiftedMulticutObjective( + rag, local_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs, +) +``` + +## External Problem Instances + +bioimage-cpp ships pooch-backed downloaders for the multicut and lifted +multicut benchmark problems used by the development comparison scripts and +the regression tests. Files are cached under `~/.cache/bioimage-cpp/`, +overridable via the `BIOIMAGE_CPP_CACHE` environment variable. + +`pooch` is an optional runtime dependency — install via the `test` or `data` +extras, e.g. `pip install bioimage-cpp[data]`. + +Multicut problems (3 samples × 2 sizes, originally from +`elf.segmentation.utils.load_multicut_problem`): + +```python +# Returns (UndirectedGraph, edge_costs) +graph, costs = bic.graph.load_multicut_problem(sample="A", size="small") +# Or just the underlying arrays +uv_ids, costs = bic.graph.load_multicut_problem_data(sample="B", size="medium") +# Or the cached file path +path = bic.graph.multicut_problem_path(sample="C", size="medium") +``` + +Valid samples are `"A"`, `"B"`, `"C"`; valid sizes are `"small"` and +`"medium"`. The legacy `load_external_multicut_problem` / +`load_external_multicut_problem_data` / `external_multicut_problem_path` +shims default to sample A, size small and continue to honor the +`BIOIMAGE_CPP_EXTERNAL_MULTICUT_PATH` and +`BIOIMAGE_CPP_EXTERNAL_MULTICUT_CACHE` environment variables. + +Lifted multicut problems (2D ISBI slice and full 3D volume, built by +`examples/segmentation/serialize_lifted_problem.py`): + +```python +problem = bic.graph.load_lifted_multicut_problem(size="2d") +# Fields: n_nodes (int), local_uvs, local_costs, lifted_uvs, lifted_costs. +graph = bic.graph.UndirectedGraph.from_edges(problem.n_nodes, problem.local_uvs) +objective = bic.graph.LiftedMulticutObjective( + graph, + problem.local_costs, + lifted_uvs=problem.lifted_uvs, + lifted_costs=problem.lifted_costs, +) +``` + +Notes: + +- Every download is integrity-checked against a SHA256 in the registry; a + corrupted cache file is detected on the next `load_*` call. +- Downloads are lazy: nothing happens until you call a loader. Re-runs are + free (the cached file is reused). +- For air-gapped use, fetch the file once on a machine with network access + and copy `~/.cache/bioimage-cpp/` to the same path on the target + machine. + +## Breadth-First Search + +Nifty has an internal `BreadthFirstSearch` template used during lifted-edge +insertion. `bioimage-cpp` exposes a Python-friendly free function: + +```python +nodes, distances = bic.graph.breadth_first_search( + graph, + source, + max_distance=3, # optional, default: full component + include_source=True, # set to False for k-hop neighborhoods excluding self +) +``` + +Both output arrays are 1D `uint64`, listing reached nodes in BFS order with +their hop distance from the source. Useful for building lifted-edge sets +manually, sampling local neighborhoods, or computing graph distances. + ## Segmentation Overlaps Nifty: diff --git a/development/graph/lifted_multicut/PERFORMANCE_NOTES.md b/development/graph/lifted_multicut/PERFORMANCE_NOTES.md new file mode 100644 index 0000000..dbacc82 --- /dev/null +++ b/development/graph/lifted_multicut/PERFORMANCE_NOTES.md @@ -0,0 +1,197 @@ +# Lifted Multicut Performance Notes + +State of the lifted-multicut solvers vs nifty on the standard benchmark +problems and notes on remaining optimization headroom. Read this before the +next round of perf work. + +## Current benchmark (3D ISBI lifted problem) + +Problem dimensions: 2462 nodes, 17 949 local edges, 21 444 lifted edges. + +| Solver | bic | nifty | Ratio | Status | +|---|---|---|---|---| +| Greedy additive | ~13 ms | ~14.7 ms | **0.88×** (faster) | Goal met | +| Kernighan-Lin (greedy + 10 outer) | ~200 ms | ~102 ms | **1.95×** (slower) | Outside 30%-of-nifty target | + +Energies match nifty to within numerical noise on both solvers (greedy diff +~0.3, KL diff ~0.08). + +## What's done + +### Greedy additive + +Three Python-side wins (~3× speedup, 39 → 13 ms): + +1. **Bulk `_add_lifted_edges` fast path** in `src/bioimage_cpp/graph/__init__.py`. + Replaced a per-row Python loop over `lifted_uvs` with one + `insert_edges` call plus `np.bincount` for the residual + collision case. +2. **`UndirectedGraph.from_unique_edges` binding** in `src/bindings/graph.cxx`. + Bypasses the per-edge hash dedup that `insert_edge` performs; used + from `_copy_graph` and from the lifted-graph construction path. +3. **Dropped defensive base-graph copy** in `LiftedMulticutObjective`. + The C++ `Objective` already only holds a `const UndirectedGraph &`; the + Python wrapper now matches. + +The C++ greedy kernel itself runs in ~4 ms — already at the algorithmic +floor for ~21 k heap operations. + +### Kernighan-Lin + +One C++ optimization (~22% speedup, 245 → 200 ms): + +1. **Pre-built per-node filtered adjacency** in + `include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx`. + `ChainScratch` gained `filtered_offset`, `filtered_count`, and + `filtered_entries` (each entry caches `{node, weight, is_base}`). + `chain_gain_init` populates these alongside the gain accumulation; the + chain loop iterates only in-pair neighbors and skips the + `bufs.in_pair[u_key]` filter check. Also caches `was_in_heap` once per + iteration so the three subsequent heap operations don't each re-query + the locator. + +## What remains — KL is the open item + +C++ KL takes ~200 ms; nifty ~102 ms. Profile breakdown +(`BIOIMAGE_PROFILE=ON`, 1 repeat) — note the inner scopes inflate total by +~80% so trust the relative shares: + +``` + cc_repartition 0.0033 s ( 0.7%) + energy_eval 0.0013 s ( 0.4%) + compute_pairs 0.0073 s ( 2.0%) + chain_init 0.0037 s ( 1.0%) + chain_gain_init 0.0812 s ( 22.1%) <-- ~36% of pair_chains + chain_loop 0.0880 s ( 23.9%) <-- ~63% of pair_chains + chain_cleanup 0.0017 s ( 0.5%) + pair_chains 0.1744 s ( 47.4%) + cluster_splits 0.0083 s ( 2.3%) +``` + +Workload statistics on the 3D problem after the greedy warm-start: + +- 643 clusters. Sizes: min 1, max 139, mean 3.8, median 2. +- 4443 base cluster pairs per outer iter. Pair sizes: mean 20, median 7, + max 198 (long tail). +- Average lifted node degree: 32. + +## Ranked future optimizations + +### 1. Pre-bucketed gain init (estimated landing point: ~150 ms total, −25%) + +**Idea.** For each outer iteration, bucket every lifted edge by its +endpoint cluster pair (sorted `(min_label, max_label)`). For pair-chain +`(A, B)` the gain init iterates only the three relevant buckets — +`(A, A)`, `(B, B)`, `(A, B)` — instead of walking the full lifted +adjacency of every node in the pair. + +**Why it would help.** Each edge currently contributes to `chain_gain_init` +exactly once per pair-chain that touches one of its endpoint clusters. +With buckets, each edge contributes O(1) per outer iter. The 80 ms +`chain_gain_init` collapses to roughly O(E) = ~5 ms per outer iter, saving +~70 ms. + +**The wrinkle.** Bucket membership goes stale as nodes move between +clusters during the sequential pair-chains. The right fix is **incremental +bucket maintenance** — every node move performs O(degree) bucket-membership +updates (remove from old bucket, push into new). Bucket entries store +their position via an `edge_id → bucket_pos` index so removal is +`swap-with-back` in O(1). Total maintenance cost per outer iter: +`O(num_moves × avg_degree)` ≈ 30 µs on our problem. + +**Complexity.** Substantial: ~200 lines, new `LiftedEdgeBuckets` struct in +`detail_kl`, hooks in `chain_loop` to call `buckets.relabel(v, old, new)` +on every committed move, careful invariant management. + +**Sanity check before implementing.** Try a quick prototype where buckets +are rebuilt fresh at the start of each outer iter (O(E) per outer iter) +and gain init uses buckets, accepting that within-outer-iter moves create +stale entries. If the resulting energy is comparable to the exact version, +the incremental maintenance is worthwhile. + +### 2. CSR adjacency layout for lifted graph (estimated: 10–15% off) + +**Idea.** `UndirectedGraph` stores adjacency as `vector>` +— one heap allocation per node. Walking adjacency for many distinct nodes +in a pair pays per-node pointer chasing. A flat CSR (`offsets[n+1]` + +`entries[2E]`) built once at the start of `kernighan_lin` would be more +cache-friendly. + +**Caveat.** The actual access pattern in `chain_gain_init` walks +adjacency for nodes in arbitrary order (whichever pair we're processing), +so spatial locality across nodes is poor regardless of layout. The win is +limited to per-node cache-line savings (one miss per node vs one per +adjacency vector header). Estimate ~10% based on rough cycle counting. + +**Complexity.** Localized: ~50 lines, build CSR in `kernighan_lin`, +replace `lifted_graph.node_adjacency(v)` calls in `chain_gain_init` with +CSR iteration. + +**Worth combining** with optimization 1, since CSR walking is what bucket +maintenance would need anyway. + +### 3. Inspect nifty's internals (estimated: unknown, possibly clarifying) + +The gap between our `chain_gain_init` and nifty's equivalent +`computeDifferences` is suspiciously ~2× per adjacency entry given that +both algorithms walk the same data. nifty's source is at: + +``` +/home/pape/Work/software/src/nifty/include/nifty/graph/opt/lifted_multicut/detail/lifted_twocut_kernighan_lin.hxx +``` + +Worth checking: +- How nifty stores `liftedGraph_.adjacency(v)` — is it CSR-like? +- Whether nifty's `referencedBy` array is `uint32` or larger (we use + `uint32`). +- Whether nifty's `differences` (= our `stash_gain`) cache line layout + differs. + +### 4. Linear-scan border instead of a heap + +**Verdict: not worth pursuing for this problem.** + +I worked through it: heap is faster than linear scan for our pair-size +distribution. The heap pop is O(log N) and heap.change is also O(log N); +for pair size 7 (the median) that's ~2× faster than O(N) linear scan, +and the gap widens for larger pairs. + +nifty uses linear scan, but that's not where its speed advantage comes +from — likely it's the adjacency-walking constants (optimization 3). + +### 5. Skip pair-chains where heap stays empty (estimated: <5 ms) + +After `chain_gain_init`, if `heap.empty()` we already skip +`chain_loop`. But we still pay for `chain_init` and the full +`chain_gain_init` walk. For pair-chains where the cluster pair has only +one alive node per side (post-staleness filtering of `cluster_to_nodes`), +we could skip earlier. Need to maintain live cluster sizes. + +Low priority — only a few ms. + +## Where to start next time + +1. Re-run `cd development/graph/lifted_multicut && python check_kernighan_lin.py --size 3d --repeats 5` to confirm the baseline hasn't drifted. +2. Build with `pip install -e . --no-build-isolation -C cmake.define.BIOIMAGE_PROFILE=ON` to get the profile breakdown back. +3. Prototype the rebuild-per-outer-iter bucket approach (optimization 1 + simplified) to validate the energy quality before committing to the + incremental maintenance version. +4. Targets: + - 30%-of-nifty: ≤132 ms. + - Match nifty: ≤102 ms. + +## Files that matter + +- `include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx` — + KL kernel; existing profile scopes wrap each phase. +- `include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx` — + greedy kernel; profile scopes already present, currently at ~4 ms + so not the focus. +- `include/bioimage_cpp/graph/lifted_multicut/objective.hxx` — objective + state, including `n_base_edges` and the lifted graph. +- `src/bioimage_cpp/graph/__init__.py::LiftedMulticutObjective` — + Python construction path (already optimized). +- `development/graph/lifted_multicut/_compatibility.py` — + bic-vs-nifty harness; uses `run_comparison(...)`. +- `tests/graph/lifted_multicut/test_external_problem.py` — regression + test on the 2D problem; energy bound is ENERGY_BOUND = -1574.5. diff --git a/development/graph/lifted_multicut/_compatibility.py b/development/graph/lifted_multicut/_compatibility.py new file mode 100644 index 0000000..9247210 --- /dev/null +++ b/development/graph/lifted_multicut/_compatibility.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import argparse +from statistics import median +from time import perf_counter +from typing import Callable + +import numpy as np + + +def parser(description: str) -> argparse.ArgumentParser: + arg_parser = argparse.ArgumentParser(description=description) + arg_parser.add_argument( + "--size", + choices=("2d", "3d"), + default="3d", + help="Lifted multicut problem instance to load (default: 3d).", + ) + arg_parser.add_argument( + "--repeats", + type=int, + default=3, + help="Number of timed repeats per implementation.", + ) + arg_parser.add_argument( + "--timeout", + type=float, + default=60.0, + help="Download timeout in seconds if the lifted problem is not cached.", + ) + arg_parser.add_argument( + "--energy-bound", + type=float, + default=None, + help=( + "Optional maximum accepted energy for both implementations. If " + "omitted, energies are reported but not asserted." + ), + ) + return arg_parser + + +def load_problem(size: str, *, timeout: float): + import bioimage_cpp as bic + import nifty + import nifty.graph.opt.lifted_multicut as nlmc + + problem = bic.graph.load_lifted_multicut_problem(size, timeout=timeout) + + bic_graph = bic.graph.UndirectedGraph.from_edges( + problem.n_nodes, problem.local_uvs + ) + + nifty_graph = nifty.graph.undirectedGraph(int(problem.n_nodes)) + nifty_graph.insertEdges(problem.local_uvs.astype(np.uint64, copy=False)) + nifty_objective = nlmc.liftedMulticutObjective(nifty_graph) + nifty_objective.setGraphEdgesCosts(problem.local_costs) + if problem.lifted_uvs.shape[0] > 0: + nifty_objective.setCosts( + problem.lifted_uvs.astype(np.uint64, copy=False), + problem.lifted_costs.astype(np.float64, copy=False), + ) + + return bic_graph, nifty_objective, problem + + +def time_call(function: Callable[[], np.ndarray], repeats: int): + timings = [] + result = None + for _ in range(repeats): + start = perf_counter() + result = function() + timings.append(perf_counter() - start) + assert result is not None + return timings, result + + +def optimize_bic_solver(make_bic_solver, bic_graph, problem): + import bioimage_cpp as bic + + objective = bic.graph.LiftedMulticutObjective( + bic_graph, + problem.local_costs, + lifted_uvs=problem.lifted_uvs, + lifted_costs=problem.lifted_costs, + ) + return make_bic_solver().optimize(objective) + + +def bic_energy(bic_graph, problem, labels: np.ndarray) -> float: + import bioimage_cpp as bic + + objective = bic.graph.LiftedMulticutObjective( + bic_graph, + problem.local_costs, + lifted_uvs=problem.lifted_uvs, + lifted_costs=problem.lifted_costs, + ) + return float(objective.energy(labels)) + + +def nifty_energy(nifty_objective, labels: np.ndarray) -> float: + return float(nifty_objective.evalNodeLabels(labels.astype(np.uint64, copy=False))) + + +def run_comparison( + name: str, + make_bic_solver, + make_nifty_solver, + args: argparse.Namespace, +) -> dict[str, float]: + bic_graph, nifty_objective, problem = load_problem(args.size, timeout=args.timeout) + + bic_timings, bic_labels = time_call( + lambda: optimize_bic_solver(make_bic_solver, bic_graph, problem), + args.repeats, + ) + nifty_timings, nifty_labels = time_call( + lambda: make_nifty_solver(nifty_objective).create(nifty_objective).optimize(), + args.repeats, + ) + + bic_score = bic_energy(bic_graph, problem, bic_labels) + nifty_score = nifty_energy(nifty_objective, np.asarray(nifty_labels)) + + if args.energy_bound is not None: + if bic_score > args.energy_bound: + raise AssertionError( + f"bioimage-cpp {name} energy {bic_score:.6f} exceeds bound " + f"{args.energy_bound:.6f}" + ) + if nifty_score > args.energy_bound: + raise AssertionError( + f"nifty {name} energy {nifty_score:.6f} exceeds bound " + f"{args.energy_bound:.6f}" + ) + + result = { + "bioimage_cpp_energy": bic_score, + "nifty_energy": nifty_score, + "energy_difference": bic_score - nifty_score, + "bioimage_cpp_median_runtime": median(bic_timings), + "nifty_median_runtime": median(nifty_timings), + } + print(f"solver: {name}") + print( + f"problem: size={args.size}, nodes={problem.n_nodes}, " + f"local edges={problem.local_uvs.shape[0]}, " + f"lifted edges={problem.lifted_uvs.shape[0]}" + ) + print(f"bioimage-cpp energy: {bic_score:.6f}") + print(f"nifty energy: {nifty_score:.6f}") + print(f"energy difference: {bic_score - nifty_score:.6f}") + print(f"bioimage-cpp median runtime [s]: {median(bic_timings):.6f}") + print(f"nifty median runtime [s]: {median(nifty_timings):.6f}") + return result diff --git a/development/graph/lifted_multicut/check_greedy_additive.py b/development/graph/lifted_multicut/check_greedy_additive.py new file mode 100644 index 0000000..c789c85 --- /dev/null +++ b/development/graph/lifted_multicut/check_greedy_additive.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import bioimage_cpp as bic + +from _compatibility import parser, run_comparison + + +def main() -> None: + args = parser( + "Compare bioimage-cpp and nifty greedy-additive lifted multicut." + ).parse_args() + run_comparison( + "lifted_greedy_additive", + lambda: bic.graph.LiftedGreedyAdditiveMulticut(), + lambda objective: objective.liftedMulticutGreedyAdditiveFactory(), + args, + ) + + +if __name__ == "__main__": + main() diff --git a/development/graph/lifted_multicut/check_kernighan_lin.py b/development/graph/lifted_multicut/check_kernighan_lin.py new file mode 100644 index 0000000..840b923 --- /dev/null +++ b/development/graph/lifted_multicut/check_kernighan_lin.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import bioimage_cpp as bic + +from _compatibility import parser, run_comparison + + +def main() -> None: + args = parser( + "Compare bioimage-cpp and nifty Kernighan-Lin lifted multicut." + ).parse_args() + # On the bioimage-cpp side, LiftedKernighanLinMulticut auto-warm-starts + # from a greedy-additive pass when the objective's current labels are the + # trivial singleton labeling. On the nifty side we make the equivalent + # warm-start explicit via the chained-solvers factory so the two + # implementations are doing the same work. + run_comparison( + "lifted_kernighan_lin", + lambda: bic.graph.LiftedKernighanLinMulticut(number_of_outer_iterations=10), + lambda objective: objective.chainedSolversFactory( + [ + objective.liftedMulticutGreedyAdditiveFactory(), + objective.liftedMulticutKernighanLinFactory( + numberOfOuterIterations=10 + ), + ] + ), + args, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/segmentation/_lifted_problem.py b/examples/segmentation/_lifted_problem.py new file mode 100644 index 0000000..b862012 --- /dev/null +++ b/examples/segmentation/_lifted_problem.py @@ -0,0 +1,181 @@ +"""Shared lifted-multicut-from-affinities pipeline. + +Used by `lifted_multicut_from_affinities.py` (visualization) and +`serialize_lifted_problem.py` (writes the problem to disk for the development +comparison scripts). Both scripts share the heightmap + watershed + +RAG + lifted-edge construction so they stay in lockstep. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from elf.segmentation.utils import load_mutex_watershed_problem +from skimage.feature import peak_local_max +from skimage.measure import label as label_components +from skimage.segmentation import watershed + +import bioimage_cpp as bic + + +@dataclass +class AffinityProblem: + """Affinity volume split into direct / long-range channels.""" + + full_affinities: np.ndarray + full_offsets: list[tuple[int, ...]] + direct_affinities: np.ndarray + direct_offsets: list[tuple[int, ...]] + long_range_affinities: np.ndarray + long_range_offsets: list[tuple[int, ...]] + + +@dataclass +class LiftedProblem: + """Built lifted-multicut problem ready for solvers or serialization.""" + + rag: bic.graph.RegionAdjacencyGraph + oversegmentation: np.ndarray + heightmap: np.ndarray + local_costs: np.ndarray + lifted_uvs: np.ndarray + lifted_costs: np.ndarray + + @property + def number_of_nodes(self) -> int: + return int(self.rag.number_of_nodes) + + @property + def local_uvs(self) -> np.ndarray: + return self.rag.uv_ids() + + +def load_affinity_problem( + data_prefix: Path, + ndim: int, + z_slice: int, +) -> AffinityProblem: + affinities, offsets = load_mutex_watershed_problem(prefix=str(data_prefix)) + offsets = [tuple(int(v) for v in offset) for offset in offsets] + if ndim == 2: + channels_2d = [index for index, offset in enumerate(offsets) if offset[0] == 0] + affinities = affinities[channels_2d, z_slice] + offsets = [offsets[index][1:] for index in channels_2d] + elif ndim != 3: + raise ValueError(f"ndim must be 2 or 3, got {ndim}") + + direct_channels = [ + index for index, offset in enumerate(offsets) + if sum(abs(v) for v in offset) == 1 + ] + long_range_channels = [ + index for index in range(len(offsets)) if index not in direct_channels + ] + return AffinityProblem( + full_affinities=np.ascontiguousarray(affinities, dtype=np.float32), + full_offsets=offsets, + direct_affinities=np.ascontiguousarray( + affinities[direct_channels], dtype=np.float32 + ), + direct_offsets=[offsets[index] for index in direct_channels], + long_range_affinities=np.ascontiguousarray( + affinities[long_range_channels], dtype=np.float32 + ), + long_range_offsets=[offsets[index] for index in long_range_channels], + ) + + +def make_heightmap(affinities: np.ndarray) -> np.ndarray: + return np.ascontiguousarray(np.mean(affinities, axis=0), dtype=np.float32) + + +def make_watershed_oversegmentation( + heightmap: np.ndarray, + *, + min_distance: int, + grid_spacing: int, + max_markers: int, +) -> np.ndarray: + coordinates = peak_local_max( + -heightmap, + min_distance=min_distance, + exclude_border=False, + num_peaks=max_markers, + ) + marker_mask = np.zeros(heightmap.shape, dtype=bool) + if len(coordinates) > 0: + marker_mask[tuple(coordinates.T)] = True + markers = label_components(marker_mask).astype(np.int32, copy=False) + + if int(markers.max()) < 2: + markers = np.zeros(heightmap.shape, dtype=np.int32) + slices = tuple(slice(None, None, grid_spacing) for _ in heightmap.shape) + marker_coordinates = np.argwhere(np.ones(heightmap.shape, dtype=bool)[slices]) + marker_coordinates *= grid_spacing + for marker_id, coord in enumerate(marker_coordinates, start=1): + markers[tuple(coord)] = marker_id + + return watershed(heightmap, markers=markers).astype(np.uint32, copy=False) + + +def build_lifted_problem( + affinity_problem: AffinityProblem, + *, + local_threshold: float, + lifted_threshold: float, + number_of_threads: int, + watershed_min_distance: int, + watershed_grid_spacing: int, + max_markers: int, +) -> LiftedProblem: + heightmap = make_heightmap(affinity_problem.direct_affinities) + oversegmentation = make_watershed_oversegmentation( + heightmap, + min_distance=watershed_min_distance, + grid_spacing=watershed_grid_spacing, + max_markers=max_markers, + ) + + rag = bic.graph.region_adjacency_graph( + oversegmentation, number_of_threads=number_of_threads + ) + + local_features = bic.graph.affinity_features( + rag, + oversegmentation, + affinity_problem.direct_affinities, + affinity_problem.direct_offsets, + number_of_threads=number_of_threads, + ) + local_costs = (local_threshold - local_features[:, 0]).astype(np.float64, copy=False) + + lifted_uvs = bic.graph.lifted_edges_from_affinities( + rag, + oversegmentation, + affinity_problem.long_range_offsets, + number_of_threads=number_of_threads, + ) + if lifted_uvs.shape[0] == 0: + lifted_costs = np.zeros(0, dtype=np.float64) + else: + lifted_features = bic.graph.lifted_affinity_features( + oversegmentation, + affinity_problem.long_range_affinities, + affinity_problem.long_range_offsets, + lifted_uvs, + number_of_threads=number_of_threads, + ) + lifted_costs = (lifted_threshold - lifted_features[:, 0]).astype( + np.float64, copy=False + ) + + return LiftedProblem( + rag=rag, + oversegmentation=oversegmentation, + heightmap=heightmap, + local_costs=np.ascontiguousarray(local_costs), + lifted_uvs=np.ascontiguousarray(lifted_uvs), + lifted_costs=np.ascontiguousarray(lifted_costs), + ) diff --git a/examples/segmentation/lifted_multicut_from_affinities.py b/examples/segmentation/lifted_multicut_from_affinities.py new file mode 100644 index 0000000..6f5f7e2 --- /dev/null +++ b/examples/segmentation/lifted_multicut_from_affinities.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import napari +from elf.io import open_file +from skimage.segmentation import find_boundaries + +import bioimage_cpp as bic + +from _lifted_problem import build_lifted_problem, load_affinity_problem + + +THIS_DIR = Path(__file__).resolve().parent +DEFAULT_DATA_PREFIX = THIS_DIR / "isbi-data-" + + +def load_raw(data_prefix: Path, ndim: int, z_slice: int): + data_path = data_prefix.with_name(data_prefix.name + "test.h5") + with open_file(data_path, "r") as f: + raw = f["raw"][z_slice] if ndim == 2 else f["raw"][:] + return raw + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Run watershed oversegmentation + RAG lifted multicut on the ISBI " + "affinity example." + ) + ) + parser.add_argument("--ndim", type=int, choices=(2, 3), default=2) + parser.add_argument("--z-slice", type=int, default=0) + parser.add_argument("--data-prefix", type=Path, default=DEFAULT_DATA_PREFIX) + parser.add_argument("--local-threshold", type=float, default=0.1) + parser.add_argument("--lifted-threshold", type=float, default=0.1) + parser.add_argument("--threads", type=int, default=0) + parser.add_argument("--watershed-min-distance", type=int, default=5) + parser.add_argument("--watershed-grid-spacing", type=int, default=12) + parser.add_argument("--max-markers", type=int, default=2048) + parser.add_argument("--kl-outer-iterations", type=int, default=10) + args = parser.parse_args() + + affinity_problem = load_affinity_problem(args.data_prefix, args.ndim, args.z_slice) + raw = load_raw(args.data_prefix, args.ndim, args.z_slice) + + lifted_problem = build_lifted_problem( + affinity_problem, + local_threshold=args.local_threshold, + lifted_threshold=args.lifted_threshold, + number_of_threads=args.threads, + watershed_min_distance=args.watershed_min_distance, + watershed_grid_spacing=args.watershed_grid_spacing, + max_markers=args.max_markers, + ) + print( + f"Built RAG: {lifted_problem.number_of_nodes} nodes, " + f"{lifted_problem.local_costs.size} local edges, " + f"{lifted_problem.lifted_uvs.shape[0]} lifted edges" + ) + + objective = bic.graph.LiftedMulticutObjective( + lifted_problem.rag, + lifted_problem.local_costs, + lifted_uvs=lifted_problem.lifted_uvs, + lifted_costs=lifted_problem.lifted_costs, + ) + solver = bic.graph.LiftedChainedSolvers( + [ + bic.graph.LiftedGreedyAdditiveMulticut(), + bic.graph.LiftedKernighanLinMulticut( + number_of_outer_iterations=args.kl_outer_iterations + ), + ] + ) + node_labels = solver.optimize(objective) + print( + f"Lifted multicut energy: {objective.energy(node_labels):.3f}, " + f"{int(node_labels.max()) + 1} segments" + ) + + segmentation = bic.graph.project_node_labels_to_pixels( + lifted_problem.rag, + lifted_problem.oversegmentation, + node_labels, + number_of_threads=args.threads, + ) + + viewer = napari.Viewer() + viewer.add_image(raw, name="raw") + viewer.add_image(affinity_problem.direct_affinities, name="direct affinities") + viewer.add_image( + affinity_problem.long_range_affinities, name="long-range affinities" + ) + viewer.add_image(lifted_problem.heightmap, name="watershed heightmap") + viewer.add_labels( + lifted_problem.oversegmentation, name="watershed oversegmentation" + ) + viewer.add_labels(segmentation, name="lifted multicut segmentation") + viewer.add_labels(find_boundaries(segmentation), name="lifted multicut boundaries") + napari.run() + + +if __name__ == "__main__": + main() diff --git a/examples/segmentation/serialize_lifted_problem.py b/examples/segmentation/serialize_lifted_problem.py new file mode 100644 index 0000000..4f2da92 --- /dev/null +++ b/examples/segmentation/serialize_lifted_problem.py @@ -0,0 +1,99 @@ +"""Serialize the ISBI lifted-multicut problem to a .npz file. + +Runs the same build pipeline as `lifted_multicut_from_affinities.py` up to +the LiftedMulticutObjective inputs, then writes them to a single `.npz` file +with fields: + + n_nodes : scalar uint64 + local_uvs : (n_local, 2) uint64 + local_costs : (n_local,) float64 + lifted_uvs : (n_lifted, 2) uint64 + lifted_costs : (n_lifted,) float64 + +The resulting file is what the development comparison scripts in +`development/graph/lifted_multicut/` load. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np + +from _lifted_problem import build_lifted_problem, load_affinity_problem + + +THIS_DIR = Path(__file__).resolve().parent +DEFAULT_DATA_PREFIX = THIS_DIR / "isbi-data-" +DEFAULT_OUTPUT = THIS_DIR / "lifted_multicut_problem.npz" + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Build the ISBI lifted multicut problem (RAG + local + lifted edges + " + "costs) and serialize it to a .npz file." + ) + ) + parser.add_argument("--ndim", type=int, choices=(2, 3), default=2) + parser.add_argument("--z-slice", type=int, default=0) + parser.add_argument("--data-prefix", type=Path, default=DEFAULT_DATA_PREFIX) + parser.add_argument("--local-threshold", type=float, default=0.1) + parser.add_argument("--lifted-threshold", type=float, default=0.1) + parser.add_argument("--threads", type=int, default=0) + parser.add_argument("--watershed-min-distance", type=int, default=5) + parser.add_argument("--watershed-grid-spacing", type=int, default=12) + parser.add_argument("--max-markers", type=int, default=2048) + parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) + args = parser.parse_args() + + affinity_problem = load_affinity_problem(args.data_prefix, args.ndim, args.z_slice) + lifted_problem = build_lifted_problem( + affinity_problem, + local_threshold=args.local_threshold, + lifted_threshold=args.lifted_threshold, + number_of_threads=args.threads, + watershed_min_distance=args.watershed_min_distance, + watershed_grid_spacing=args.watershed_grid_spacing, + max_markers=args.max_markers, + ) + + local_uvs = np.ascontiguousarray(lifted_problem.local_uvs.astype(np.uint64, copy=False)) + local_costs = np.ascontiguousarray( + lifted_problem.local_costs.astype(np.float64, copy=False) + ) + lifted_uvs = np.ascontiguousarray(lifted_problem.lifted_uvs.astype(np.uint64, copy=False)) + lifted_costs = np.ascontiguousarray( + lifted_problem.lifted_costs.astype(np.float64, copy=False) + ) + + n_nodes = np.uint64(lifted_problem.number_of_nodes) + output = args.output + output.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + output, + n_nodes=n_nodes, + local_uvs=local_uvs, + local_costs=local_costs, + lifted_uvs=lifted_uvs, + lifted_costs=lifted_costs, + ) + + print(f"Wrote lifted multicut problem to {output}") + print(f" number of nodes: {int(n_nodes)}") + print(f" number of local edges: {local_uvs.shape[0]}") + print(f" number of lifted edges: {lifted_uvs.shape[0]}") + print( + f" local cost range: [{float(local_costs.min()):+.3f}, " + f"{float(local_costs.max()):+.3f}]" + ) + if lifted_costs.size: + print( + f" lifted cost range: [{float(lifted_costs.min()):+.3f}, " + f"{float(lifted_costs.max()):+.3f}]" + ) + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/graph/breadth_first_search.hxx b/include/bioimage_cpp/graph/breadth_first_search.hxx new file mode 100644 index 0000000..0f140d3 --- /dev/null +++ b/include/bioimage_cpp/graph/breadth_first_search.hxx @@ -0,0 +1,118 @@ +#pragma once + +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +// One (node, distance) entry from a breadth-first search. +struct BfsEntry { + std::uint64_t node; + std::uint64_t distance; +}; + +// Reusable scratch state for `breadth_first_search`. Reset via `reset(graph)` +// before each call so the visited-flag and distance buffers grow once and stay +// allocated across many BFS invocations on the same graph. +class BfsWorkspace { +public: + BfsWorkspace() = default; + + void reset(const UndirectedGraph &graph) { + const auto n_nodes = static_cast(graph.number_of_nodes()); + visited_.assign(n_nodes, 0); + distance_.assign(n_nodes, 0); + } + + [[nodiscard]] std::vector &visited() { return visited_; } + [[nodiscard]] std::vector &distance() { return distance_; } + +private: + std::vector visited_; + std::vector distance_; +}; + +// Distance value reported for the source node itself (distance == 0). +inline constexpr std::uint64_t bfs_source_distance = 0; + +// Sentinel for "no maximum distance" — the BFS expands until the entire +// connected component of `source` has been reported. +inline constexpr std::uint64_t bfs_no_max_distance = + std::numeric_limits::max(); + +// Run a breadth-first search on `graph` starting from `source`, reporting every +// node reached within `max_distance` hops (inclusive) in BFS order. +// +// The `source` node itself is reported with distance 0. When +// `include_source` is false the source is excluded from the output (useful for +// "nodes within k hops, excluding self" queries, e.g. lifted-edge insertion). +// When `max_distance == bfs_no_max_distance` the search expands until the +// entire connected component is visited. +// +// `workspace` lets the caller reuse internal buffers across calls on the same +// graph; pass a fresh workspace for a one-off call. +inline std::vector breadth_first_search( + const UndirectedGraph &graph, + const std::uint64_t source, + const std::uint64_t max_distance, + const bool include_source, + BfsWorkspace &workspace +) { + if (source >= graph.number_of_nodes()) { + throw std::invalid_argument( + "source must be < number_of_nodes" + ); + } + workspace.reset(graph); + auto &visited = workspace.visited(); + auto &distance = workspace.distance(); + + std::vector result; + std::queue queue; + queue.push(source); + visited[static_cast(source)] = 1; + distance[static_cast(source)] = 0; + if (include_source) { + result.push_back({source, 0}); + } + + while (!queue.empty()) { + const auto node = queue.front(); + queue.pop(); + const auto node_distance = distance[static_cast(node)]; + if (node_distance >= max_distance) { + continue; + } + const auto next_distance = node_distance + 1; + for (const auto adjacency : graph.node_adjacency(node)) { + const auto neighbor = adjacency.node; + if (visited[static_cast(neighbor)] != 0) { + continue; + } + visited[static_cast(neighbor)] = 1; + distance[static_cast(neighbor)] = next_distance; + result.push_back({neighbor, next_distance}); + queue.push(neighbor); + } + } + return result; +} + +inline std::vector breadth_first_search( + const UndirectedGraph &graph, + const std::uint64_t source, + const std::uint64_t max_distance = bfs_no_max_distance, + const bool include_source = true +) { + BfsWorkspace workspace; + return breadth_first_search(graph, source, max_distance, include_source, workspace); +} + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/lifted_from_affinities.hxx b/include/bioimage_cpp/graph/lifted_from_affinities.hxx new file mode 100644 index 0000000..4fed26c --- /dev/null +++ b/include/bioimage_cpp/graph/lifted_from_affinities.hxx @@ -0,0 +1,338 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/edge_hash.hxx" +#include "bioimage_cpp/detail/grid.hxx" +#include "bioimage_cpp/detail/threading.hxx" +#include "bioimage_cpp/graph/feature_accumulation.hxx" +#include "bioimage_cpp/graph/region_adjacency_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +namespace detail_lifted { + +// An offset is "long-range" if it moves by more than a single grid step along +// any axis or by a step on more than one axis. 1-hop offsets (sum |o| == 1) +// already correspond to local RAG edges; skipping them keeps lifted-edge +// discovery focused on edges that the local RAG cannot represent. +inline bool is_long_range(const std::vector &offset) { + std::ptrdiff_t l1 = 0; + for (const auto value : offset) { + l1 += value < 0 ? -value : value; + } + return l1 > 1; +} + +inline void validate_affinity_inputs( + const ConstArrayView &affinities_or_dummy, + const std::vector &labels_shape, + const std::vector> &offsets, + const bool has_affinities +) { + if (has_affinities) { + if (affinities_or_dummy.ndim() != static_cast(labels_shape.size()) + 1) { + throw std::invalid_argument( + "affinities must have shape (channels, *labels.shape)" + ); + } + if (static_cast(affinities_or_dummy.shape[0]) != offsets.size()) { + throw std::invalid_argument( + "offsets length must match affinities channel count" + ); + } + for (std::size_t axis = 0; axis < labels_shape.size(); ++axis) { + if (affinities_or_dummy.shape[axis + 1] != labels_shape[axis]) { + throw std::invalid_argument( + "affinities spatial shape must match labels shape" + ); + } + } + } + for (const auto &offset : offsets) { + if (offset.size() != labels_shape.size()) { + throw std::invalid_argument("each offset must match labels ndim"); + } + } +} + +template +void discover_lifted_chunk( + const RegionAdjacencyGraph &rag, + const LabelT *labels, + const std::vector> &offsets, + const std::vector &long_range_channels, + const std::vector &shape, + const std::vector &strides, + const std::size_t node_begin, + const std::size_t node_end, + std::unordered_set &out +) { + for (const auto channel : long_range_channels) { + const auto &offset = offsets[channel]; + for (std::uint64_t node = node_begin; node < node_end; ++node) { + std::uint64_t target = 0; + if (!bioimage_cpp::detail::valid_offset_target(node, offset, shape, strides, target)) { + continue; + } + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + continue; + } + if (rag.find_edge(u, v) >= 0) { + continue; + } + out.insert(bioimage_cpp::detail::edge_key(u, v)); + } + } +} + +} // namespace detail_lifted + +// Discover lifted edges implied by long-range offsets on the affinity grid. +// +// Walks every grid coordinate together with each long-range offset (offsets +// whose L1 norm is > 1; 1-hop offsets are silently skipped). When labels at +// (p, p + offset) differ and the (u, v) pair is not already a local RAG edge, +// (u, v) is recorded as a lifted edge. Returns the deduplicated set sorted +// lexicographically with `u < v`. +template +std::vector lifted_edges_from_offsets( + const RegionAdjacencyGraph &rag, + const ConstArrayView &labels, + const std::vector> &offsets, + const std::size_t number_of_threads +) { + if (labels.ndim() != 2 && labels.ndim() != 3) { + throw std::invalid_argument("labels must be a 2D or 3D array"); + } + detail_lifted::validate_affinity_inputs( + ConstArrayView{nullptr, {}, {}}, + labels.shape, + offsets, + /*has_affinities=*/false + ); + + std::vector long_range_channels; + long_range_channels.reserve(offsets.size()); + for (std::size_t channel = 0; channel < offsets.size(); ++channel) { + if (detail_lifted::is_long_range(offsets[channel])) { + long_range_channels.push_back(channel); + } + } + if (long_range_channels.empty()) { + return {}; + } + + const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); + const auto n_threads = bioimage_cpp::detail::normalize_thread_count( + number_of_threads, number_of_nodes + ); + const auto strides = bioimage_cpp::detail::c_order_strides(labels.shape); + + using EdgeSet = std::unordered_set< + bioimage_cpp::detail::Edge, bioimage_cpp::detail::EdgeHash + >; + std::vector per_thread(n_threads); + + bioimage_cpp::detail::parallel_for_chunks( + n_threads, + number_of_nodes, + [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { + detail_lifted::discover_lifted_chunk( + rag, labels.data, offsets, long_range_channels, + labels.shape, strides, begin, end, per_thread[thread_id] + ); + } + ); + + EdgeSet merged; + std::size_t total = 0; + for (const auto &set : per_thread) { + total += set.size(); + } + merged.reserve(total); + for (auto &set : per_thread) { + merged.insert(set.begin(), set.end()); + } + + std::vector result(merged.begin(), merged.end()); + std::sort(result.begin(), result.end()); + return result; +} + +namespace detail_lifted { + +template +void scan_lifted_affinity_chunk( + const std::unordered_map + &lifted_index, + const LabelT *labels, + const ValueT *affinities, + const std::vector> &offsets, + const std::vector &long_range_channels, + const std::vector &shape, + const std::size_t node_begin, + const std::size_t node_end, + std::vector &stats +) { + const auto strides = bioimage_cpp::detail::c_order_strides(shape); + const auto number_of_nodes = static_cast(detail_features::number_of_pixels(shape)); + for (const auto channel : long_range_channels) { + const auto &offset = offsets[channel]; + const auto channel_offset = static_cast(channel) * number_of_nodes; + for (std::uint64_t node = node_begin; node < node_end; ++node) { + std::uint64_t target = 0; + if (!bioimage_cpp::detail::valid_offset_target(node, offset, shape, strides, target)) { + continue; + } + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + continue; + } + const auto found = lifted_index.find(bioimage_cpp::detail::edge_key(u, v)); + if (found == lifted_index.end()) { + continue; + } + stats[found->second].add(affinities[channel_offset + node]); + } + } +} + +} // namespace detail_lifted + +// Accumulate affinity statistics onto a caller-supplied lifted edge set. +// +// `lifted_uvs` lists the (u, v) pairs to bin into, one row per lifted edge; +// it is typically the output of `lifted_edges_from_offsets`. Pixel pairs +// (p, p + offset) whose endpoints differ and whose (u, v) appears in +// `lifted_uvs` contribute their affinity value to that lifted edge's stats. +// Pairs that hit a local-only or unknown edge are silently skipped, so a +// local edge that happens to be reachable via a long-range offset is not +// contaminated by long-range affinities. +// +// 1-hop offsets are skipped automatically so callers can pass the full +// offset list without pre-filtering. +template +void accumulate_lifted_affinity_features( + const ConstArrayView &labels, + const ConstArrayView &affinities, + const std::vector> &offsets, + const std::vector &lifted_uvs, + const bool compute_complex_features, + const std::size_t number_of_threads, + const ArrayView &out +) { + if (labels.ndim() != 2 && labels.ndim() != 3) { + throw std::invalid_argument("labels must be a 2D or 3D array"); + } + if (affinities.ndim() != labels.ndim() + 1) { + throw std::invalid_argument( + "affinities must have shape (channels, *labels.shape)" + ); + } + if (static_cast(affinities.shape[0]) != offsets.size()) { + throw std::invalid_argument( + "offsets length must match affinities channel count" + ); + } + for (std::size_t axis = 0; axis < labels.shape.size(); ++axis) { + if (affinities.shape[axis + 1] != labels.shape[axis]) { + throw std::invalid_argument( + "affinities spatial shape must match labels shape" + ); + } + } + for (const auto &offset : offsets) { + if (offset.size() != static_cast(labels.ndim())) { + throw std::invalid_argument("each offset must match labels ndim"); + } + } + + const auto expected_features = compute_complex_features ? 12 : 2; + if (out.shape != std::vector{ + static_cast(lifted_uvs.size()), expected_features}) { + throw std::invalid_argument( + "out shape must be (number_of_lifted_edges, number_of_features)" + ); + } + + std::vector long_range_channels; + long_range_channels.reserve(offsets.size()); + for (std::size_t channel = 0; channel < offsets.size(); ++channel) { + if (detail_lifted::is_long_range(offsets[channel])) { + long_range_channels.push_back(channel); + } + } + + std::unordered_map + lifted_index; + lifted_index.reserve(lifted_uvs.size()); + for (std::size_t i = 0; i < lifted_uvs.size(); ++i) { + const auto key = bioimage_cpp::detail::edge_key(lifted_uvs[i].first, lifted_uvs[i].second); + if (!lifted_index.emplace(key, i).second) { + throw std::invalid_argument("lifted_uvs must not contain duplicate edges"); + } + } + + const auto number_of_lifted = lifted_uvs.size(); + if (long_range_channels.empty() || number_of_lifted == 0) { + if (compute_complex_features) { + std::vector empty(number_of_lifted); + detail_features::write_complex_features(empty, out); + } else { + std::vector empty(number_of_lifted); + detail_features::write_simple_features(empty, out); + } + return; + } + + const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); + const auto n_threads = bioimage_cpp::detail::normalize_thread_count( + number_of_threads, number_of_nodes + ); + + const auto run_scan = [&](auto &per_thread_stats) { + bioimage_cpp::detail::parallel_for_chunks( + n_threads, + number_of_nodes, + [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { + detail_lifted::scan_lifted_affinity_chunk( + lifted_index, labels.data, affinities.data, + offsets, long_range_channels, labels.shape, + begin, end, per_thread_stats[thread_id] + ); + } + ); + }; + + if (compute_complex_features) { + std::vector> per_thread_stats( + n_threads, + std::vector(number_of_lifted) + ); + run_scan(per_thread_stats); + auto stats = detail_features::merge_stats(per_thread_stats, number_of_lifted); + detail_features::write_complex_features(stats, out); + } else { + std::vector> per_thread_stats( + n_threads, + std::vector(number_of_lifted) + ); + run_scan(per_thread_stats); + auto stats = detail_features::merge_stats(per_thread_stats, number_of_lifted); + detail_features::write_simple_features(stats, out); + } +} + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/lifted_multicut.hxx b/include/bioimage_cpp/graph/lifted_multicut.hxx new file mode 100644 index 0000000..d2720d8 --- /dev/null +++ b/include/bioimage_cpp/graph/lifted_multicut.hxx @@ -0,0 +1,5 @@ +#pragma once + +#include "bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx" +#include "bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx" +#include "bioimage_cpp/graph/lifted_multicut/objective.hxx" diff --git a/include/bioimage_cpp/graph/lifted_multicut/detail.hxx b/include/bioimage_cpp/graph/lifted_multicut/detail.hxx new file mode 100644 index 0000000..051bc41 --- /dev/null +++ b/include/bioimage_cpp/graph/lifted_multicut/detail.hxx @@ -0,0 +1,276 @@ +#pragma once + +#include "bioimage_cpp/detail/indexed_heap.hxx" +#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/graph/connected_components.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph::lifted_multicut::detail { + +inline constexpr std::size_t no_edge = std::numeric_limits::max(); + +// Per-super-node adjacency entry. +struct NeighborEntry { + std::size_t neighbor; + std::size_t edge_id; +}; + +// Per-edge data stored in a flat vector indexed by stable edge_id. Same shape +// as multicut::detail::DynamicEdge but with `is_lifted` replacing the multicut +// constraint flag, since lifted edges have different heap semantics: +// * lifted edges are excluded from the heap until they become non-lifted; +// * lifted propagates iff *both* inputs are lifted (vs. constraint, which +// propagates if *either* input is constrained). +struct DynamicEdge { + std::size_t u = 0; + std::size_t v = 0; + double weight = 0.0; + unsigned char is_lifted = 0; +}; + +using EdgeHeap = bioimage_cpp::detail::DenseIndexedHeap; + +struct DynamicGraph { + DynamicGraph() = default; + + explicit DynamicGraph(const UndirectedGraph &lifted_graph) { + reset(lifted_graph); + } + + void reset(const UndirectedGraph &lifted_graph) { + const auto n_nodes = static_cast(lifted_graph.number_of_nodes()); + const auto n_edges = static_cast(lifted_graph.number_of_edges()); + + for (auto &adj : adjacency) { + adj.clear(); + } + adjacency.resize(n_nodes); + for (std::uint64_t node = 0; node < lifted_graph.number_of_nodes(); ++node) { + const auto degree = lifted_graph.node_adjacency(node).size(); + adjacency[static_cast(node)].reserve(degree); + } + + alive.assign(n_nodes, true); + alive_count = n_nodes; + scratch_edge_id.assign(n_nodes, no_edge); + edges.resize(n_edges); + } + + std::vector> adjacency; + std::vector edges; + std::vector alive; + std::size_t alive_count; + std::vector scratch_edge_id; +}; + +// Initialize the dynamic graph from a lifted graph + per-edge weights, where +// the first `n_base_edges` entries of the lifted graph are the base edges. +// Optional Gaussian noise is added to the weights, mirroring the multicut +// greedy additive flow. +// +// Lifted edges enter the dynamic graph but are *not* pushed onto the heap (the +// driver guards them via `is_lifted`). Non-lifted (base) edges are heap-pushed +// at their (possibly noisy) weight. +inline void initialize_dynamic_graph( + const UndirectedGraph &lifted_graph, + const std::vector &weights, + const std::uint64_t n_base_edges, + DynamicGraph &dynamic_graph, + EdgeHeap &heap, + const bool add_noise = false, + const int seed = 42, + const double sigma = 1.0 +) { + const auto n_edges = static_cast(lifted_graph.number_of_edges()); + heap.reset_capacity(n_edges); + + std::vector heap_entries; + heap_entries.reserve(static_cast(n_base_edges)); + + std::mt19937 generator(seed); + std::normal_distribution noise(0.0, sigma); + for (std::uint64_t edge = 0; edge < lifted_graph.number_of_edges(); ++edge) { + const auto uv = lifted_graph.uv(edge); + const auto u = static_cast(uv.first); + const auto v = static_cast(uv.second); + double weight = weights[static_cast(edge)]; + if (add_noise) { + weight += noise(generator); + } + const auto edge_id = static_cast(edge); + auto &e = dynamic_graph.edges[edge_id]; + e.u = u; + e.v = v; + e.weight = weight; + e.is_lifted = (edge < n_base_edges) ? 0 : 1; + dynamic_graph.adjacency[u].push_back({v, edge_id}); + dynamic_graph.adjacency[v].push_back({u, edge_id}); + if (e.is_lifted == 0) { + heap_entries.push_back({edge_id, weight}); + } + } + + heap.build_heap(std::move(heap_entries)); +} + +namespace internal { + +inline bool erase_by_neighbor(std::vector &list, const std::size_t target) { + for (std::size_t i = 0; i < list.size(); ++i) { + if (list[i].neighbor == target) { + list[i] = list.back(); + list.pop_back(); + return true; + } + } + return false; +} + +inline void rename_neighbor( + std::vector &list, + const std::size_t from_node, + const std::size_t to_node +) { + for (auto &entry : list) { + if (entry.neighbor == from_node) { + entry.neighbor = to_node; + return; + } + } +} + +} // namespace internal + +// Contract the edge between u and v in the dynamic graph. Folds the +// smaller-degree super-node into the larger one. Heap is kept in sync with the +// following lifted-aware rules for parallel-edge merges: +// * both inputs lifted → result lifted (stays out of heap). +// * both inputs base → result base, priority = summed weight. +// * one of each kind → result base, edge re-enters or is updated in the +// heap with priority = summed weight. +// +// Precondition: the edge being contracted is non-lifted (the driver only +// pops base edges off the heap). +inline std::size_t merge_dynamic_nodes( + DynamicGraph &dynamic_graph, + bioimage_cpp::detail::UnionFind &sets, + EdgeHeap &heap, + std::size_t u, + std::size_t v +) { + u = static_cast(sets.find(u)); + v = static_cast(sets.find(v)); + if (u == v) { + return u; + } + + auto stable = u; + auto removed = v; + if (dynamic_graph.adjacency[stable].size() < dynamic_graph.adjacency[removed].size()) { + std::swap(stable, removed); + } + sets.merge_to(stable, removed); + + for (const auto &entry : dynamic_graph.adjacency[stable]) { + dynamic_graph.scratch_edge_id[entry.neighbor] = entry.edge_id; + } + + const auto contracted_edge_id = dynamic_graph.scratch_edge_id[removed]; + heap.erase(contracted_edge_id); + dynamic_graph.scratch_edge_id[removed] = no_edge; + internal::erase_by_neighbor(dynamic_graph.adjacency[stable], removed); + + const auto removed_neighbors = dynamic_graph.adjacency[removed]; + + for (const auto &entry : removed_neighbors) { + const auto neighbor = entry.neighbor; + const auto removed_edge_id = entry.edge_id; + if (neighbor == stable) { + continue; + } + + const auto existing_id = dynamic_graph.scratch_edge_id[neighbor]; + if (existing_id == no_edge) { + // No stable-side edge to `neighbor`: re-key the removed-side edge + // by replacing `removed` with `stable`. Lifted flag and weight are + // preserved, so the heap state stays consistent (lifted edges stay + // off the heap; base edges keep their existing heap entry). + dynamic_graph.adjacency[stable].push_back({neighbor, removed_edge_id}); + dynamic_graph.scratch_edge_id[neighbor] = removed_edge_id; + internal::rename_neighbor(dynamic_graph.adjacency[neighbor], removed, stable); + auto &e = dynamic_graph.edges[removed_edge_id]; + if (e.u == removed) { + e.u = stable; + } else { + e.v = stable; + } + } else { + // Stable already has an edge to `neighbor`. Sum the weights and + // resolve the lifted flag. + auto &keep = dynamic_graph.edges[existing_id]; + const auto &fold = dynamic_graph.edges[removed_edge_id]; + const bool keep_was_lifted = keep.is_lifted != 0; + const bool fold_was_lifted = fold.is_lifted != 0; + keep.weight += fold.weight; + const bool result_lifted = keep_was_lifted && fold_was_lifted; + keep.is_lifted = result_lifted ? 1 : 0; + + internal::erase_by_neighbor(dynamic_graph.adjacency[neighbor], removed); + + if (result_lifted) { + // Both inputs were lifted: the result is still lifted, so the + // heap should not contain either side. (Both were already off.) + } else if (keep_was_lifted && !fold_was_lifted) { + // The fold side was the base edge and lives on the heap; we + // are dropping that id. The keep side (formerly lifted) needs + // to be inserted into the heap with the summed weight. + heap.erase(removed_edge_id); + heap.push(existing_id, keep.weight); + } else if (!keep_was_lifted && fold_was_lifted) { + // The keep side is on the heap; the fold side was lifted and + // is not. Update the keep entry's priority to the summed weight. + heap.change(existing_id, keep.weight); + } else { + // Both base edges: drop the removed-side heap entry and + // refresh the kept entry's priority. + heap.erase(removed_edge_id); + heap.change(existing_id, keep.weight); + } + } + } + + for (const auto &entry : dynamic_graph.adjacency[stable]) { + dynamic_graph.scratch_edge_id[entry.neighbor] = no_edge; + } + + dynamic_graph.adjacency[removed].clear(); + dynamic_graph.alive[removed] = false; + --dynamic_graph.alive_count; + return stable; +} + +inline std::vector labels_from_sets( + bioimage_cpp::detail::UnionFind &sets, + const UndirectedGraph &graph +) { + return dense_labels_from_union_find(sets, graph.number_of_nodes()); +} + +inline std::size_t stop_node_count( + const UndirectedGraph &graph, + const double node_num_stop +) { + return node_num_stop >= 1.0 + ? static_cast(node_num_stop) + : static_cast(double(graph.number_of_nodes()) * node_num_stop + 0.5); +} + +} // namespace bioimage_cpp::graph::lifted_multicut::detail diff --git a/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx b/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx new file mode 100644 index 0000000..d28e089 --- /dev/null +++ b/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx @@ -0,0 +1,146 @@ +#pragma once + +#include "bioimage_cpp/detail/profile.hxx" +#include "bioimage_cpp/graph/lifted_multicut/detail.hxx" +#include "bioimage_cpp/graph/lifted_multicut/objective.hxx" + +#include +#include +#include + +namespace bioimage_cpp::graph::lifted_multicut { + +// Reusable scratch state for `lifted_greedy_additive`. +struct GreedyAdditiveWorkspace { + detail::DynamicGraph dynamic_graph; + bioimage_cpp::detail::UnionFind union_find{0}; + detail::EdgeHeap heap; + + void reset(const UndirectedGraph &lifted_graph) { + dynamic_graph.reset(lifted_graph); + union_find.reset(static_cast(lifted_graph.number_of_nodes())); + heap.reset_capacity(static_cast(lifted_graph.number_of_edges())); + } +}; + +// Greedy-additive multicut on the lifted graph. Identical contraction flow to +// multicut::greedy_additive, with one twist: lifted edges are never contracted +// directly. They influence the energy via weight accumulation during merges — +// a lifted edge between super-nodes contributes its weight to the contracted +// edge whenever the two super-nodes also share a non-lifted edge, and the +// combined edge stops being lifted as soon as a base-graph path connects its +// endpoints. +inline std::vector greedy_additive( + const UndirectedGraph &lifted_graph, + const std::vector &weights, + const std::uint64_t n_base_edges, + const double weight_stop, + const double node_num_stop, + const bool add_noise, + const int seed, + const double sigma, + GreedyAdditiveWorkspace &workspace +) { + BIOIMAGE_PROFILE_INIT(profile); + validate_weights(lifted_graph, weights); + { + BIOIMAGE_PROFILE_SCOPE(profile, "workspace_reset"); + workspace.reset(lifted_graph); + } + auto &dynamic_graph = workspace.dynamic_graph; + auto &sets = workspace.union_find; + auto &heap = workspace.heap; + { + BIOIMAGE_PROFILE_SCOPE(profile, "initialize"); + detail::initialize_dynamic_graph( + lifted_graph, weights, n_base_edges, dynamic_graph, heap, add_noise, seed, sigma + ); + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "contraction_loop"); + while (!heap.empty() && dynamic_graph.alive_count > 1) { + const auto top = heap.top(); + if (top.priority <= weight_stop) { + break; + } + if (node_num_stop > 0.0 + && dynamic_graph.alive_count <= detail::stop_node_count(lifted_graph, node_num_stop)) { + break; + } + const auto &edge = dynamic_graph.edges[top.key]; + detail::merge_dynamic_nodes(dynamic_graph, sets, heap, edge.u, edge.v); + } + } + std::vector labels; + { + BIOIMAGE_PROFILE_SCOPE(profile, "labels_from_sets"); + labels = detail::labels_from_sets(sets, lifted_graph); + } + BIOIMAGE_PROFILE_REPORT(profile); + return labels; +} + +inline std::vector greedy_additive( + const UndirectedGraph &lifted_graph, + const std::vector &weights, + const std::uint64_t n_base_edges, + const double weight_stop, + const double node_num_stop, + const bool add_noise, + const int seed, + const double sigma +) { + GreedyAdditiveWorkspace workspace; + return greedy_additive( + lifted_graph, + weights, + n_base_edges, + weight_stop, + node_num_stop, + add_noise, + seed, + sigma, + workspace + ); +} + +class GreedyAdditiveSolver final : public SolverBase { +public: + GreedyAdditiveSolver( + const double weight_stop = 0.0, + const double node_num_stop = -1.0, + const bool add_noise = false, + const int seed = 42, + const double sigma = 1.0 + ) + : weight_stop_(weight_stop), + node_num_stop_(node_num_stop), + add_noise_(add_noise), + seed_(seed), + sigma_(sigma) {} + + std::vector optimize(Objective &objective) const override { + auto labels = greedy_additive( + objective.lifted_graph(), + objective.weights(), + objective.number_of_base_edges(), + weight_stop_, + node_num_stop_, + add_noise_, + seed_, + sigma_ + ); + objective.set_labels(labels); + return labels; + } + +private: + double weight_stop_; + double node_num_stop_; + bool add_noise_; + int seed_; + double sigma_; +}; + +} // namespace bioimage_cpp::graph::lifted_multicut diff --git a/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx b/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx new file mode 100644 index 0000000..b309816 --- /dev/null +++ b/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx @@ -0,0 +1,521 @@ +#pragma once + +#include "bioimage_cpp/detail/edge_hash.hxx" +#include "bioimage_cpp/detail/indexed_heap.hxx" +#include "bioimage_cpp/detail/profile.hxx" +#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/graph/connected_components.hxx" +#include "bioimage_cpp/graph/lifted_multicut/objective.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph::lifted_multicut { + +namespace detail_kl { + +struct ClusterPair { + std::uint64_t a; + std::uint64_t b; +}; + +// Build the set of cluster pairs that share at least one base-graph edge. +// These are the only pairs eligible for the two-cut chain in lifted KL — +// merging non-base-adjacent clusters is undone by the post-iteration +// connected-component repartition. +inline std::vector compute_base_cluster_pairs( + const UndirectedGraph &base_graph, + const std::vector &labels +) { + using bioimage_cpp::detail::Edge; + using bioimage_cpp::detail::EdgeHash; + using bioimage_cpp::detail::edge_key; + + std::unordered_map table; + table.reserve(static_cast(base_graph.number_of_edges())); + for (std::uint64_t edge = 0; edge < base_graph.number_of_edges(); ++edge) { + const auto uv = base_graph.uv(edge); + const auto la = labels[static_cast(uv.first)]; + const auto lb = labels[static_cast(uv.second)]; + if (la == lb) { + continue; + } + table.emplace(edge_key(la, lb), std::uint8_t{1}); + } + + std::vector result; + result.reserve(table.size()); + for (const auto &entry : table) { + result.push_back({entry.first.first, entry.first.second}); + } + std::sort(result.begin(), result.end(), [](const ClusterPair &lhs, const ClusterPair &rhs) { + return std::tie(lhs.a, lhs.b) < std::tie(rhs.a, rhs.b); + }); + return result; +} + +inline std::vector> build_cluster_to_nodes( + const std::vector &labels, + const std::uint64_t number_of_clusters +) { + std::vector> result(static_cast(number_of_clusters)); + for (std::uint64_t node = 0; node < labels.size(); ++node) { + result[static_cast(labels[static_cast(node)])].push_back(node); + } + return result; +} + +// Same per-node scratch shape as multicut::detail_kl::ChainBuffers but +// duplicated to keep the two KL files independent (the helpers are small). +struct ChainBuffers { + std::vector in_pair; + std::vector moved; + std::vector cross_count; + std::vector stash_gain; + + explicit ChainBuffers(const std::size_t n_nodes) + : in_pair(n_nodes, 0), + moved(n_nodes, 0), + cross_count(n_nodes, 0), + stash_gain(n_nodes, 0.0) {} +}; + +// Per-edge entry of the per-node filtered adjacency built once during +// `chain_gain_init` and consumed by `chain_loop`. Lifted-graph adjacency is +// walked exactly once per chain (during init) and the few entries that +// survive the in-pair filter are appended here, with their weight and base- +// vs-lifted classification cached. The chain loop then iterates only these +// surviving entries — for typical small pairs in a large graph that's a +// >10× reduction in inner-loop iterations versus re-walking +// ``lifted_graph.node_adjacency(v)`` from scratch on every move. +struct FilteredAdj { + std::uint64_t node; + double weight; + bool is_base; +}; + +struct ChainScratch { + std::vector queue_nodes; + bioimage_cpp::detail::DenseIndexedHeap heap; + + // Per-node CSR-style range into ``filtered_entries``. ``filtered_count[v]`` + // is set whenever v is part of the current chain's pair; readers must + // therefore index only into nodes known to be in-pair (i.e. nodes popped + // from the heap, all of which were pushed during init). + std::vector filtered_offset; + std::vector filtered_count; + std::vector filtered_entries; + + explicit ChainScratch(const std::size_t n_nodes) + : heap(n_nodes), + filtered_offset(n_nodes, 0), + filtered_count(n_nodes, 0) { + filtered_entries.reserve(n_nodes * 4); + } +}; + +// Run a Kernighan-Lin move-chain on (cluster_a, cluster_b), restricted by the +// base graph for connectivity (cross_count) and driven by the lifted graph +// for energy gains. +// +// As a parallel branch, also consider fully merging cluster_b into cluster_a +// — that is the limit of a chain that moves every B-node into A. The merge +// alternative is committed if its lifted-cut gain beats the best chain prefix +// gain. This matches nifty's two-cut behaviour. +// +// Returns the committed cumulative gain (>= 0). +template +inline double run_chain( + const UndirectedGraph &base_graph, + const UndirectedGraph &lifted_graph, + const std::vector &lifted_weights, + const std::uint64_t n_base_edges, + std::vector &labels, + std::vector> &cluster_to_nodes, + ChainBuffers &bufs, + ChainScratch &scratch, + const std::uint64_t cluster_a, + const std::uint64_t cluster_b, + const double epsilon, + [[maybe_unused]] ProfilerT &profile +) { + if (cluster_a == cluster_b) { + return 0.0; + } + + auto &queue_nodes = scratch.queue_nodes; + auto &heap = scratch.heap; + auto &filtered_entries = scratch.filtered_entries; + queue_nodes.clear(); + heap.clear(); + filtered_entries.clear(); + + std::size_t live_a = 0; + std::size_t live_b = 0; + { + BIOIMAGE_PROFILE_SCOPE(profile, "chain_init"); + const auto &stale_a = cluster_to_nodes[static_cast(cluster_a)]; + const auto &stale_b = cluster_to_nodes[static_cast(cluster_b)]; + queue_nodes.reserve(stale_a.size() + stale_b.size()); + + for (const auto v : stale_a) { + if (labels[static_cast(v)] == cluster_a) { + queue_nodes.push_back(v); + bufs.in_pair[static_cast(v)] = 1; + ++live_a; + } + } + for (const auto v : stale_b) { + if (labels[static_cast(v)] == cluster_b) { + queue_nodes.push_back(v); + bufs.in_pair[static_cast(v)] = 1; + ++live_b; + } + } + } + if (live_a + live_b < 2 || (live_a == 1 && live_b == 1)) { + for (const auto v : queue_nodes) { + bufs.in_pair[static_cast(v)] = 0; + } + return 0.0; + } + + const bool is_split = (live_b == 0); + + // First pass: compute initial gains (over lifted graph) and cross_count + // (over base graph), and accumulate the cross-side lifted-weight sum used + // by the merge alternative. + double gain_from_merging_double = 0.0; + { + BIOIMAGE_PROFILE_SCOPE(profile, "chain_gain_init"); + for (const auto v : queue_nodes) { + double w_to_a = 0.0; + double w_to_b = 0.0; + std::uint32_t cross = 0; + const auto v_label = labels[static_cast(v)]; + const auto v_key = static_cast(v); + const auto filter_start = filtered_entries.size(); + for (const auto adj : lifted_graph.node_adjacency(v)) { + const auto u_key = static_cast(adj.node); + if (!bufs.in_pair[u_key]) { + continue; + } + const auto w = lifted_weights[static_cast(adj.edge)]; + const auto u_label = labels[u_key]; + const bool is_base = adj.edge < n_base_edges; + filtered_entries.push_back({adj.node, w, is_base}); + if (u_label == cluster_a) { + w_to_a += w; + } else { + w_to_b += w; + } + if (u_label != v_label) { + if (is_base) { + ++cross; + } + gain_from_merging_double += w; + } + } + scratch.filtered_offset[v_key] = filter_start; + scratch.filtered_count[v_key] = + static_cast(filtered_entries.size() - filter_start); + const double gain_v = (v_label == cluster_a) ? (w_to_b - w_to_a) : (w_to_a - w_to_b); + bufs.stash_gain[v_key] = gain_v; + bufs.cross_count[v_key] = cross; + if (is_split || cross > 0) { + heap.push(v_key, gain_v); + } + } + } // end chain_gain_init scope + const double gain_from_merging = is_split ? 0.0 : 0.5 * gain_from_merging_double; + + struct Move { + std::uint64_t node; + std::uint64_t new_label; + }; + std::vector chain; + chain.reserve(queue_nodes.size()); + + double cumulative = 0.0; + double best_cumulative = 0.0; + std::size_t best_prefix = 0; + + { + BIOIMAGE_PROFILE_SCOPE(profile, "chain_loop"); + while (!heap.empty()) { + const auto top = heap.pop(); + const auto v = static_cast(top.key); + const auto gain_v = top.priority; + const auto v_key = static_cast(v); + const auto old_label = labels[v_key]; + const auto new_label = (old_label == cluster_a) ? cluster_b : cluster_a; + + bufs.moved[v_key] = 1; + cumulative += gain_v; + chain.push_back({v, new_label}); + + if (cumulative > best_cumulative + epsilon) { + best_cumulative = cumulative; + best_prefix = chain.size(); + } + + // Walk the pre-built in-pair adjacency; each entry is guaranteed + // ``bufs.in_pair[u_key] == 1`` so we only need to filter on + // ``bufs.moved``. ``was_in_heap`` is cached so the three subsequent + // heap operations don't each re-query the locator. + const auto offset = scratch.filtered_offset[v_key]; + const auto count = scratch.filtered_count[v_key]; + for (std::uint32_t i = 0; i < count; ++i) { + const auto &fa = scratch.filtered_entries[offset + i]; + const auto u_key = static_cast(fa.node); + if (bufs.moved[u_key]) { + continue; + } + const auto u_label = labels[u_key]; + const double delta = (u_label == old_label) ? 2.0 * fa.weight : -2.0 * fa.weight; + bufs.stash_gain[u_key] += delta; + const bool was_in_heap = heap.contains(u_key); + if (was_in_heap) { + heap.change(u_key, bufs.stash_gain[u_key]); + } + // Border maintenance is base-graph-only. Lifted-only edges affect + // the gain but cannot make a node bordered. + if (fa.is_base) { + if (u_label == old_label) { + ++bufs.cross_count[u_key]; + if (!is_split && !was_in_heap) { + heap.push(u_key, bufs.stash_gain[u_key]); + } + } else { + if (bufs.cross_count[u_key] > 0) { + --bufs.cross_count[u_key]; + } + if (!is_split && bufs.cross_count[u_key] == 0 && was_in_heap) { + heap.erase(u_key); + } + } + } + } + } + + } // end chain_loop scope + + { + BIOIMAGE_PROFILE_SCOPE(profile, "chain_cleanup"); + for (const auto v : queue_nodes) { + const auto v_key = static_cast(v); + bufs.in_pair[v_key] = 0; + bufs.moved[v_key] = 0; + bufs.cross_count[v_key] = 0; + } + } // end chain_cleanup scope + + // Decide between best chain prefix and full-merge alternative. + if (gain_from_merging > best_cumulative + epsilon && gain_from_merging > epsilon) { + // Move all live B-nodes into A. + for (const auto v : queue_nodes) { + if (labels[static_cast(v)] == cluster_b) { + labels[static_cast(v)] = cluster_a; + cluster_to_nodes[static_cast(cluster_a)].push_back(v); + } + } + return gain_from_merging; + } + + if (best_cumulative > epsilon) { + for (std::size_t i = 0; i < best_prefix; ++i) { + const auto v = chain[i].node; + const auto new_label = chain[i].new_label; + labels[static_cast(v)] = new_label; + cluster_to_nodes[static_cast(new_label)].push_back(v); + } + return best_cumulative; + } + return 0.0; +} + +// Re-split labels so that every cluster is connected in the base graph. +// Returns the new labeling (dense relabeled). +inline std::vector enforce_base_connectivity( + const UndirectedGraph &base_graph, + const std::vector &labels +) { + const auto n_edges = static_cast(base_graph.number_of_edges()); + std::vector mask(n_edges); + for (std::uint64_t edge = 0; edge < base_graph.number_of_edges(); ++edge) { + const auto uv = base_graph.uv(edge); + mask[static_cast(edge)] = + labels[static_cast(uv.first)] + == labels[static_cast(uv.second)] + ? 1 : 0; + } + return connected_components(base_graph, mask.empty() ? nullptr : mask.data()); +} + +} // namespace detail_kl + +inline std::vector kernighan_lin( + const UndirectedGraph &base_graph, + const UndirectedGraph &lifted_graph, + const std::vector &lifted_weights, + const std::uint64_t n_base_edges, + std::vector labels, + const std::uint64_t number_of_outer_iterations, + const double epsilon +) { + BIOIMAGE_PROFILE_INIT(profile); + validate_weights(lifted_graph, lifted_weights); + validate_labels(base_graph, labels); + + // Make sure every cluster is base-graph connected before we start; the + // chain assumes this invariant. + { + BIOIMAGE_PROFILE_SCOPE(profile, "cc_repartition"); + labels = detail_kl::enforce_base_connectivity(base_graph, labels); + labels = dense_relabel(labels); + } + + double current_energy; + { + BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); + current_energy = energy(lifted_graph, lifted_weights, labels); + } + auto last_good = labels; + double last_good_energy = current_energy; + + const auto n_nodes = static_cast(base_graph.number_of_nodes()); + detail_kl::ChainBuffers bufs(n_nodes); + detail_kl::ChainScratch scratch(n_nodes); + + for (std::uint64_t iteration = 0; iteration < number_of_outer_iterations; ++iteration) { + bool improved = false; + + std::vector pairs; + std::uint64_t number_of_clusters = 0; + std::vector> cluster_to_nodes; + { + BIOIMAGE_PROFILE_SCOPE(profile, "compute_pairs"); + pairs = detail_kl::compute_base_cluster_pairs(base_graph, labels); + number_of_clusters = labels.empty() + ? std::uint64_t{0} + : (*std::max_element(labels.begin(), labels.end()) + 1); + cluster_to_nodes = detail_kl::build_cluster_to_nodes(labels, number_of_clusters); + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "pair_chains"); + for (const auto &pair : pairs) { + const auto delta = detail_kl::run_chain( + base_graph, + lifted_graph, + lifted_weights, + n_base_edges, + labels, + cluster_to_nodes, + bufs, + scratch, + pair.a, + pair.b, + epsilon, + profile + ); + if (delta > epsilon) { + improved = true; + } + } + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "cluster_splits"); + std::uint64_t next_label = number_of_clusters; + for (std::uint64_t cluster = 0; cluster < number_of_clusters; ++cluster) { + while (true) { + if (next_label >= cluster_to_nodes.size()) { + cluster_to_nodes.resize(static_cast(next_label) + 1); + } + const auto delta = detail_kl::run_chain( + base_graph, + lifted_graph, + lifted_weights, + n_base_edges, + labels, + cluster_to_nodes, + bufs, + scratch, + cluster, + next_label, + epsilon, + profile + ); + if (delta <= epsilon) { + break; + } + improved = true; + ++next_label; + } + } + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "cc_repartition"); + labels = detail_kl::enforce_base_connectivity(base_graph, labels); + labels = dense_relabel(labels); + } + + double new_energy; + { + BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); + new_energy = energy(lifted_graph, lifted_weights, labels); + } + if (new_energy + epsilon < last_good_energy) { + last_good = labels; + last_good_energy = new_energy; + current_energy = new_energy; + } else { + labels = last_good; + current_energy = last_good_energy; + break; + } + + if (!improved) { + break; + } + } + BIOIMAGE_PROFILE_REPORT(profile); + return labels; +} + +class KernighanLinSolver final : public SolverBase { +public: + KernighanLinSolver( + const std::uint64_t number_of_outer_iterations = 100, + const double epsilon = 1.0e-6 + ) + : number_of_outer_iterations_(number_of_outer_iterations), + epsilon_(epsilon) {} + + std::vector optimize(Objective &objective) const override { + auto labels = kernighan_lin( + objective.graph(), + objective.lifted_graph(), + objective.weights(), + objective.number_of_base_edges(), + objective.labels(), + number_of_outer_iterations_, + epsilon_ + ); + objective.set_labels(labels); + return labels; + } + +private: + std::uint64_t number_of_outer_iterations_; + double epsilon_; +}; + +} // namespace bioimage_cpp::graph::lifted_multicut diff --git a/include/bioimage_cpp/graph/lifted_multicut/objective.hxx b/include/bioimage_cpp/graph/lifted_multicut/objective.hxx new file mode 100644 index 0000000..67ba1df --- /dev/null +++ b/include/bioimage_cpp/graph/lifted_multicut/objective.hxx @@ -0,0 +1,264 @@ +#pragma once + +#include "bioimage_cpp/detail/relabel.hxx" +#include "bioimage_cpp/graph/breadth_first_search.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph::lifted_multicut { + +inline void validate_weights( + const UndirectedGraph &lifted_graph, + const std::vector &weights +) { + if (weights.size() != static_cast(lifted_graph.number_of_edges())) { + throw std::invalid_argument( + "edge weights length must match lifted graph number_of_edges" + ); + } +} + +inline void validate_labels( + const UndirectedGraph &graph, + const std::vector &labels +) { + if (labels.size() != static_cast(graph.number_of_nodes())) { + throw std::invalid_argument("labels length must match graph number_of_nodes"); + } +} + +inline std::vector singleton_labels(const UndirectedGraph &graph) { + std::vector labels(static_cast(graph.number_of_nodes())); + std::iota(labels.begin(), labels.end(), std::uint64_t{0}); + return labels; +} + +using bioimage_cpp::detail::dense_relabel; + +// Energy of a node labeling on a lifted graph: sum of weights across cut +// edges. The lifted graph's edge set is base ∪ lifted, so both contribute. +inline double energy( + const UndirectedGraph &lifted_graph, + const std::vector &weights, + const std::vector &labels +) { + validate_weights(lifted_graph, weights); + validate_labels(lifted_graph, labels); + double result = 0.0; + for (std::uint64_t edge = 0; edge < lifted_graph.number_of_edges(); ++edge) { + const auto uv = lifted_graph.uv(edge); + if (labels[static_cast(uv.first)] + != labels[static_cast(uv.second)]) { + result += weights[static_cast(edge)]; + } + } + return result; +} + +// Build a lifted graph as a superset of the base graph. The first +// `base_graph.number_of_edges()` edges of the returned graph are the base +// edges in the same order; subsequent edges are the lifted edges. Duplicate +// (u, v) pairs across base and lifted are not inserted twice — the lifted +// "edge" simply contributes its weight to the corresponding base edge. +// +// Returns: +// - lifted_graph: the constructed UndirectedGraph. +// - weights: vector of length `lifted_graph.number_of_edges()` with the base +// weights followed by lifted contributions. Weights for (u, v) that +// coincide with a base edge are added to the base weight when +// `overwrite_existing == false`, or replace it when `true`. +struct LiftedGraphBuild { + UndirectedGraph lifted_graph; + std::vector weights; +}; + +inline LiftedGraphBuild build_lifted_graph( + const UndirectedGraph &base_graph, + const std::vector &base_weights, + const std::vector> &lifted_uvs, + const std::vector &lifted_weights, + const bool overwrite_existing = false +) { + validate_weights(base_graph, base_weights); + if (lifted_uvs.size() != lifted_weights.size()) { + throw std::invalid_argument( + "lifted_uvs and lifted_weights must have the same length" + ); + } + + UndirectedGraph lifted_graph( + base_graph.number_of_nodes(), + base_graph.number_of_edges() + static_cast(lifted_uvs.size()) + ); + for (std::uint64_t edge = 0; edge < base_graph.number_of_edges(); ++edge) { + const auto uv = base_graph.uv(edge); + lifted_graph.insert_edge(uv.first, uv.second); + } + std::vector weights = base_weights; + weights.resize(static_cast(lifted_graph.number_of_edges()), 0.0); + + for (std::size_t i = 0; i < lifted_uvs.size(); ++i) { + const auto u = lifted_uvs[i].first; + const auto v = lifted_uvs[i].second; + const auto pre = lifted_graph.number_of_edges(); + const auto edge = lifted_graph.insert_edge(u, v); + const auto weight = lifted_weights[i]; + if (lifted_graph.number_of_edges() > pre) { + weights.push_back(weight); + } else { + if (overwrite_existing) { + weights[static_cast(edge)] = weight; + } else { + weights[static_cast(edge)] += weight; + } + } + } + return LiftedGraphBuild{std::move(lifted_graph), std::move(weights)}; +} + +// Lifted multicut objective. Stores a reference to the base graph and owns the +// lifted graph (= base ∪ lifted edges) and the per-lifted-edge weights. The +// node-labeling lives over the base graph's node set. +// +// Invariants: +// - lifted_graph().uv(e) == base_graph.uv(e) for every base edge id e. +// - weights().size() == lifted_graph().number_of_edges(). +class Objective { +public: + Objective( + const UndirectedGraph &base_graph, + std::vector base_weights + ) + : base_graph_(base_graph), + lifted_graph_(build_base_only_(base_graph, base_weights)), + weights_(std::move(base_weights)), + labels_(singleton_labels(base_graph)), + n_base_edges_(base_graph.number_of_edges()) { + validate_weights(lifted_graph_, weights_); + } + + Objective( + const UndirectedGraph &base_graph, + std::vector base_weights, + const std::vector> &lifted_uvs, + const std::vector &lifted_weights, + const bool overwrite_existing = false + ) + : base_graph_(base_graph), + lifted_graph_(), + weights_(), + labels_(singleton_labels(base_graph)), + n_base_edges_(base_graph.number_of_edges()) { + auto build = build_lifted_graph( + base_graph, base_weights, lifted_uvs, lifted_weights, overwrite_existing + ); + lifted_graph_ = std::move(build.lifted_graph); + weights_ = std::move(build.weights); + } + + [[nodiscard]] const UndirectedGraph &graph() const { return base_graph_; } + [[nodiscard]] const UndirectedGraph &lifted_graph() const { return lifted_graph_; } + [[nodiscard]] const std::vector &weights() const { return weights_; } + [[nodiscard]] std::uint64_t number_of_base_edges() const { return n_base_edges_; } + [[nodiscard]] std::uint64_t number_of_lifted_edges() const { + return lifted_graph_.number_of_edges() - n_base_edges_; + } + [[nodiscard]] const std::vector &labels() const { return labels_; } + + // Insert or update a lifted edge. Returns the edge id in the lifted graph + // and whether the edge is newly inserted. If the (u, v) pair already + // exists (as a base edge or as a previously inserted lifted edge), the + // weight is accumulated unless `overwrite` is true. Inserting a brand-new + // edge appends it after all existing edges, so base edges remain at the + // head of the lifted graph. + std::pair set_cost( + const std::uint64_t u, + const std::uint64_t v, + const double weight, + const bool overwrite = false + ) { + const auto pre = lifted_graph_.number_of_edges(); + const auto edge = lifted_graph_.insert_edge(u, v); + if (lifted_graph_.number_of_edges() > pre) { + weights_.push_back(weight); + return {edge, true}; + } + if (overwrite) { + weights_[static_cast(edge)] = weight; + } else { + weights_[static_cast(edge)] += weight; + } + return {edge, false}; + } + + // Insert lifted edges between every pair of base-graph nodes within + // `max_distance` hops (excluding the source). Each new edge is created + // with weight 0; callers are expected to call `set_cost` afterwards or + // pre-populate weights another way. Existing edges are left untouched. + void insert_lifted_edges_bfs(const std::uint64_t max_distance) { + BfsWorkspace workspace; + for (std::uint64_t source = 0; source < base_graph_.number_of_nodes(); ++source) { + const auto reached = breadth_first_search( + base_graph_, source, max_distance, false, workspace + ); + for (const auto entry : reached) { + if (entry.node > source) { + set_cost(source, entry.node, 0.0, false); + } + } + } + } + + void set_labels(std::vector labels) { + validate_labels(base_graph_, labels); + labels_ = dense_relabel(labels); + } + + void reset_labels() { + labels_ = singleton_labels(base_graph_); + } + + [[nodiscard]] double eval() const { + return energy(lifted_graph_, weights_, labels_); + } + + [[nodiscard]] double eval(const std::vector &labels) const { + return energy(lifted_graph_, weights_, labels); + } + +private: + static UndirectedGraph build_base_only_( + const UndirectedGraph &base_graph, + const std::vector &base_weights + ) { + validate_weights(base_graph, base_weights); + UndirectedGraph lifted_graph( + base_graph.number_of_nodes(), base_graph.number_of_edges() + ); + for (std::uint64_t edge = 0; edge < base_graph.number_of_edges(); ++edge) { + const auto uv = base_graph.uv(edge); + lifted_graph.insert_edge(uv.first, uv.second); + } + return lifted_graph; + } + + const UndirectedGraph &base_graph_; + UndirectedGraph lifted_graph_; + std::vector weights_; + std::vector labels_; + std::uint64_t n_base_edges_; +}; + +class SolverBase { +public: + virtual ~SolverBase() = default; + virtual std::vector optimize(Objective &objective) const = 0; +}; + +} // namespace bioimage_cpp::graph::lifted_multicut diff --git a/pyproject.toml b/pyproject.toml index 40aece5..ecbd508 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,8 @@ classifiers = [ ] [project.optional-dependencies] -test = ["pytest"] +test = ["pytest", "pooch"] +data = ["pooch"] [tool.scikit-build] cmake.version = ">=3.21" diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 35dca71..62653e5 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -1,9 +1,12 @@ #include "graph.hxx" #include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/graph/breadth_first_search.hxx" #include "bioimage_cpp/graph/connected_components.hxx" #include "bioimage_cpp/graph/edge_weighted_watershed.hxx" #include "bioimage_cpp/graph/feature_accumulation.hxx" +#include "bioimage_cpp/graph/lifted_from_affinities.hxx" +#include "bioimage_cpp/graph/lifted_multicut.hxx" #include "bioimage_cpp/graph/multicut.hxx" #include "bioimage_cpp/graph/multicut/fusion_move.hxx" #include "bioimage_cpp/graph/multicut/greedy_additive.hxx" @@ -216,6 +219,30 @@ Graph graph_from_edges(const std::uint64_t number_of_nodes, ConstUInt64Array uvs return graph; } +// Bulk-construct a graph from a pre-deduplicated (u, v) array, bypassing the +// per-edge hash dedup that ``insert_edge`` performs. The caller asserts that +// no (u, v) pair appears twice in ``uvs`` and that ``u != v`` in every row. +// Edges receive ids matching their position in ``uvs``. This is the fast path +// for copying an existing graph (its ``uv_ids()`` are unique by construction). +Graph graph_from_unique_edges(const std::uint64_t number_of_nodes, ConstUInt64Array uvs) { + require_uv_array(uvs, "uvs"); + std::vector edges; + edges.reserve(static_cast(uvs.shape(0))); + const auto *in = uvs.data(); + for (std::size_t index = 0; index < uvs.shape(0); ++index) { + const auto u = in[2 * index]; + const auto v = in[2 * index + 1]; + if (u == v) { + throw std::invalid_argument("self edges are not supported"); + } + if (u >= number_of_nodes || v >= number_of_nodes) { + throw std::out_of_range("edge endpoint exceeds number_of_nodes"); + } + edges.emplace_back(u, v); + } + return Graph::from_sorted_unique_edges(number_of_nodes, std::move(edges)); +} + Graph graph_deserialize(ConstUInt64Array serialization) { if (serialization.ndim() != 1) { throw std::invalid_argument("serialization must be a 1D uint64 array"); @@ -414,6 +441,108 @@ UInt64Array multicut_kernighan_lin( return vector_to_uint64_array(label_vector); } +std::pair graph_breadth_first_search( + const Graph &graph, + const std::uint64_t source, + const std::uint64_t max_distance, + const bool include_source +) { + std::vector entries; + { + nb::gil_scoped_release release; + entries = graph::breadth_first_search(graph, source, max_distance, include_source); + } + auto nodes = make_uint64_array({entries.size()}); + auto distances = make_uint64_array({entries.size()}); + auto *node_data = nodes.data(); + auto *distance_data = distances.data(); + for (std::size_t index = 0; index < entries.size(); ++index) { + node_data[index] = entries[index].node; + distance_data[index] = entries[index].distance; + } + return {nodes, distances}; +} + +double lifted_multicut_energy( + const Graph &lifted_graph, + ConstDoubleArray weights, + ConstUInt64Array labels +) { + const auto weight_vector = + double_array_to_vector(weights, "edge_weights", lifted_graph.number_of_edges()); + const auto label_vector = + uint64_array_to_vector(labels, "labels", lifted_graph.number_of_nodes()); + nb::gil_scoped_release release; + return graph::lifted_multicut::energy(lifted_graph, weight_vector, label_vector); +} + +UInt64Array lifted_multicut_greedy_additive( + const Graph &lifted_graph, + ConstDoubleArray weights, + const std::uint64_t n_base_edges, + const double weight_stop, + const double node_num_stop, + const bool add_noise, + const int seed, + const double sigma +) { + const auto weight_vector = + double_array_to_vector(weights, "edge_weights", lifted_graph.number_of_edges()); + if (n_base_edges > lifted_graph.number_of_edges()) { + throw std::invalid_argument( + "n_base_edges must be <= lifted graph number_of_edges" + ); + } + std::vector labels; + { + nb::gil_scoped_release release; + labels = graph::lifted_multicut::greedy_additive( + lifted_graph, + weight_vector, + n_base_edges, + weight_stop, + node_num_stop, + add_noise, + seed, + sigma + ); + } + return vector_to_uint64_array(labels); +} + +UInt64Array lifted_multicut_kernighan_lin( + const Graph &base_graph, + const Graph &lifted_graph, + ConstDoubleArray weights, + const std::uint64_t n_base_edges, + ConstUInt64Array initial_labels, + const std::uint64_t number_of_outer_iterations, + const double epsilon +) { + const auto weight_vector = + double_array_to_vector(weights, "edge_weights", lifted_graph.number_of_edges()); + auto label_vector = + uint64_array_to_vector(initial_labels, "initial_labels", base_graph.number_of_nodes()); + if (n_base_edges > lifted_graph.number_of_edges()) { + throw std::invalid_argument( + "n_base_edges must be <= lifted graph number_of_edges" + ); + } + { + nb::gil_scoped_release release; + label_vector = graph::lifted_multicut::kernighan_lin( + base_graph, + lifted_graph, + weight_vector, + n_base_edges, + std::move(label_vector), + number_of_outer_iterations, + epsilon + ); + } + return vector_to_uint64_array(label_vector); +} + UInt64Array multicut_fusion_move( const Graph &graph, ConstDoubleArray costs, @@ -579,6 +708,89 @@ DoubleArray accumulate_affinity_features_t( return result; } +template +UInt64Array lifted_edges_from_affinities_t( + const Rag &rag, + LabelArray labels, + const std::vector> &offsets, + const std::size_t number_of_threads +) { + ConstArrayView labels_view{ + labels.data(), + ndarray_shape(labels), + {}, + }; + + std::vector lifted_edges; + { + nb::gil_scoped_release release; + lifted_edges = graph::lifted_edges_from_offsets( + rag, labels_view, offsets, number_of_threads + ); + } + auto result = make_uint64_array({lifted_edges.size(), 2}); + auto *data = result.data(); + for (std::size_t index = 0; index < lifted_edges.size(); ++index) { + data[2 * index] = lifted_edges[index].first; + data[2 * index + 1] = lifted_edges[index].second; + } + return result; +} + +template +DoubleArray accumulate_lifted_affinity_features_t( + LabelArray labels, + ConstDoubleArray affinities, + const std::vector> &offsets, + ConstUInt64Array lifted_uvs, + const bool compute_complex_features, + const std::size_t number_of_threads +) { + if (lifted_uvs.ndim() != 2 || lifted_uvs.shape(1) != 2) { + throw std::invalid_argument("lifted_uvs must have shape (n_lifted, 2)"); + } + const auto n_lifted = lifted_uvs.shape(0); + const auto number_of_features = compute_complex_features ? 12 : 2; + auto result = make_double_array({n_lifted, number_of_features}); + + std::vector lifted_edges(n_lifted); + const auto *uv_data = lifted_uvs.data(); + for (std::size_t index = 0; index < n_lifted; ++index) { + lifted_edges[index] = {uv_data[2 * index], uv_data[2 * index + 1]}; + } + + ConstArrayView labels_view{ + labels.data(), + ndarray_shape(labels), + {}, + }; + ConstArrayView affinities_view{ + affinities.data(), + ndarray_shape(affinities), + {}, + }; + ArrayView out_view{ + result.data(), + { + static_cast(n_lifted), + static_cast(number_of_features), + }, + {}, + }; + + nb::gil_scoped_release release; + graph::accumulate_lifted_affinity_features( + labels_view, + affinities_view, + offsets, + lifted_edges, + compute_complex_features, + number_of_threads, + out_view + ); + return result; +} + template UInt64Array project_node_labels_to_pixels_t( const Rag &rag, @@ -667,6 +879,12 @@ void bind_graph(nb::module_ &m) { nb::arg("number_of_nodes"), nb::arg("uvs") ) + .def_static( + "from_unique_edges", + &graph_from_unique_edges, + nb::arg("number_of_nodes"), + nb::arg("uvs") + ) .def_static( "deserialize", &graph_deserialize, @@ -693,6 +911,14 @@ void bind_graph(nb::module_ &m) { nb::class_(m, "RegionAdjacencyGraph") .def_prop_ro("shape", &Rag::shape); + m.def( + "_breadth_first_search", + &graph_breadth_first_search, + nb::arg("graph"), + nb::arg("source"), + nb::arg("max_distance"), + nb::arg("include_source") + ); m.def("_connected_components", &graph_connected_components, nb::arg("graph")); m.def( "_connected_components_masked", @@ -840,6 +1066,62 @@ void bind_graph(nb::module_ &m) { nb::arg("seed") = 0 ); + m.def( + "_lifted_multicut_energy", + &lifted_multicut_energy, + nb::arg("lifted_graph"), + nb::arg("edge_weights"), + nb::arg("labels") + ); + m.def( + "_lifted_multicut_greedy_additive", + &lifted_multicut_greedy_additive, + nb::arg("lifted_graph"), + nb::arg("edge_weights"), + nb::arg("n_base_edges"), + nb::arg("weight_stop"), + nb::arg("node_num_stop"), + nb::arg("add_noise"), + nb::arg("seed"), + nb::arg("sigma") + ); + m.def( + "_lifted_multicut_kernighan_lin", + &lifted_multicut_kernighan_lin, + nb::arg("base_graph"), + nb::arg("lifted_graph"), + nb::arg("edge_weights"), + nb::arg("n_base_edges"), + nb::arg("initial_labels"), + nb::arg("number_of_outer_iterations"), + nb::arg("epsilon") + ); + + // Lifted multicut sub-solver hierarchy. Same shape as the multicut sub- + // solver bindings — opaque to Python, used by future fusion-move drivers. + nb::class_(m, "_LiftedMulticutSolverBase"); + nb::class_< + graph::lifted_multicut::GreedyAdditiveSolver, + graph::lifted_multicut::SolverBase + >(m, "_GreedyAdditiveLiftedMulticutSubSolver") + .def( + nb::init(), + nb::arg("weight_stop") = 0.0, + nb::arg("node_num_stop") = -1.0, + nb::arg("add_noise") = false, + nb::arg("seed") = 42, + nb::arg("sigma") = 1.0 + ); + nb::class_< + graph::lifted_multicut::KernighanLinSolver, + graph::lifted_multicut::SolverBase + >(m, "_KernighanLinLiftedMulticutSubSolver") + .def( + nb::init(), + nb::arg("number_of_outer_iterations") = 100, + nb::arg("epsilon") = 1.0e-6 + ); + m.def( "_multicut_fusion_move", &multicut_fusion_move, @@ -959,6 +1241,80 @@ void bind_graph(nb::module_ &m) { nb::arg("compute_complex_features"), nb::arg("number_of_threads") ); + m.def( + "_lifted_edges_from_affinities_uint32", + &lifted_edges_from_affinities_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("number_of_threads") + ); + m.def( + "_lifted_edges_from_affinities_uint64", + &lifted_edges_from_affinities_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("number_of_threads") + ); + m.def( + "_lifted_edges_from_affinities_int32", + &lifted_edges_from_affinities_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("number_of_threads") + ); + m.def( + "_lifted_edges_from_affinities_int64", + &lifted_edges_from_affinities_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("number_of_threads") + ); + + m.def( + "_accumulate_lifted_affinity_features_uint32", + &accumulate_lifted_affinity_features_t, + nb::arg("labels"), + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("lifted_uvs"), + nb::arg("compute_complex_features"), + nb::arg("number_of_threads") + ); + m.def( + "_accumulate_lifted_affinity_features_uint64", + &accumulate_lifted_affinity_features_t, + nb::arg("labels"), + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("lifted_uvs"), + nb::arg("compute_complex_features"), + nb::arg("number_of_threads") + ); + m.def( + "_accumulate_lifted_affinity_features_int32", + &accumulate_lifted_affinity_features_t, + nb::arg("labels"), + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("lifted_uvs"), + nb::arg("compute_complex_features"), + nb::arg("number_of_threads") + ); + m.def( + "_accumulate_lifted_affinity_features_int64", + &accumulate_lifted_affinity_features_t, + nb::arg("labels"), + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("lifted_uvs"), + nb::arg("compute_complex_features"), + nb::arg("number_of_threads") + ); + m.def( "_project_node_labels_to_pixels_uint32", &project_node_labels_to_pixels_t, diff --git a/src/bioimage_cpp/_data.py b/src/bioimage_cpp/_data.py new file mode 100644 index 0000000..13734f0 --- /dev/null +++ b/src/bioimage_cpp/_data.py @@ -0,0 +1,132 @@ +"""Pooch-based registry for downloadable problem instances. + +Caches files under ``~/.cache/bioimage-cpp/`` by default, overridable with +``BIOIMAGE_CPP_CACHE``. Pooch is imported lazily so the runtime does not gain +a hard dependency on it; ``fetch`` raises a clear error if pooch is missing. + +Registered files +---------------- + +Multicut problems (text files with rows ``u v cost``; originate from +``elf.segmentation.utils.load_multicut_problem``): + +- ``multicut_problem_A_small.txt`` ... ``multicut_problem_C_medium.txt`` + (3 samples × 2 sizes = 6 problems). + +Lifted multicut problems (``.npz`` files written by +``examples/segmentation/serialize_lifted_problem.py``): + +- ``lifted_multicut_problem_2d.npz`` — 2D ISBI slice (small, ~756 nodes). +- ``lifted_multicut_problem_3d.npz`` — full 3D ISBI volume (large, ~18k nodes). +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Optional + +DEFAULT_CACHE_DIR = Path.home() / ".cache" / "bioimage-cpp" +CACHE_ENV_VAR = "BIOIMAGE_CPP_CACHE" + + +# Each entry is filename -> (url, sha256). To refresh a hash, delete the +# corresponding file under :func:`cache_dir`, re-download it, and run +# ``sha256sum`` on the cached file. +_REGISTRY: dict[str, tuple[str, Optional[str]]] = { + "multicut_problem_A_small.txt": ( + "https://oc.embl.de/index.php/s/yVKwyQ8VoPXYkft/download", + "eeb1083557a20f7ce1ece28f5c613cc8ce5bf6231cd74aadbeb8a5012c6f8ef0", + ), + "multicut_problem_A_medium.txt": ( + "https://oc.embl.de/index.php/s/ztnwjmv0bmd3mnS/download", + "a8cdd23fcd911ad62b1b859b242bac28d16e7cdc3920137116b05672c4a6ec8a", + ), + "multicut_problem_B_small.txt": ( + "https://oc.embl.de/index.php/s/QKYA2EoMXqxQuO4/download", + "abd2c040234f20b107cc237b2c87120058d78e2c5e3ba2b95bc12b3b4d433aa5", + ), + "multicut_problem_B_medium.txt": ( + "https://oc.embl.de/index.php/s/yuk7VwCvgZC017q/download", + "6a8406c774553753e49103531945c32170587cc0d20d0459c866b47de5b014ec", + ), + "multicut_problem_C_small.txt": ( + "https://oc.embl.de/index.php/s/eDZprDwT2cXFAe0/download", + "6db8336c0ba3f75e3f9432628ac13b156fb9e43f75307cdda11469927ed1a108", + ), + "multicut_problem_C_medium.txt": ( + "https://oc.embl.de/index.php/s/hGyqlkenHfsq5P4/download", + "130d1be14d69f8bfb5d20d1375452291db7ba620e2f03bf9ffbe52d1f577f0dc", + ), + "lifted_multicut_problem_2d.npz": ( + "https://owncloud.gwdg.de/index.php/s/QikYgJzbVxD5q8q/download", + "27f10d9b7b2405cf64fab49c9065291455f2f1364224bb94a255c4cc72798240", + ), + "lifted_multicut_problem_3d.npz": ( + "https://owncloud.gwdg.de/index.php/s/ZVzDy8Xb0Dr2Ell/download", + "269ce644e2b9f8259f7f2ff827d5808ac5c9bfe6ca0444e298290f23867dce8a", + ), +} + + +def cache_dir() -> Path: + """Return the cache directory used for downloaded problem instances.""" + override = os.environ.get(CACHE_ENV_VAR) + return Path(override).expanduser() if override else DEFAULT_CACHE_DIR + + +def registered_files() -> list[str]: + """List every filename available via :func:`fetch`.""" + return sorted(_REGISTRY.keys()) + + +def fetch(filename: str, *, timeout: Optional[float] = None) -> Path: + """Return the local path to a registered file, downloading on first call. + + Parameters + ---------- + filename: + Filename as listed by :func:`registered_files`. + timeout: + Optional HTTP timeout in seconds, forwarded to pooch's downloader. + + Raises + ------ + FileNotFoundError + If ``filename`` is not registered. + ModuleNotFoundError + If ``pooch`` is not installed. Install with ``pip install pooch``. + RuntimeError + If the download fails. The underlying ``HTTPError`` is chained. + """ + if filename not in _REGISTRY: + registered = registered_files() + raise FileNotFoundError( + f"{filename!r} is not in the bioimage-cpp data registry. " + f"Available: {registered}" + ) + try: + import pooch + except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "pooch is required to download bioimage-cpp problem instances. " + "Install it with `pip install pooch`." + ) from error + + url, sha256 = _REGISTRY[filename] + fetcher = pooch.create( + path=cache_dir(), + base_url="", + registry={filename: sha256}, + urls={filename: url}, + ) + downloader = None + if timeout is not None: + downloader = pooch.HTTPDownloader(timeout=float(timeout)) + try: + local_path = fetcher.fetch(filename, downloader=downloader) + except Exception as error: + raise RuntimeError( + f"could not download {filename} from {url}: {error}" + ) from error + return Path(local_path) diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index 83496b9..3ebd6b1 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -10,9 +10,15 @@ from ._external import ( DEFAULT_EXTERNAL_MULTICUT_PROBLEM_PATH, EXTERNAL_MULTICUT_PROBLEM_URL, + LiftedMulticutProblem, external_multicut_problem_path, + lifted_multicut_problem_path, load_external_multicut_problem, load_external_multicut_problem_data, + load_lifted_multicut_problem, + load_multicut_problem, + load_multicut_problem_data, + multicut_problem_path, ) _REGION_ADJACENCY_GRAPH_BY_DTYPE = { @@ -36,6 +42,20 @@ np.dtype("int64"): _core._accumulate_affinity_features_int64, } +_LIFTED_EDGES_FROM_AFFINITIES_BY_DTYPE = { + np.dtype("uint32"): _core._lifted_edges_from_affinities_uint32, + np.dtype("uint64"): _core._lifted_edges_from_affinities_uint64, + np.dtype("int32"): _core._lifted_edges_from_affinities_int32, + np.dtype("int64"): _core._lifted_edges_from_affinities_int64, +} + +_LIFTED_AFFINITY_FEATURES_BY_DTYPE = { + np.dtype("uint32"): _core._accumulate_lifted_affinity_features_uint32, + np.dtype("uint64"): _core._accumulate_lifted_affinity_features_uint64, + np.dtype("int32"): _core._accumulate_lifted_affinity_features_int32, + np.dtype("int64"): _core._accumulate_lifted_affinity_features_int64, +} + _EDGE_WEIGHTED_WATERSHED_BY_DTYPE = { (np.dtype("float32"), np.dtype("uint32")): _core._edge_weighted_watershed_float32_uint32, (np.dtype("float32"), np.dtype("uint64")): _core._edge_weighted_watershed_float32_uint64, @@ -145,11 +165,17 @@ def _as_serialization_array(serialization) -> np.ndarray: return np.ascontiguousarray(array) -def _copy_graph(graph: UndirectedGraph | RegionAdjacencyGraph) -> UndirectedGraph: - copied = UndirectedGraph(int(graph.number_of_nodes), int(graph.number_of_edges)) - if graph.number_of_edges: - copied.insert_edges(graph.uv_ids()) - return copied +def _copy_graph(graph: UndirectedGraph | RegionAdjacencyGraph) -> _core.UndirectedGraph: + # `uv_ids()` always returns a unique list (graphs deduplicate on insert), + # so we can use the bulk constructor that skips per-edge hash dedup — + # significantly faster than `insert_edges` for large graphs. The result + # is a ``_core.UndirectedGraph``; downstream code (objectives, solvers, + # validators) uses base-class methods that work identically. + if graph.number_of_edges == 0: + return _core.UndirectedGraph(int(graph.number_of_nodes)) + return _core.UndirectedGraph.from_unique_edges( + int(graph.number_of_nodes), graph.uv_ids() + ) def _as_edge_costs(edge_costs, graph: UndirectedGraph | RegionAdjacencyGraph) -> np.ndarray: @@ -229,6 +255,49 @@ def connected_components( ) +def breadth_first_search( + graph: UndirectedGraph | RegionAdjacencyGraph, + source: int, + *, + max_distance: int | None = None, + include_source: bool = True, +) -> tuple[np.ndarray, np.ndarray]: + """Breadth-first search from ``source`` on ``graph``. + + Returns ``(nodes, distances)`` — two 1D ``uint64`` arrays of equal length, + listing every reachable node within ``max_distance`` hops (inclusive) in + BFS order along with its hop distance from the source. + + Parameters + ---------- + graph: + :class:`UndirectedGraph` or :class:`RegionAdjacencyGraph`. + source: + Source node id. + max_distance: + Maximum hop distance from ``source`` to report. ``None`` (default) + means no limit — the search expands until the entire connected + component of ``source`` is visited. + include_source: + If ``True`` (default), the source itself is reported with distance 0. + Set to ``False`` for "nodes within k hops, excluding self" queries. + """ + if int(source) < 0 or int(source) >= int(graph.number_of_nodes): + raise ValueError( + f"source must be in [0, number_of_nodes), got source={source}, " + f"number_of_nodes={int(graph.number_of_nodes)}" + ) + if max_distance is None: + limit = (1 << 64) - 1 + else: + if int(max_distance) < 0: + raise ValueError("max_distance must be non-negative") + limit = int(max_distance) + return _core._breadth_first_search( + graph, int(source), limit, bool(include_source) + ) + + def edge_weighted_watershed( graph: UndirectedGraph | RegionAdjacencyGraph, edge_weights, @@ -699,6 +768,391 @@ def optimize(self, objective: MulticutObjective) -> np.ndarray: return objective.labels +class LiftedMulticutObjective: + """Lifted multicut objective. + + Stores a base graph + base edge costs together with an internal *lifted + graph* that is a superset of the base graph (base edges occupy ids + ``0 .. base.number_of_edges - 1``; lifted edges follow). The energy of a + node labeling is the sum of base + lifted edge weights across cut edges. + + The lifted edges can be supplied either as explicit ``(uvs, costs)`` + arrays, via a ``bfs_distance=k`` constructor argument that inserts a + zero-weight lifted edge for every pair of nodes within ``k`` hops of each + other in the base graph, or by calling :meth:`set_cost` after construction. + """ + + def __init__( + self, + graph: UndirectedGraph | RegionAdjacencyGraph, + edge_costs, + *, + lifted_uvs=None, + lifted_costs=None, + bfs_distance: int | None = None, + overwrite_existing: bool = False, + initial_labels=None, + ): + # The objective holds a reference to the user's base graph (no + # defensive copy — the C++ ``Objective`` already keeps a const + # reference). The user is expected to treat the input graph as + # immutable while the objective is alive; mutations are visible to + # the objective and may produce undefined behaviour. + base_graph = graph + base_costs = _as_edge_costs(edge_costs, base_graph) + + # Use the bulk constructor for the lifted graph's base portion to + # bypass the per-edge hash dedup that ``insert_edges`` performs. + if int(base_graph.number_of_edges) > 0: + lifted_graph = _core.UndirectedGraph.from_unique_edges( + int(base_graph.number_of_nodes), base_graph.uv_ids() + ) + else: + lifted_graph = _core.UndirectedGraph(int(base_graph.number_of_nodes)) + + weights_list = [base_costs.copy()] + + if bfs_distance is not None: + distance = int(bfs_distance) + if distance < 1: + raise ValueError("bfs_distance must be >= 1") + bfs_uvs = [] + for source in range(int(base_graph.number_of_nodes)): + nodes, _ = breadth_first_search( + base_graph, + source, + max_distance=distance, + include_source=False, + ) + if nodes.size == 0: + continue + tail = nodes[nodes > source] + if tail.size == 0: + continue + source_column = np.full(tail.size, source, dtype=np.uint64) + bfs_uvs.append(np.stack([source_column, tail], axis=1)) + if bfs_uvs: + bfs_uv_array = np.ascontiguousarray(np.concatenate(bfs_uvs, axis=0)) + _add_lifted_edges( + lifted_graph, + weights_list, + bfs_uv_array, + np.zeros(bfs_uv_array.shape[0], dtype=np.float64), + overwrite_existing=overwrite_existing, + ) + + if lifted_uvs is not None or lifted_costs is not None: + if lifted_uvs is None or lifted_costs is None: + raise ValueError( + "lifted_uvs and lifted_costs must be provided together" + ) + uv_array = _as_uv_array(lifted_uvs, "lifted_uvs") + cost_array = np.asarray(lifted_costs, dtype=np.float64) + if cost_array.ndim != 1: + raise ValueError("lifted_costs must be a 1D array") + if cost_array.shape[0] != uv_array.shape[0]: + raise ValueError( + "lifted_uvs and lifted_costs must have the same length, got " + f"lifted_uvs.shape[0]={uv_array.shape[0]}, " + f"lifted_costs.shape[0]={cost_array.shape[0]}" + ) + _add_lifted_edges( + lifted_graph, + weights_list, + uv_array, + np.ascontiguousarray(cost_array), + overwrite_existing=overwrite_existing, + ) + + self._base_graph = base_graph + self._lifted_graph = lifted_graph + self._n_base_edges = int(base_graph.number_of_edges) + self._weights = np.ascontiguousarray(np.concatenate(weights_list)) \ + if len(weights_list) > 1 else weights_list[0] + if initial_labels is None: + self._labels = np.arange(base_graph.number_of_nodes, dtype=np.uint64) + else: + self._labels = _as_node_labels(initial_labels, base_graph) + + @property + def graph(self) -> UndirectedGraph: + return self._base_graph + + @property + def lifted_graph(self) -> UndirectedGraph: + return self._lifted_graph + + @property + def weights(self) -> np.ndarray: + return self._weights + + @property + def number_of_base_edges(self) -> int: + return self._n_base_edges + + @property + def number_of_lifted_edges(self) -> int: + return int(self._lifted_graph.number_of_edges) - self._n_base_edges + + @property + def labels(self) -> np.ndarray: + return self._labels + + @labels.setter + def labels(self, labels) -> None: + self._labels = _as_node_labels(labels, self._base_graph) + + def set_labels(self, labels) -> None: + self.labels = labels + + def reset_labels(self) -> None: + self._labels = np.arange(self._base_graph.number_of_nodes, dtype=np.uint64) + + def set_cost( + self, + u: int, + v: int, + weight: float, + *, + overwrite: bool = False, + ) -> tuple[int, bool]: + """Insert or update a single lifted edge weight. + + Returns ``(edge_id, is_new)`` — the lifted-graph edge id and whether a + new edge was inserted. If the edge already exists (as a base edge or a + previously inserted lifted edge), the weight is accumulated unless + ``overwrite=True``. + """ + pre = int(self._lifted_graph.number_of_edges) + edge = int(self._lifted_graph.insert_edge(int(u), int(v))) + if int(self._lifted_graph.number_of_edges) > pre: + self._weights = np.concatenate( + [self._weights, np.asarray([float(weight)], dtype=np.float64)] + ) + return edge, True + if overwrite: + self._weights[edge] = float(weight) + else: + self._weights[edge] = self._weights[edge] + float(weight) + return edge, False + + def energy(self, labels=None) -> float: + label_array = ( + self._labels if labels is None else _as_node_labels(labels, self._base_graph) + ) + return float( + _core._lifted_multicut_energy(self._lifted_graph, self._weights, label_array) + ) + + +def _add_lifted_edges( + lifted_graph: UndirectedGraph, + weights_list: list[np.ndarray], + lifted_uvs: np.ndarray, + lifted_costs: np.ndarray, + *, + overwrite_existing: bool, +) -> None: + """Insert lifted edges into an UndirectedGraph and update the weights list. + + Mirrors the C++ ``build_lifted_graph`` semantics: brand-new edges append + their cost to ``weights_list``; existing edges (duplicates of a base edge + or of a previously inserted lifted edge) either overwrite or accumulate + their weight in place. ``weights_list`` is expected to hold the current + weights array as its single element on entry; on exit it contains either + one or two ndarray entries (the existing weights plus the new tail). + """ + if lifted_uvs.shape[0] == 0: + return + # In-place updates require a single flat working buffer; coalesce first. + if len(weights_list) > 1: + weights_list[:] = [np.ascontiguousarray(np.concatenate(weights_list))] + + # Fast path: bulk-insert and detect uniqueness from the row count delta. + # For the typical case — ``lifted_uvs`` produced by + # ``lifted_edges_from_affinities`` or by the BFS constructor — every row + # is a brand-new edge, so the delta equals the input length and we can + # append the weights array directly. No ``find_edges`` calls are needed. + if not overwrite_existing: + pre_count = int(lifted_graph.number_of_edges) + lifted_graph.insert_edges(lifted_uvs) + post_count = int(lifted_graph.number_of_edges) + n_new = post_count - pre_count + + if n_new == lifted_uvs.shape[0]: + weights_list.append( + np.ascontiguousarray(lifted_costs.astype(np.float64, copy=False)) + ) + return + + # Some rows collided with existing edges or with each other. Use + # find_edges to recover the per-row edge id (insertion is already + # done; this is just a lookup). + edge_ids = np.asarray(lifted_graph.find_edges(lifted_uvs)) + lifted_costs_f64 = lifted_costs.astype(np.float64, copy=False) + working = weights_list[0] + + collision_mask = edge_ids < pre_count + if collision_mask.any(): + np.add.at( + working, + edge_ids[collision_mask].astype(np.intp, copy=False), + lifted_costs_f64[collision_mask], + ) + + new_mask = ~collision_mask + if new_mask.any(): + slot = (edge_ids[new_mask] - pre_count).astype(np.int64, copy=False) + new_weights = np.bincount( + slot, weights=lifted_costs_f64[new_mask], minlength=n_new + ).astype(np.float64, copy=False) + else: + new_weights = np.zeros(n_new, dtype=np.float64) + + weights_list[0] = working + if n_new > 0: + weights_list.append(new_weights) + return + + # Slow path: per-row Python loop for ``overwrite_existing=True`` (rare). + # Order-sensitive last-write-wins semantics on collisions. + working = weights_list[0] + new_costs: list[float] = [] + for index in range(lifted_uvs.shape[0]): + u = int(lifted_uvs[index, 0]) + v = int(lifted_uvs[index, 1]) + weight = float(lifted_costs[index]) + pre = int(lifted_graph.number_of_edges) + edge = int(lifted_graph.insert_edge(u, v)) + if int(lifted_graph.number_of_edges) > pre: + new_costs.append(weight) + else: + working[edge] = weight + if new_costs: + weights_list[0] = working + weights_list.append(np.asarray(new_costs, dtype=np.float64)) + else: + weights_list[0] = working + + +class LiftedMulticutSolver(ABC): + """Base class for lifted multicut solvers.""" + + @abstractmethod + def optimize(self, objective: LiftedMulticutObjective) -> np.ndarray: + """Optimize ``objective`` and return the node labeling.""" + + +class LiftedGreedyAdditiveMulticut(LiftedMulticutSolver): + def __init__( + self, + *, + weight_stop: float = 0.0, + node_num_stop: float = -1.0, + add_noise: bool = False, + seed: int = 42, + sigma: float = 1.0, + ): + self.weight_stop = float(weight_stop) + self.node_num_stop = float(node_num_stop) + self.add_noise = bool(add_noise) + self.seed = int(seed) + self.sigma = float(sigma) + + def optimize(self, objective: LiftedMulticutObjective) -> np.ndarray: + labels = _core._lifted_multicut_greedy_additive( + objective.lifted_graph, + objective.weights, + objective.number_of_base_edges, + self.weight_stop, + self.node_num_stop, + self.add_noise, + self.seed, + self.sigma, + ) + objective.labels = labels + return objective.labels + + def _build_cpp_sub_solver(self): + return _core._GreedyAdditiveLiftedMulticutSubSolver( + weight_stop=self.weight_stop, + node_num_stop=self.node_num_stop, + add_noise=self.add_noise, + seed=self.seed, + sigma=self.sigma, + ) + + +class LiftedKernighanLinMulticut(LiftedMulticutSolver): + def __init__( + self, + *, + number_of_outer_iterations: int = 100, + epsilon: float = 1.0e-6, + ): + self.number_of_outer_iterations = int(number_of_outer_iterations) + if self.number_of_outer_iterations < 0: + raise ValueError("number_of_outer_iterations must be non-negative") + self.epsilon = float(epsilon) + + def optimize(self, objective: LiftedMulticutObjective) -> np.ndarray: + initial_labels = objective.labels + if np.array_equal( + initial_labels, + np.arange(objective.graph.number_of_nodes, dtype=np.uint64), + ): + initial_labels = _core._lifted_multicut_greedy_additive( + objective.lifted_graph, + objective.weights, + objective.number_of_base_edges, + 0.0, + -1.0, + False, + 42, + 1.0, + ) + labels = _core._lifted_multicut_kernighan_lin( + objective.graph, + objective.lifted_graph, + objective.weights, + objective.number_of_base_edges, + initial_labels, + self.number_of_outer_iterations, + self.epsilon, + ) + objective.labels = labels + return objective.labels + + def _build_cpp_sub_solver(self): + return _core._KernighanLinLiftedMulticutSubSolver( + number_of_outer_iterations=self.number_of_outer_iterations, + epsilon=self.epsilon, + ) + + +class LiftedChainedSolvers(LiftedMulticutSolver): + """Chain of lifted multicut solvers run in sequence on the same objective. + + Each solver's output labeling is fed to the next via the shared + :class:`LiftedMulticutObjective`. Typical use: ``[LiftedGreedyAdditiveMulticut(), + LiftedKernighanLinMulticut(...)]`` for a fast warm-start followed by a + local refinement. + """ + + def __init__(self, solvers): + self.solvers = list(solvers) + if len(self.solvers) == 0: + raise ValueError("solvers must contain at least one solver") + if not all(isinstance(solver, LiftedMulticutSolver) for solver in self.solvers): + raise TypeError("all solvers must inherit from LiftedMulticutSolver") + + def optimize(self, objective: LiftedMulticutObjective) -> np.ndarray: + labels = objective.labels + for solver in self.solvers: + labels = solver.optimize(objective) + return labels + + def region_adjacency_graph( labels: np.ndarray, *, @@ -812,6 +1266,112 @@ def affinity_features_complex( ) +def lifted_edges_from_affinities( + rag: RegionAdjacencyGraph, + labels: np.ndarray, + offsets, + *, + number_of_threads: int = 0, +) -> np.ndarray: + """Discover lifted edges implied by long-range affinity offsets. + + Walks every grid coordinate together with each long-range offset (1-hop + offsets are skipped automatically). When the labels at ``(p, p + offset)`` + differ and ``(labels[p], labels[p + offset])`` is not already a local + edge of ``rag``, the pair is recorded as a lifted edge. + + Parameters + ---------- + rag: + :class:`RegionAdjacencyGraph` built from ``labels``. + labels: + 2D or 3D label array. Supported dtypes: ``uint32``, ``uint64``, + ``int32``, ``int64``. + offsets: + Sequence of per-channel offsets. Each offset must have length equal + to ``labels.ndim``. Offsets with L1 norm ``<= 1`` are skipped, so + callers can pass the full offset list of an affinity volume without + pre-filtering. + + Returns + ------- + np.ndarray + ``(n_lifted, 2)`` ``uint64`` array of ``(u, v)`` pairs with + ``u < v``, sorted lexicographically. + """ + label_array = _normalize_labels(labels) + if tuple(int(size) for size in rag.shape) != label_array.shape: + raise ValueError( + "rag shape must match labels shape, got " + f"rag shape={tuple(rag.shape)}, labels shape={label_array.shape}" + ) + + normalized_offsets = [tuple(int(value) for value in offset) for offset in offsets] + if any(len(offset) != label_array.ndim for offset in normalized_offsets): + raise ValueError("each offset must have length matching labels ndim") + + run = _LIFTED_EDGES_FROM_AFFINITIES_BY_DTYPE[label_array.dtype] + return run( + rag, + label_array, + normalized_offsets, + _normalize_number_of_threads(number_of_threads), + ) + + +def lifted_affinity_features( + labels: np.ndarray, + affinities: np.ndarray, + offsets, + lifted_uvs, + *, + number_of_threads: int = 0, +) -> np.ndarray: + """Compute mean and size features for affinity links across lifted edges. + + Affinity values at pixel pairs ``(p, p + offset)`` whose labels match a + row of ``lifted_uvs`` are binned into that lifted edge. Pixel pairs that + fall on a non-lifted edge (or no edge at all) are silently skipped, so + a local edge that is also reachable by a long-range offset is not + contaminated by long-range affinities. + + 1-hop offsets are skipped automatically. + + The returned array has shape ``(len(lifted_uvs), 2)`` with columns + ``SIMPLE_EDGE_FEATURE_NAMES`` (``mean``, ``size``). + """ + return _accumulate_lifted_affinity_features( + labels, + affinities, + offsets, + lifted_uvs, + compute_complex_features=False, + number_of_threads=number_of_threads, + ) + + +def lifted_affinity_features_complex( + labels: np.ndarray, + affinities: np.ndarray, + offsets, + lifted_uvs, + *, + number_of_threads: int = 0, +) -> np.ndarray: + """Complex affinity features for links across lifted edges. + + Output columns: ``COMPLEX_EDGE_FEATURE_NAMES``. + """ + return _accumulate_lifted_affinity_features( + labels, + affinities, + offsets, + lifted_uvs, + compute_complex_features=True, + number_of_threads=number_of_threads, + ) + + def project_node_labels_to_pixels( rag: RegionAdjacencyGraph, labels: np.ndarray, @@ -912,6 +1472,47 @@ def _accumulate_affinity_features( ) +def _accumulate_lifted_affinity_features( + labels: np.ndarray, + affinities: np.ndarray, + offsets, + lifted_uvs, + *, + compute_complex_features: bool, + number_of_threads: int, +) -> np.ndarray: + label_array = _normalize_labels(labels) + affinity_array = np.asarray(affinities, dtype=np.float64) + if affinity_array.ndim != label_array.ndim + 1: + raise ValueError("affinities must have shape (channels, *labels.shape)") + if affinity_array.shape[1:] != label_array.shape: + raise ValueError( + "affinities spatial shape must match labels shape, got " + f"affinities shape={affinity_array.shape}, labels shape={label_array.shape}" + ) + + normalized_offsets = [tuple(int(value) for value in offset) for offset in offsets] + if len(normalized_offsets) != affinity_array.shape[0]: + raise ValueError( + "offsets length must match affinities channel count, got " + f"offsets length={len(normalized_offsets)}, channels={affinity_array.shape[0]}" + ) + if any(len(offset) != label_array.ndim for offset in normalized_offsets): + raise ValueError("each offset must have length matching labels ndim") + + lifted_uv_array = _as_uv_array(lifted_uvs, "lifted_uvs") + + run = _LIFTED_AFFINITY_FEATURES_BY_DTYPE[label_array.dtype] + return run( + label_array, + np.ascontiguousarray(affinity_array), + normalized_offsets, + lifted_uv_array, + bool(compute_complex_features), + _normalize_number_of_threads(number_of_threads), + ) + + def _normalize_labels(labels: np.ndarray) -> np.ndarray: array = np.asarray(labels) if array.ndim not in (2, 3): @@ -945,6 +1546,12 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "GreedyAdditiveProposalGenerator", "GreedyFixationMulticut", "KernighanLinMulticut", + "LiftedChainedSolvers", + "LiftedGreedyAdditiveMulticut", + "LiftedKernighanLinMulticut", + "LiftedMulticutObjective", + "LiftedMulticutProblem", + "LiftedMulticutSolver", "MulticutDecomposer", "MulticutObjective", "MulticutSolver", @@ -955,13 +1562,22 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "WatershedProposalGenerator", "affinity_features", "affinity_features_complex", + "breadth_first_search", "connected_components", "edge_map_features", "edge_map_features_complex", "edge_weighted_watershed", + "lifted_affinity_features", + "lifted_affinity_features_complex", + "lifted_edges_from_affinities", "external_multicut_problem_path", + "lifted_multicut_problem_path", "load_external_multicut_problem", "load_external_multicut_problem_data", + "load_lifted_multicut_problem", + "load_multicut_problem", + "load_multicut_problem_data", + "multicut_problem_path", "project_node_labels_to_pixels", "region_adjacency_graph", "undirected_graph", diff --git a/src/bioimage_cpp/graph/_external.py b/src/bioimage_cpp/graph/_external.py index 5401281..5846d5a 100644 --- a/src/bioimage_cpp/graph/_external.py +++ b/src/bioimage_cpp/graph/_external.py @@ -1,17 +1,162 @@ -"""External graph problems used for development and regression checks.""" +"""External graph problems used for development and regression checks. + +Downloads go through the shared pooch registry in :mod:`bioimage_cpp._data` +(cached under ``~/.cache/bioimage-cpp/``). The module exposes: + +- :func:`load_multicut_problem` / :func:`load_multicut_problem_data` / + :func:`multicut_problem_path` for the 6 multicut problems (3 samples × 2 + sizes from ``elf.segmentation.utils.load_multicut_problem``). +- :func:`load_lifted_multicut_problem` / + :func:`lifted_multicut_problem_path` for the 2D and 3D lifted multicut + problems built by ``examples/segmentation/serialize_lifted_problem.py``. + +A legacy compatibility layer (``load_external_multicut_problem`` and friends) +delegates to sample A, size small, which is the problem the regression test +uses. +""" from __future__ import annotations import os -import urllib.request +from dataclasses import dataclass from pathlib import Path -from urllib.error import URLError +from typing import Optional import numpy as np +from bioimage_cpp._data import cache_dir, fetch + +VALID_SAMPLES = ("A", "B", "C") +VALID_SIZES = ("small", "medium") +VALID_LIFTED_SIZES = ("2d", "3d") + + +# Legacy constants. The URL points at the canonical sample-A-small download +# used since the first public version. The default path now resolves to the +# new shared cache location instead of /tmp. EXTERNAL_MULTICUT_PROBLEM_URL = "https://oc.embl.de/index.php/s/yVKwyQ8VoPXYkft/download" -DEFAULT_EXTERNAL_MULTICUT_PROBLEM_PATH = Path("/tmp/bioimage_cpp_external_multicut_problem.txt") +DEFAULT_EXTERNAL_MULTICUT_PROBLEM_PATH = cache_dir() / "multicut_problem_A_small.txt" + + +@dataclass +class LiftedMulticutProblem: + """Raw arrays describing a lifted multicut problem.""" + + n_nodes: int + local_uvs: np.ndarray + local_costs: np.ndarray + lifted_uvs: np.ndarray + lifted_costs: np.ndarray + + +# Multicut helpers -------------------------------------------------------- + + +def _validate_sample_size(sample: str, size: str) -> None: + if sample not in VALID_SAMPLES: + raise ValueError( + f"sample must be one of {VALID_SAMPLES}, got {sample!r}" + ) + if size not in VALID_SIZES: + raise ValueError(f"size must be one of {VALID_SIZES}, got {size!r}") + + +def multicut_problem_path( + sample: str = "A", + size: str = "small", + *, + timeout: Optional[float] = None, +) -> Path: + """Return the cached path to a multicut problem, downloading if needed.""" + _validate_sample_size(sample, size) + return fetch(f"multicut_problem_{sample}_{size}.txt", timeout=timeout) + + +def load_multicut_problem_data( + sample: str = "A", + size: str = "small", + *, + timeout: Optional[float] = None, +) -> tuple[np.ndarray, np.ndarray]: + """Load ``(uv_ids, costs)`` for one of the multicut problems.""" + path = multicut_problem_path(sample, size, timeout=timeout) + problem = np.genfromtxt(path) + uv_ids = np.ascontiguousarray(problem[:, :2].astype(np.uint64, copy=False)) + costs = np.ascontiguousarray(problem[:, -1].astype(np.float64, copy=False)) + return uv_ids, costs + + +def load_multicut_problem( + sample: str = "A", + size: str = "small", + *, + timeout: Optional[float] = None, +): + """Load a multicut problem as a bioimage-cpp graph and edge-cost vector.""" + from . import UndirectedGraph + + uv_ids, costs = load_multicut_problem_data(sample, size, timeout=timeout) + graph = UndirectedGraph.from_edges(int(uv_ids.max()) + 1, uv_ids) + return graph, costs + + +# Lifted multicut helpers ------------------------------------------------ + + +def _validate_lifted_size(size: str) -> None: + if size not in VALID_LIFTED_SIZES: + raise ValueError( + f"size must be one of {VALID_LIFTED_SIZES}, got {size!r}" + ) + + +def lifted_multicut_problem_path( + size: str = "2d", + *, + timeout: Optional[float] = None, +) -> Path: + """Return the cached path to the lifted multicut problem file.""" + _validate_lifted_size(size) + return fetch(f"lifted_multicut_problem_{size}.npz", timeout=timeout) + + +def load_lifted_multicut_problem( + size: str = "2d", + *, + timeout: Optional[float] = None, +) -> LiftedMulticutProblem: + """Load a lifted multicut problem as :class:`LiftedMulticutProblem`. + + Parameters + ---------- + size: + ``"2d"`` for the small ISBI 2D slice (~756 nodes, fast, used by the + regression test) or ``"3d"`` for the full ISBI volume (~18k nodes, + used by the development comparison scripts). + timeout: + Optional HTTP timeout in seconds for the download. + """ + path = lifted_multicut_problem_path(size, timeout=timeout) + data = np.load(path) + return LiftedMulticutProblem( + n_nodes=int(data["n_nodes"]), + local_uvs=np.ascontiguousarray( + data["local_uvs"].astype(np.uint64, copy=False) + ), + local_costs=np.ascontiguousarray( + data["local_costs"].astype(np.float64, copy=False) + ), + lifted_uvs=np.ascontiguousarray( + data["lifted_uvs"].astype(np.uint64, copy=False) + ), + lifted_costs=np.ascontiguousarray( + data["lifted_costs"].astype(np.float64, copy=False) + ), + ) + + +# Legacy shims (defaults to sample A small) ------------------------------ def external_multicut_problem_path( @@ -20,38 +165,41 @@ def external_multicut_problem_path( download: bool = True, timeout: float = 30.0, ) -> Path: - """Return the local path for the external multicut problem. + """Return the local path for the default multicut problem (A, small). - The path can be supplied explicitly, via - ``BIOIMAGE_CPP_EXTERNAL_MULTICUT_PATH``, or via the cache path - ``BIOIMAGE_CPP_EXTERNAL_MULTICUT_CACHE``. If no existing file is found and - ``download`` is true, the problem is downloaded into the cache path. + Honors ``BIOIMAGE_CPP_EXTERNAL_MULTICUT_PATH`` (explicit existing file) + and ``BIOIMAGE_CPP_EXTERNAL_MULTICUT_CACHE`` (explicit cache path) for + backwards compatibility; otherwise routes through the shared pooch + cache via :func:`multicut_problem_path`. """ explicit_path = path or os.environ.get("BIOIMAGE_CPP_EXTERNAL_MULTICUT_PATH") if explicit_path is not None: resolved = Path(explicit_path) if not resolved.exists(): - raise FileNotFoundError(f"external multicut problem does not exist: {resolved}") + raise FileNotFoundError( + f"external multicut problem does not exist: {resolved}" + ) return resolved - cache_path = Path( - os.environ.get( - "BIOIMAGE_CPP_EXTERNAL_MULTICUT_CACHE", - str(DEFAULT_EXTERNAL_MULTICUT_PROBLEM_PATH), - ) - ) - if cache_path.exists(): - return cache_path + legacy_cache = os.environ.get("BIOIMAGE_CPP_EXTERNAL_MULTICUT_CACHE") + if legacy_cache is not None: + legacy_path = Path(legacy_cache) + if legacy_path.exists(): + return legacy_path + if not download: + raise FileNotFoundError( + f"external multicut problem does not exist: {legacy_path}" + ) + if not download: - raise FileNotFoundError(f"external multicut problem does not exist: {cache_path}") + candidate = DEFAULT_EXTERNAL_MULTICUT_PROBLEM_PATH + if not candidate.exists(): + raise FileNotFoundError( + f"external multicut problem does not exist: {candidate}" + ) + return candidate - cache_path.parent.mkdir(parents=True, exist_ok=True) - try: - with urllib.request.urlopen(EXTERNAL_MULTICUT_PROBLEM_URL, timeout=timeout) as response: - cache_path.write_bytes(response.read()) - except URLError as error: - raise RuntimeError(f"could not download external multicut problem: {error}") from error - return cache_path + return multicut_problem_path(timeout=timeout) def load_external_multicut_problem_data( @@ -60,11 +208,9 @@ def load_external_multicut_problem_data( download: bool = True, timeout: float = 30.0, ) -> tuple[np.ndarray, np.ndarray]: - """Load external multicut problem edge ids and costs.""" + """Load default multicut problem edge ids and costs (sample A, small).""" problem_path = external_multicut_problem_path( - path, - download=download, - timeout=timeout, + path, download=download, timeout=timeout ) problem = np.genfromtxt(problem_path) uv_ids = np.ascontiguousarray(problem[:, :2].astype(np.uint64, copy=False)) @@ -78,13 +224,11 @@ def load_external_multicut_problem( download: bool = True, timeout: float = 30.0, ): - """Load the external multicut problem as a bioimage-cpp graph and costs.""" + """Load the default multicut problem as a graph + cost vector.""" from . import UndirectedGraph uv_ids, costs = load_external_multicut_problem_data( - path, - download=download, - timeout=timeout, + path, download=download, timeout=timeout ) graph = UndirectedGraph.from_edges(int(uv_ids.max()) + 1, uv_ids) return graph, costs diff --git a/tests/graph/lifted_multicut/__init__.py b/tests/graph/lifted_multicut/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/graph/lifted_multicut/_helpers.py b/tests/graph/lifted_multicut/_helpers.py new file mode 100644 index 0000000..dee34ef --- /dev/null +++ b/tests/graph/lifted_multicut/_helpers.py @@ -0,0 +1,11 @@ +import numpy as np + + +def same_partition(labels, expected): + labels = np.asarray(labels) + expected = np.asarray(expected) + assert labels.shape == expected.shape + np.testing.assert_array_equal( + labels[:, None] == labels[None, :], + expected[:, None] == expected[None, :], + ) diff --git a/tests/graph/lifted_multicut/conftest.py b/tests/graph/lifted_multicut/conftest.py new file mode 100644 index 0000000..17063c7 --- /dev/null +++ b/tests/graph/lifted_multicut/conftest.py @@ -0,0 +1,52 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +@pytest.fixture +def chain_with_lifted(): + """4-node chain 0-1-2-3 with attractive base edges and a single + repulsive lifted edge across (0, 3). The repulsive lifted edge should + cause the chain to be split somewhere along the way.""" + base = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [1, 2], [2, 3]]) + base_costs = np.array([2.0, 2.0, 2.0], dtype=np.float64) + lifted_uvs = np.array([[0, 3]], dtype=np.uint64) + lifted_costs = np.array([-10.0], dtype=np.float64) + return base, base_costs, lifted_uvs, lifted_costs + + +@pytest.fixture +def disjoint_clusters_with_attractive_lifted(): + """Two two-node base components (0-1) and (2-3) with attractive base edges. + A single strongly attractive lifted edge across (0, 2) should NOT cause + the two components to be reported as one cluster — the solvers keep every + cluster base-graph connected.""" + base = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [2, 3]]) + base_costs = np.array([1.0, 1.0], dtype=np.float64) + lifted_uvs = np.array([[0, 2]], dtype=np.uint64) + lifted_costs = np.array([20.0], dtype=np.float64) + return base, base_costs, lifted_uvs, lifted_costs + + +@pytest.fixture +def triangle_with_lifted(): + """3-node triangle. Base edges all attractive; one lifted edge (0, 1) doubles + up over the base (0, 1) edge — exercising the dedup-and-accumulate path.""" + base = bic.graph.UndirectedGraph.from_edges(3, [[0, 1], [1, 2], [0, 2]]) + base_costs = np.array([1.0, 1.0, 1.0], dtype=np.float64) + lifted_uvs = np.array([[0, 1]], dtype=np.uint64) + lifted_costs = np.array([5.0], dtype=np.float64) + return base, base_costs, lifted_uvs, lifted_costs + + +@pytest.fixture +def small_chain_bfs_problem(): + """Long chain where BFS-based lifted-edge construction yields a known + edge count: 5 nodes, chain edges, bfs_distance=2 should add lifted edges + for every node-pair within 2 hops that isn't already a base edge.""" + base = bic.graph.UndirectedGraph.from_edges( + 5, [[0, 1], [1, 2], [2, 3], [3, 4]] + ) + base_costs = np.ones(4, dtype=np.float64) + return base, base_costs diff --git a/tests/graph/lifted_multicut/test_external_problem.py b/tests/graph/lifted_multicut/test_external_problem.py new file mode 100644 index 0000000..0a2603a --- /dev/null +++ b/tests/graph/lifted_multicut/test_external_problem.py @@ -0,0 +1,51 @@ +import pytest + +import bioimage_cpp as bic + + +# Energy bound for the 2D ISBI lifted multicut problem. Every shipped solver +# (including the chained greedy + KL combination) should reach an energy at +# most this high. Baseline measurements at the time of writing: +# singleton labelling : -1263.9 +# LiftedGreedyAdditive : -1575.0 +# LiftedKernighanLin (10) : -1575.2 +# chained (greedy + KL) : -1575.2 +# Margin of ~0.5 catches meaningful regressions while tolerating numerical drift. +ENERGY_BOUND = -1574.5 + + +def _load_problem(): + try: + return bic.graph.load_lifted_multicut_problem("2d", timeout=30) + except (FileNotFoundError, ModuleNotFoundError, RuntimeError) as error: + pytest.skip(str(error)) + + +def test_solvers_on_2d_lifted_problem(): + problem = _load_problem() + graph = bic.graph.UndirectedGraph.from_edges(problem.n_nodes, problem.local_uvs) + + solvers = [ + bic.graph.LiftedGreedyAdditiveMulticut(), + bic.graph.LiftedKernighanLinMulticut(number_of_outer_iterations=10), + bic.graph.LiftedChainedSolvers( + [ + bic.graph.LiftedGreedyAdditiveMulticut(), + bic.graph.LiftedKernighanLinMulticut(number_of_outer_iterations=10), + ] + ), + ] + + for solver in solvers: + objective = bic.graph.LiftedMulticutObjective( + graph, + problem.local_costs, + lifted_uvs=problem.lifted_uvs, + lifted_costs=problem.lifted_costs, + ) + labels = solver.optimize(objective) + assert labels.shape == (graph.number_of_nodes,) + assert objective.energy(labels) <= ENERGY_BOUND, ( + f"{type(solver).__name__} energy {objective.energy(labels):.3f} " + f"exceeds regression bound {ENERGY_BOUND}" + ) diff --git a/tests/graph/lifted_multicut/test_greedy_additive.py b/tests/graph/lifted_multicut/test_greedy_additive.py new file mode 100644 index 0000000..ac5a881 --- /dev/null +++ b/tests/graph/lifted_multicut/test_greedy_additive.py @@ -0,0 +1,88 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + +from ._helpers import same_partition + + +def test_greedy_additive_stops_before_lifted_repulsive_dominates(chain_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + + labels = bic.graph.LiftedGreedyAdditiveMulticut().optimize(objective) + # As base edges get contracted, the lifted -10 weight is folded into the + # remaining contracted edge to the unmerged endpoint. Once the summed + # weight drops below `weight_stop=0` greedy stops, leaving the chain + # split — which is the optimal labeling (energy -8 vs 0 for the fully + # merged labeling). + assert labels[0] != labels[3] + assert objective.energy(labels) == pytest.approx(-8.0) + np.testing.assert_array_equal(objective.labels, labels) + + +def test_greedy_additive_keeps_base_disconnected_clusters( + disjoint_clusters_with_attractive_lifted, +): + base, base_costs, lifted_uvs, lifted_costs = disjoint_clusters_with_attractive_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + + labels = bic.graph.LiftedGreedyAdditiveMulticut().optimize(objective) + # Lifted (0, 2) is attractive but there's no base path to merge across, + # so greedy-additive must keep the two base components separate. + assert labels[0] == labels[1] + assert labels[2] == labels[3] + assert labels[0] != labels[2] + + +def test_greedy_additive_respects_weight_stop(chain_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + + labels = bic.graph.LiftedGreedyAdditiveMulticut(weight_stop=3.0).optimize(objective) + # Every base edge has weight 2.0, below the stop threshold, so no contractions. + same_partition(labels, [0, 1, 2, 3]) + + +def test_greedy_additive_matches_multicut_when_no_lifted_edges(): + # With zero lifted edges, the lifted greedy-additive on the lifted graph + # should yield the same labeling as plain multicut greedy-additive on the + # base graph. + base = bic.graph.UndirectedGraph.from_edges( + 6, + [ + [0, 1], [0, 3], [1, 2], [1, 4], [2, 5], [3, 4], [4, 5], + ], + ) + base_costs = np.array([5, -20, 5, 5, -20, 5, 5], dtype=np.float64) + + mc_objective = bic.graph.MulticutObjective(base, base_costs) + mc_labels = bic.graph.GreedyAdditiveMulticut().optimize(mc_objective) + + lmc_objective = bic.graph.LiftedMulticutObjective(base, base_costs) + lmc_labels = bic.graph.LiftedGreedyAdditiveMulticut().optimize(lmc_objective) + + same_partition(lmc_labels, mc_labels) + + +def test_greedy_additive_lifted_attractive_merges_through_base_path(): + # Chain 0-1-2 with a strong attractive lifted edge (0, 2). + # Even though greedy can't contract a lifted edge, contracting either base + # edge folds the lifted weight into the remaining contracted edge, which + # then becomes a normal heap candidate. + base = bic.graph.UndirectedGraph.from_edges(3, [[0, 1], [1, 2]]) + base_costs = np.array([1.0, 1.0], dtype=np.float64) + lifted_uvs = np.array([[0, 2]], dtype=np.uint64) + lifted_costs = np.array([10.0], dtype=np.float64) + + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + labels = bic.graph.LiftedGreedyAdditiveMulticut().optimize(objective) + same_partition(labels, [0, 0, 0]) diff --git a/tests/graph/lifted_multicut/test_kernighan_lin.py b/tests/graph/lifted_multicut/test_kernighan_lin.py new file mode 100644 index 0000000..bdf8a65 --- /dev/null +++ b/tests/graph/lifted_multicut/test_kernighan_lin.py @@ -0,0 +1,99 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + +from ._helpers import same_partition + + +def test_kl_splits_chain_along_repulsive_lifted_edge(chain_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + + # Warm-starting from a single cluster (everything merged) and running KL + # should split the chain so that the repulsive (0, 3) lifted edge is cut. + objective.labels = np.zeros(4, dtype=np.uint64) + labels = bic.graph.LiftedKernighanLinMulticut( + number_of_outer_iterations=10 + ).optimize(objective) + + # Energy should improve over the single-cluster labeling. + assert objective.energy(labels) < objective.energy(np.zeros(4, dtype=np.uint64)) + # The lifted edge endpoints must end up in different clusters. + assert labels[0] != labels[3] + + +def test_kl_keeps_base_disconnected_clusters_separate( + disjoint_clusters_with_attractive_lifted, +): + base, base_costs, lifted_uvs, lifted_costs = disjoint_clusters_with_attractive_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + + labels = bic.graph.LiftedKernighanLinMulticut( + number_of_outer_iterations=10 + ).optimize(objective) + # Lifted (0, 2) is attractive but the base graph offers no path between + # the two components; lifted KL must keep every cluster base-graph + # connected. + assert labels[0] == labels[1] + assert labels[2] == labels[3] + assert labels[0] != labels[2] + + +def test_kl_matches_multicut_without_lifted_edges(): + base = bic.graph.UndirectedGraph.from_edges( + 6, + [ + [0, 1], [0, 3], [1, 2], [1, 4], [2, 5], [3, 4], [4, 5], + ], + ) + base_costs = np.array([5, -20, 5, 5, -20, 5, 5], dtype=np.float64) + + mc_objective = bic.graph.MulticutObjective(base, base_costs) + mc_labels = bic.graph.KernighanLinMulticut( + number_of_outer_iterations=10 + ).optimize(mc_objective) + + lmc_objective = bic.graph.LiftedMulticutObjective(base, base_costs) + lmc_labels = bic.graph.LiftedKernighanLinMulticut( + number_of_outer_iterations=10 + ).optimize(lmc_objective) + + same_partition(lmc_labels, mc_labels) + + +def test_kl_lifted_attractive_overrides_repulsive_base(): + # Chain 0-1-2 with one weakly attractive base edge and one strongly + # repulsive base edge; an attractive lifted (0, 2) tips the balance back + # toward keeping everything in one cluster. KL warm-starts from the + # greedy result {0,1}|{2} (which only sees the base graph) and must + # discover the merge once the lifted edge is considered. + base = bic.graph.UndirectedGraph.from_edges(3, [[0, 1], [1, 2]]) + base_costs = np.array([1.0, -5.0], dtype=np.float64) + lifted_uvs = np.array([[0, 2]], dtype=np.uint64) + lifted_costs = np.array([10.0], dtype=np.float64) + + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + labels = bic.graph.LiftedKernighanLinMulticut( + number_of_outer_iterations=10 + ).optimize(objective) + + same_partition(labels, [0, 0, 0]) + assert objective.energy(labels) == pytest.approx(0.0) + + +def test_kl_warm_starts_from_singleton(): + base = bic.graph.UndirectedGraph.from_edges(3, [[0, 1], [1, 2]]) + base_costs = np.array([2.0, 2.0], dtype=np.float64) + objective = bic.graph.LiftedMulticutObjective(base, base_costs) + # Singleton labels trigger an internal greedy-additive warm start. + labels = bic.graph.LiftedKernighanLinMulticut( + number_of_outer_iterations=5 + ).optimize(objective) + same_partition(labels, [0, 0, 0]) diff --git a/tests/graph/lifted_multicut/test_objective.py b/tests/graph/lifted_multicut/test_objective.py new file mode 100644 index 0000000..47da69a --- /dev/null +++ b/tests/graph/lifted_multicut/test_objective.py @@ -0,0 +1,157 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_objective_no_lifted_edges_matches_base(chain_with_lifted): + base, base_costs, _, _ = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective(base, base_costs) + + # Default labels (singleton): every base edge is cut. + assert objective.number_of_base_edges == int(base.number_of_edges) + assert objective.number_of_lifted_edges == 0 + assert objective.energy() == pytest.approx(float(base_costs.sum())) + + # Same labeling under a multicut objective should yield the same energy. + mc_objective = bic.graph.MulticutObjective(base, base_costs) + assert objective.energy() == pytest.approx(mc_objective.energy()) + + +def test_objective_with_lifted_edges(chain_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + assert objective.number_of_base_edges == 3 + assert objective.number_of_lifted_edges == 1 + + # All singleton labels: every edge is cut -> 2+2+2+(-10) = -4 + assert objective.energy() == pytest.approx(-4.0) + + # Two-cluster split {0,1,2} | {3}: cuts 2-3 base edge (cost 2) + # and 0-3 lifted edge (cost -10). + assert objective.energy([0, 0, 0, 1]) == pytest.approx(-8.0) + + # Three-cluster {0,1}|{2}|{3}: cuts 1-2 base (2), 2-3 base (2), + # 0-3 lifted (-10) = -6. + assert objective.energy([0, 0, 1, 2]) == pytest.approx(-6.0) + + +def test_objective_lifted_edge_over_base_accumulates(triangle_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = triangle_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + + # The lifted (0, 1) overlaps the base (0, 1): no new edge in the lifted + # graph, but the base edge weight should now be 1 + 5 = 6. + assert objective.number_of_lifted_edges == 0 + # Cut everything (singleton labels): 6 + 1 + 1 = 8 + assert objective.energy() == pytest.approx(8.0) + + +def test_objective_lifted_edge_over_base_overwrite(triangle_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = triangle_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, + base_costs, + lifted_uvs=lifted_uvs, + lifted_costs=lifted_costs, + overwrite_existing=True, + ) + + # With overwrite, the (0, 1) base weight becomes 5 (the lifted value), + # not 1 + 5 = 6. + assert objective.energy() == pytest.approx(7.0) + + +def test_objective_set_cost_inserts_and_accumulates(chain_with_lifted): + base, base_costs, _, _ = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective(base, base_costs) + + edge, is_new = objective.set_cost(0, 3, -5.0) + assert is_new + assert objective.number_of_lifted_edges == 1 + assert int(objective.lifted_graph.find_edge(0, 3)) == edge + + edge2, is_new2 = objective.set_cost(0, 3, -3.0) + assert not is_new2 + assert edge2 == edge + assert objective.weights[edge] == pytest.approx(-8.0) + + edge3, is_new3 = objective.set_cost(0, 3, 1.0, overwrite=True) + assert not is_new3 + assert objective.weights[edge3] == pytest.approx(1.0) + + +def test_objective_bfs_distance_inserts_within_k_hops(small_chain_bfs_problem): + base, base_costs = small_chain_bfs_problem + + objective = bic.graph.LiftedMulticutObjective(base, base_costs, bfs_distance=2) + # 5-node chain at distance 2: lifted edges are (0,2), (1,3), (2,4) — 3 edges. + assert objective.number_of_lifted_edges == 3 + lifted_uvs = objective.lifted_graph.uv_ids()[objective.number_of_base_edges:] + pairs = {tuple(sorted(map(int, uv))) for uv in lifted_uvs} + assert pairs == {(0, 2), (1, 3), (2, 4)} + + # Default lifted weights are zero so the energy equals the multicut energy. + expected = bic.graph.MulticutObjective(base, base_costs).energy() + assert objective.energy() == pytest.approx(expected) + + +def test_objective_bfs_distance_combines_with_explicit_lifted_costs(small_chain_bfs_problem): + base, base_costs = small_chain_bfs_problem + lifted_uvs = np.array([[0, 4]], dtype=np.uint64) + lifted_costs = np.array([-3.0], dtype=np.float64) + objective = bic.graph.LiftedMulticutObjective( + base, + base_costs, + bfs_distance=2, + lifted_uvs=lifted_uvs, + lifted_costs=lifted_costs, + ) + # BFS adds 3 edges; (0, 4) is new (distance 4) → 4 lifted edges total. + assert objective.number_of_lifted_edges == 4 + assert objective.energy([0, 0, 0, 0, 1]) == pytest.approx(1.0 - 3.0) + + +def test_objective_labels_and_reset(chain_with_lifted): + base, base_costs, lifted_uvs, lifted_costs = chain_with_lifted + objective = bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=lifted_uvs, lifted_costs=lifted_costs + ) + objective.labels = [0, 0, 1, 1] + np.testing.assert_array_equal(objective.labels, np.array([0, 0, 1, 1], dtype=np.uint64)) + objective.reset_labels() + np.testing.assert_array_equal( + objective.labels, np.arange(4, dtype=np.uint64) + ) + + +def test_objective_validation_errors(): + base = bic.graph.UndirectedGraph.from_edges(3, [[0, 1], [1, 2]]) + base_costs = np.array([1.0, 2.0], dtype=np.float64) + + with pytest.raises(ValueError, match="edge_costs"): + bic.graph.LiftedMulticutObjective(base, np.array([1.0], dtype=np.float64)) + + obj = bic.graph.LiftedMulticutObjective(base, base_costs) + with pytest.raises(ValueError, match="labels"): + obj.labels = [0, 1] + + with pytest.raises(ValueError, match="lifted_uvs and lifted_costs"): + bic.graph.LiftedMulticutObjective( + base, base_costs, lifted_uvs=np.array([[0, 2]], dtype=np.uint64) + ) + + with pytest.raises(ValueError, match="same length"): + bic.graph.LiftedMulticutObjective( + base, + base_costs, + lifted_uvs=np.array([[0, 2]], dtype=np.uint64), + lifted_costs=np.array([1.0, 2.0], dtype=np.float64), + ) + + with pytest.raises(ValueError, match="bfs_distance"): + bic.graph.LiftedMulticutObjective(base, base_costs, bfs_distance=0) diff --git a/tests/graph/test_breadth_first_search.py b/tests/graph/test_breadth_first_search.py new file mode 100644 index 0000000..87d9d84 --- /dev/null +++ b/tests/graph/test_breadth_first_search.py @@ -0,0 +1,71 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_bfs_chain_distances(): + graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [1, 2], [2, 3]]) + nodes, distances = bic.graph.breadth_first_search(graph, 0) + np.testing.assert_array_equal(nodes, np.array([0, 1, 2, 3], dtype=np.uint64)) + np.testing.assert_array_equal(distances, np.array([0, 1, 2, 3], dtype=np.uint64)) + + +def test_bfs_max_distance_limits_expansion(): + graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [1, 2], [2, 3]]) + nodes, distances = bic.graph.breadth_first_search(graph, 0, max_distance=2) + np.testing.assert_array_equal(nodes, np.array([0, 1, 2], dtype=np.uint64)) + np.testing.assert_array_equal(distances, np.array([0, 1, 2], dtype=np.uint64)) + + +def test_bfs_exclude_source(): + graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [1, 2], [2, 3]]) + nodes, distances = bic.graph.breadth_first_search( + graph, 1, max_distance=1, include_source=False + ) + np.testing.assert_array_equal(sorted(nodes.tolist()), [0, 2]) + np.testing.assert_array_equal(distances, np.array([1, 1], dtype=np.uint64)) + + +def test_bfs_branching_distances(): + # Star: 0 -- 1, 0 -- 2, 1 -- 3 + graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [0, 2], [1, 3]]) + nodes, distances = bic.graph.breadth_first_search(graph, 0) + order = np.argsort(nodes) + np.testing.assert_array_equal(nodes[order], np.array([0, 1, 2, 3], dtype=np.uint64)) + np.testing.assert_array_equal(distances[order], np.array([0, 1, 1, 2], dtype=np.uint64)) + + +def test_bfs_disconnected_component_not_reached(): + graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [2, 3]]) + nodes, distances = bic.graph.breadth_first_search(graph, 0) + np.testing.assert_array_equal(sorted(nodes.tolist()), [0, 1]) + np.testing.assert_array_equal(distances, np.array([0, 1], dtype=np.uint64)) + + +def test_bfs_invalid_source_raises(): + graph = bic.graph.UndirectedGraph.from_edges(2, [[0, 1]]) + with pytest.raises(ValueError, match="source"): + bic.graph.breadth_first_search(graph, 2) + + +def test_bfs_negative_max_distance_raises(): + graph = bic.graph.UndirectedGraph.from_edges(2, [[0, 1]]) + with pytest.raises(ValueError, match="max_distance"): + bic.graph.breadth_first_search(graph, 0, max_distance=-1) + + +def test_bfs_zero_max_distance_returns_only_source(): + graph = bic.graph.UndirectedGraph.from_edges(2, [[0, 1]]) + nodes, distances = bic.graph.breadth_first_search(graph, 0, max_distance=0) + np.testing.assert_array_equal(nodes, np.array([0], dtype=np.uint64)) + np.testing.assert_array_equal(distances, np.array([0], dtype=np.uint64)) + + +def test_bfs_zero_max_distance_exclude_source_returns_empty(): + graph = bic.graph.UndirectedGraph.from_edges(2, [[0, 1]]) + nodes, distances = bic.graph.breadth_first_search( + graph, 0, max_distance=0, include_source=False + ) + assert nodes.shape == (0,) + assert distances.shape == (0,) diff --git a/tests/graph/test_rag_lifted_features.py b/tests/graph/test_rag_lifted_features.py new file mode 100644 index 0000000..6a6e50f --- /dev/null +++ b/tests/graph/test_rag_lifted_features.py @@ -0,0 +1,307 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +@pytest.fixture +def four_quadrant_labels_2d(): + """4x4 image with four 2x2 quadrants labelled 0/1/2/3.""" + labels = np.array( + [ + [0, 0, 1, 1], + [0, 0, 1, 1], + [2, 2, 3, 3], + [2, 2, 3, 3], + ], + dtype=np.uint32, + ) + return labels + + +@pytest.fixture +def four_quadrant_rag_2d(four_quadrant_labels_2d): + return bic.graph.region_adjacency_graph(four_quadrant_labels_2d) + + +def test_lifted_edges_discovered_via_diagonal_offset( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + # Diagonal offset (2, 2): every hit connects label 0 (top-left) with + # label 3 (bottom-right). (0, 3) is not in the RAG -> lifted edge. + lifted = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, 2)] + ) + assert lifted.dtype == np.uint64 + np.testing.assert_array_equal(lifted, np.array([[0, 3]], dtype=np.uint64)) + + +def test_lifted_edges_anti_diagonal_offset( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + # (2, -2): connects label 1 (top-right) with label 2 (bottom-left). + lifted = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, -2)] + ) + np.testing.assert_array_equal(lifted, np.array([[1, 2]], dtype=np.uint64)) + + +def test_lifted_edges_skip_one_hop_offsets( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + # 1-hop offsets only hit pairs that already exist in the RAG. The + # function must skip them and return an empty result. + lifted = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, + four_quadrant_labels_2d, + [(0, 1), (1, 0), (0, -1), (-1, 0)], + ) + assert lifted.shape == (0, 2) + + +def test_lifted_edges_mixed_offsets(four_quadrant_labels_2d, four_quadrant_rag_2d): + # Mixing a long-range diagonal with 1-hop offsets: only the long-range + # contribution produces lifted edges. + lifted = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, + four_quadrant_labels_2d, + [(0, 1), (2, 2), (1, 0), (2, -2)], + ) + np.testing.assert_array_equal( + lifted, np.array([[0, 3], [1, 2]], dtype=np.uint64) + ) + + +def test_lifted_edges_pair_already_in_rag_skipped(): + # 3x3 with two segments: long-range offset reaches a pair that is + # already a local edge -> not reported as lifted. + labels = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 0], + ], + dtype=np.uint32, + ) + rag = bic.graph.region_adjacency_graph(labels) + lifted = bic.graph.lifted_edges_from_affinities( + rag, labels, [(2, 0), (0, 2), (2, 2)] + ) + assert lifted.shape == (0, 2) + + +def test_lifted_edges_empty_offsets(four_quadrant_labels_2d, four_quadrant_rag_2d): + lifted = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [] + ) + assert lifted.shape == (0, 2) + + +def test_lifted_edges_3d(): + # 3D label volume with two slabs across z and two columns across x. + labels = np.zeros((4, 1, 4), dtype=np.uint32) + labels[0, 0, :2] = 0 + labels[0, 0, 2:] = 1 + labels[1, 0, :2] = 0 + labels[1, 0, 2:] = 1 + labels[2, 0, :2] = 2 + labels[2, 0, 2:] = 3 + labels[3, 0, :2] = 2 + labels[3, 0, 2:] = 3 + rag = bic.graph.region_adjacency_graph(labels) + + # Diagonal (z, y, x) = (2, 0, 2) connects (0,*,0..1) with (2,*,2..3) + # i.e. label 0 with label 3 -> lifted (0, 3) not in RAG. + lifted = bic.graph.lifted_edges_from_affinities(rag, labels, [(2, 0, 2)]) + np.testing.assert_array_equal(lifted, np.array([[0, 3]], dtype=np.uint64)) + + +def test_lifted_edges_dedup_across_offsets( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + # (2, 2) and (-2, -2) discover the same (0, 3) pair from opposite ends. + lifted = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, 2), (-2, -2)] + ) + np.testing.assert_array_equal(lifted, np.array([[0, 3]], dtype=np.uint64)) + + +def test_lifted_edges_threads(four_quadrant_labels_2d, four_quadrant_rag_2d): + lifted_single = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, + four_quadrant_labels_2d, + [(2, 2), (2, -2)], + number_of_threads=1, + ) + lifted_multi = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, + four_quadrant_labels_2d, + [(2, 2), (2, -2)], + number_of_threads=4, + ) + np.testing.assert_array_equal(lifted_single, lifted_multi) + + +def test_lifted_affinity_features_basic( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + # Make a (1, 4, 4) affinity volume with known values. The diagonal + # (2, 2) offset has 4 valid hits in a 4x4 grid: (0,0), (0,1), (1,0), + # (1,1) all hitting the (0, 3) lifted edge. + affinities = np.zeros((1, 4, 4), dtype=np.float64) + affinities[0, 0, 0] = 0.1 + affinities[0, 0, 1] = 0.2 + affinities[0, 1, 0] = 0.3 + affinities[0, 1, 1] = 0.4 + + lifted_uvs = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, 2)] + ) + features = bic.graph.lifted_affinity_features( + four_quadrant_labels_2d, affinities, [(2, 2)], lifted_uvs + ) + assert features.shape == (1, 2) + assert features[0, 0] == pytest.approx(0.25) # mean of [0.1, 0.2, 0.3, 0.4] + assert features[0, 1] == pytest.approx(4.0) + + +def test_lifted_affinity_features_complex_columns( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + affinities = np.full((1, 4, 4), 0.5, dtype=np.float64) + lifted_uvs = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, 2)] + ) + features = bic.graph.lifted_affinity_features_complex( + four_quadrant_labels_2d, affinities, [(2, 2)], lifted_uvs + ) + assert features.shape == (1, 12) + # mean + assert features[0, 0] == pytest.approx(0.5) + # median, min, max + assert features[0, 1] == pytest.approx(0.5) + assert features[0, 3] == pytest.approx(0.5) + assert features[0, 4] == pytest.approx(0.5) + # count + assert features[0, -1] == pytest.approx(4.0) + + +def test_lifted_affinity_features_skips_one_hop( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + # A 1-hop offset would contribute lots of affinity hits to local edges, + # but the lifted accumulator must not bin any of those onto our lifted + # edge. + affinities = np.full((2, 4, 4), 0.7, dtype=np.float64) + # Channel 0 is 1-hop (should be skipped); channel 1 is the long-range + # diagonal that defines the lifted edge. + offsets = [(0, 1), (2, 2)] + lifted_uvs = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, offsets + ) + features = bic.graph.lifted_affinity_features( + four_quadrant_labels_2d, affinities, offsets, lifted_uvs + ) + # Only the 4 diagonal hits count. + assert features[0, 1] == pytest.approx(4.0) + assert features[0, 0] == pytest.approx(0.7) + + +def test_lifted_affinity_features_skips_local_hits(): + # If a long-range offset happens to connect a pair that is also a local + # edge, the lifted accumulator must not bin the affinity onto a lifted + # edge (the local edge is not in the lifted set). + labels = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 0], + ], + dtype=np.uint32, + ) + rag = bic.graph.region_adjacency_graph(labels) + # Offset (2, 2): in a 3x3 grid only (0, 0) -> (2, 2), both label 0. + # So no diff, no accumulation. We pass a lifted edge that doesn't exist + # in the data to make sure no accumulation happens. + lifted_uvs = np.array([[0, 1]], dtype=np.uint64) # actually a local edge! + affinities = np.full((1, 3, 3), 0.9, dtype=np.float64) + features = bic.graph.lifted_affinity_features( + labels, affinities, [(2, 2)], lifted_uvs + ) + # The (0,0)->(2,2) hit has matching labels, so no count. + assert features[0, 1] == pytest.approx(0.0) + + +def test_lifted_affinity_features_empty_lifted_uvs(four_quadrant_labels_2d): + affinities = np.full((1, 4, 4), 0.3, dtype=np.float64) + empty_uvs = np.zeros((0, 2), dtype=np.uint64) + features = bic.graph.lifted_affinity_features( + four_quadrant_labels_2d, affinities, [(2, 2)], empty_uvs + ) + assert features.shape == (0, 2) + + +def test_lifted_affinity_features_threads( + four_quadrant_labels_2d, four_quadrant_rag_2d +): + affinities = np.linspace(0.0, 1.0, num=16).reshape(1, 4, 4).astype(np.float64) + lifted_uvs = bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, 2)] + ) + single = bic.graph.lifted_affinity_features( + four_quadrant_labels_2d, affinities, [(2, 2)], lifted_uvs, + number_of_threads=1, + ) + multi = bic.graph.lifted_affinity_features( + four_quadrant_labels_2d, affinities, [(2, 2)], lifted_uvs, + number_of_threads=4, + ) + np.testing.assert_array_almost_equal(single, multi) + + +def test_lifted_affinity_features_3d(): + labels = np.zeros((4, 1, 4), dtype=np.uint32) + labels[0, 0, :2] = 0 + labels[0, 0, 2:] = 1 + labels[1, 0, :2] = 0 + labels[1, 0, 2:] = 1 + labels[2, 0, :2] = 2 + labels[2, 0, 2:] = 3 + labels[3, 0, :2] = 2 + labels[3, 0, 2:] = 3 + rag = bic.graph.region_adjacency_graph(labels) + + affinities = np.full((1, 4, 1, 4), 0.42, dtype=np.float64) + lifted_uvs = bic.graph.lifted_edges_from_affinities(rag, labels, [(2, 0, 2)]) + features = bic.graph.lifted_affinity_features( + labels, affinities, [(2, 0, 2)], lifted_uvs + ) + # (0,0,0)->(2,0,2): label 0 -> 3, hit + # (0,0,1)->(2,0,3): label 0 -> 3, hit + # (1,0,0)->(3,0,2): label 0 -> 3, hit + # (1,0,1)->(3,0,3): label 0 -> 3, hit + assert features[0, 0] == pytest.approx(0.42) + assert features[0, 1] == pytest.approx(4.0) + + +def test_lifted_affinity_features_validation(): + labels = np.zeros((4, 4), dtype=np.uint32) + affinities = np.zeros((1, 4, 4), dtype=np.float64) + bad_affinities = np.zeros((2, 4, 4), dtype=np.float64) + lifted_uvs = np.zeros((1, 2), dtype=np.uint64) + + with pytest.raises(ValueError, match="channel count"): + bic.graph.lifted_affinity_features( + labels, bad_affinities, [(2, 2)], lifted_uvs + ) + with pytest.raises(ValueError, match="ndim"): + bic.graph.lifted_affinity_features( + labels, affinities, [(2, 2, 0)], lifted_uvs + ) + + +def test_lifted_edges_validation(four_quadrant_labels_2d, four_quadrant_rag_2d): + with pytest.raises(ValueError, match="ndim"): + bic.graph.lifted_edges_from_affinities( + four_quadrant_rag_2d, four_quadrant_labels_2d, [(2, 2, 0)] + )