From d46f31ce1d3f4a2f5d8fbcf4095ce297e646b862 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Fri, 15 May 2026 08:24:54 +0200 Subject: [PATCH 1/5] Implement edge-weighted graph watershed --- MIGRATION_GUIDE.md | 44 +++ .../graph/edge_weighted_watershed.hxx | 105 +++++++ src/bindings/graph.cxx | 70 +++++ src/bioimage_cpp/graph/__init__.py | 90 ++++++ tests/graph/test_edge_weighted_watershed.py | 294 ++++++++++++++++++ 5 files changed, 603 insertions(+) create mode 100644 include/bioimage_cpp/graph/edge_weighted_watershed.hxx create mode 100644 tests/graph/test_edge_weighted_watershed.py diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index d952472..b77e306 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -254,6 +254,50 @@ Notes: - `number_of_threads=0` uses the library default; pass a positive integer for a fixed thread count. +## Edge-Weighted Watershed + +Nifty: + +```python +import nifty.graph as ng + +labels = ng.edgeWeightedWatershedsSegmentation(graph, edge_weights, seeds) +``` + +bioimage-cpp: + +```python +import bioimage_cpp as bic + +labels = bic.graph.edge_weighted_watershed(graph, edge_weights, seeds) +``` + +Notes: + +- Only the Kruskal variant of nifty's algorithm is provided. Edges are visited + in ascending weight order; two distinct components merge iff at least one is + unlabeled (seed label `0`), so seed boundaries are preserved. +- `graph` may be an `UndirectedGraph` or a `RegionAdjacencyGraph`. +- `edge_weights` must be 1D with length `graph.number_of_edges`. Supported + dtypes are `float32` and `float64`; other floating dtypes are promoted to + `float32` (matching nifty, whose Python binding is `float32`-only). Non-float + dtypes raise `TypeError`. +- `seeds` must be 1D with length `graph.number_of_nodes`. Supported dtypes are + `uint32`, `uint64`, `int32`, `int64`. The value `0` marks unlabeled nodes; + non-zero ids are propagated along low-weight paths. Signed seed arrays must + not contain negative values. +- The output is 1D with length `graph.number_of_nodes` and the same dtype as + `seeds`. Seed label values are preserved (no dense relabeling). Nodes that + no seed can reach remain `0`. + +Intentional differences vs. nifty: + +- No priority-queue variant — only the simpler sort + union-find Kruskal flow. + For the same input it matches nifty's default behavior (which also dispatches + to the Kruskal implementation). +- No carving / background-bias variant. Build a carving prior into the edge + weights before calling the function if needed. + ## Multicut Nifty exposes multicut through an objective + factory-style solver hierarchy. diff --git a/include/bioimage_cpp/graph/edge_weighted_watershed.hxx b/include/bioimage_cpp/graph/edge_weighted_watershed.hxx new file mode 100644 index 0000000..29c0506 --- /dev/null +++ b/include/bioimage_cpp/graph/edge_weighted_watershed.hxx @@ -0,0 +1,105 @@ +#pragma once + +#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +// Kruskal-style edge-weighted seeded watershed. +// +// Edges are visited in ascending weight order. Two distinct components are +// merged iff at least one of them is unlabeled (seed label 0); the non-zero +// seed label then propagates. Two distinct already-labeled components are +// never merged, so seed boundaries are preserved. Nodes that no seed can +// reach retain label 0. +// +// `WeightT` is any type orderable with `<` (typically float32 or float64). +// `SeedT` is any integer label type. Signed types must not contain negative +// values; this is checked at the boundary. +template +inline std::vector edge_weighted_watershed( + const UndirectedGraph &graph, + const std::vector &edge_weights, + const std::vector &seeds +) { + const auto number_of_nodes = graph.number_of_nodes(); + const auto number_of_edges = graph.number_of_edges(); + + if (edge_weights.size() != static_cast(number_of_edges)) { + throw std::invalid_argument( + "edge_weights length must equal number_of_edges, got " + + std::to_string(edge_weights.size()) + " for number_of_edges=" + + std::to_string(number_of_edges) + ); + } + if (seeds.size() != static_cast(number_of_nodes)) { + throw std::invalid_argument( + "seeds length must equal number_of_nodes, got " + + std::to_string(seeds.size()) + " for number_of_nodes=" + + std::to_string(number_of_nodes) + ); + } + if constexpr (std::is_signed_v) { + for (const auto value : seeds) { + if (value < 0) { + throw std::invalid_argument("seeds must not contain negative values"); + } + } + } + + std::vector labels(seeds); + + if (number_of_edges == 0) { + return labels; + } + + std::vector order(static_cast(number_of_edges)); + for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { + order[static_cast(edge)] = edge; + } + std::stable_sort( + order.begin(), + order.end(), + [&edge_weights](const std::uint64_t a, const std::uint64_t b) { + return edge_weights[static_cast(a)] < + edge_weights[static_cast(b)]; + } + ); + + detail::UnionFind sets(static_cast(number_of_nodes)); + + constexpr SeedT zero{0}; + for (const auto edge : order) { + const auto uv = graph.uv(edge); + const auto ru = sets.find(uv.first); + const auto rv = sets.find(uv.second); + if (ru == rv) { + continue; + } + const auto lu = labels[static_cast(ru)]; + const auto lv = labels[static_cast(rv)]; + if (lu != zero && lv != zero) { + continue; + } + const auto new_label = std::max(lu, lv); + const auto new_root = sets.unite_roots(ru, rv); + labels[static_cast(new_root)] = new_label; + } + + for (std::uint64_t node = 0; node < number_of_nodes; ++node) { + const auto root = sets.find(node); + labels[static_cast(node)] = labels[static_cast(root)]; + } + + return labels; +} + +} // namespace bioimage_cpp::graph diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 5055e4f..1a85395 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -2,6 +2,7 @@ #include "bioimage_cpp/array_view.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/multicut.hxx" #include "bioimage_cpp/graph/node_label_projection.hxx" @@ -12,6 +13,7 @@ #include #include +#include #include #include #include @@ -285,6 +287,55 @@ UInt64Array graph_connected_components_masked(const Graph &graph, ConstUInt8Arra return vector_to_uint64_array(labels); } +template +using ConstArray1D = nb::ndarray; + +template +std::vector array_1d_to_vector( + ConstArray1D array, + const char *argument_name, + const std::uint64_t expected_size +) { + if (array.ndim() != 1) { + throw std::invalid_argument(std::string(argument_name) + " must be a 1D array"); + } + if (array.shape(0) != static_cast(expected_size)) { + throw std::invalid_argument( + std::string(argument_name) + " length must match expected size" + ); + } + const auto *data = array.data(); + return std::vector(data, data + array.shape(0)); +} + +template +nb::ndarray vector_to_array_1d(const std::vector &values) { + const std::size_t size = values.size(); + auto *data = new T[size]; + std::copy(values.begin(), values.end(), data); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + const std::array shape{size}; + return nb::ndarray(data, 1, shape.data(), owner); +} + +template +nb::ndarray graph_edge_weighted_watershed_t( + const Graph &graph, + ConstArray1D edge_weights, + ConstArray1D seeds +) { + const auto weight_vector = + array_1d_to_vector(edge_weights, "edge_weights", graph.number_of_edges()); + const auto seed_vector = + array_1d_to_vector(seeds, "seeds", graph.number_of_nodes()); + std::vector labels; + { + nb::gil_scoped_release release; + labels = graph::edge_weighted_watershed(graph, weight_vector, seed_vector); + } + return vector_to_array_1d(labels); +} + double multicut_energy(const Graph &graph, ConstDoubleArray costs, ConstUInt64Array labels) { const auto cost_vector = double_array_to_vector(costs, "edge_costs", graph.number_of_edges()); const auto label_vector = uint64_array_to_vector(labels, "labels", graph.number_of_nodes()); @@ -605,6 +656,25 @@ void bind_graph(nb::module_ &m) { nb::arg("graph"), nb::arg("edge_mask") ); + const auto register_watershed = [&m]( + const char *name + ) { + m.def( + name, + &graph_edge_weighted_watershed_t, + nb::arg("graph"), + nb::arg("edge_weights"), + nb::arg("seeds") + ); + }; + register_watershed.operator()("_edge_weighted_watershed_float32_uint32"); + register_watershed.operator()("_edge_weighted_watershed_float32_uint64"); + register_watershed.operator()("_edge_weighted_watershed_float32_int32"); + register_watershed.operator()("_edge_weighted_watershed_float32_int64"); + register_watershed.operator()("_edge_weighted_watershed_float64_uint32"); + register_watershed.operator()("_edge_weighted_watershed_float64_uint64"); + register_watershed.operator()("_edge_weighted_watershed_float64_int32"); + register_watershed.operator()("_edge_weighted_watershed_float64_int64"); m.def( "_multicut_energy", &multicut_energy, diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index a2f05fc..1916663 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -36,6 +36,17 @@ np.dtype("int64"): _core._accumulate_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, + (np.dtype("float32"), np.dtype("int32")): _core._edge_weighted_watershed_float32_int32, + (np.dtype("float32"), np.dtype("int64")): _core._edge_weighted_watershed_float32_int64, + (np.dtype("float64"), np.dtype("uint32")): _core._edge_weighted_watershed_float64_uint32, + (np.dtype("float64"), np.dtype("uint64")): _core._edge_weighted_watershed_float64_uint64, + (np.dtype("float64"), np.dtype("int32")): _core._edge_weighted_watershed_float64_int32, + (np.dtype("float64"), np.dtype("int64")): _core._edge_weighted_watershed_float64_int64, +} + _PROJECT_NODE_LABELS_TO_PIXELS_BY_DTYPE = { np.dtype("uint32"): _core._project_node_labels_to_pixels_uint32, np.dtype("uint64"): _core._project_node_labels_to_pixels_uint64, @@ -159,6 +170,17 @@ def _as_node_labels(labels, graph: UndirectedGraph | RegionAdjacencyGraph) -> np return np.ascontiguousarray(array) +def _as_1d_array(values, dtype, name: str, expected_size: int) -> np.ndarray: + array = np.asarray(values, dtype=dtype) + if array.ndim != 1: + raise ValueError(f"{name} must be a 1D array") + if array.shape[0] != expected_size: + raise ValueError( + f"{name} length must be {expected_size}, got {array.shape[0]}" + ) + return np.ascontiguousarray(array) + + def _dense_labels(labels) -> np.ndarray: labels = np.asarray(labels, dtype=np.uint64) _, dense = np.unique(labels, return_inverse=True) @@ -207,6 +229,73 @@ def connected_components( ) +def edge_weighted_watershed( + graph: UndirectedGraph | RegionAdjacencyGraph, + edge_weights, + seeds, +) -> np.ndarray: + """Kruskal-style edge-weighted seeded watershed on an undirected graph. + + Edges are visited in ascending weight order. Two distinct components are + merged iff at least one of them is unlabeled (seed label ``0``); the + non-zero seed label then propagates. Two distinct already-labeled + components are never merged, so seed boundaries are preserved. + + Parameters + ---------- + graph: + :class:`UndirectedGraph` or :class:`RegionAdjacencyGraph`. + edge_weights: + 1D array of length ``graph.number_of_edges``. Supported dtypes are + ``float32`` and ``float64``. Other floating dtypes are cast to + ``float32`` (matches nifty); other dtypes raise ``TypeError``. + seeds: + 1D array of length ``graph.number_of_nodes``. Supported dtypes are + ``uint32``, ``uint64``, ``int32``, ``int64``. ``0`` marks unlabeled + nodes; positive ids are seed labels and propagate along low-weight + paths. Signed seed arrays must not contain negative values. + + Returns + ------- + np.ndarray + 1D array of length ``graph.number_of_nodes`` with the same dtype as + ``seeds``. Nodes reachable from a seed receive that seed's label; + unreachable nodes remain ``0``. Seed label values are preserved (no + dense relabeling). + """ + weight_array = np.asarray(edge_weights) + if weight_array.dtype not in (np.dtype("float32"), np.dtype("float64")): + if np.issubdtype(weight_array.dtype, np.floating): + weight_array = weight_array.astype(np.float32, copy=False) + else: + raise TypeError( + "edge_weights must have dtype float32 or float64, got " + f"dtype={weight_array.dtype}" + ) + + seed_array = np.asarray(seeds) + if seed_array.dtype not in ( + np.dtype("uint32"), + np.dtype("uint64"), + np.dtype("int32"), + np.dtype("int64"), + ): + raise TypeError( + "seeds must have dtype uint32, uint64, int32, or int64, got " + f"dtype={seed_array.dtype}" + ) + + weight_array = _as_1d_array( + weight_array, weight_array.dtype, "edge_weights", int(graph.number_of_edges) + ) + seed_array = _as_1d_array( + seed_array, seed_array.dtype, "seeds", int(graph.number_of_nodes) + ) + + run = _EDGE_WEIGHTED_WATERSHED_BY_DTYPE[(weight_array.dtype, seed_array.dtype)] + return run(graph, weight_array, seed_array) + + class MulticutObjective: """Multicut objective for an undirected graph and edge costs.""" @@ -672,6 +761,7 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "connected_components", "edge_map_features", "edge_map_features_complex", + "edge_weighted_watershed", "external_multicut_problem_path", "load_external_multicut_problem", "load_external_multicut_problem_data", diff --git a/tests/graph/test_edge_weighted_watershed.py b/tests/graph/test_edge_weighted_watershed.py new file mode 100644 index 0000000..f925b99 --- /dev/null +++ b/tests/graph/test_edge_weighted_watershed.py @@ -0,0 +1,294 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def _path_graph(n: int) -> bic.graph.UndirectedGraph: + edges = np.array([[i, i + 1] for i in range(n - 1)], dtype=np.uint64) + return bic.graph.UndirectedGraph.from_edges(n, edges) + + +def test_propagates_along_lowest_weight_path(): + graph = _path_graph(5) + # Weights: 0-1=0.1, 1-2=0.5, 2-3=0.2, 3-4=0.1 + # Seeds at node 0 (label 1) and node 4 (label 2). + # Ascending sort: edges (0-1, 3-4) tie, then (2-3), then (1-2). + # After 0-1 and 3-4: nodes 0,1 -> 1; nodes 3,4 -> 2. + # 2-3: node 2 unlabeled, joins label 2 component. + # 1-2: both labeled, no merge — boundary stays. + weights = np.array([0.1, 0.5, 0.2, 0.1], dtype=np.float64) + seeds = np.array([1, 0, 0, 0, 2], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + assert labels.dtype == np.uint64 + np.testing.assert_array_equal(labels, np.array([1, 1, 2, 2, 2], dtype=np.uint64)) + + +def test_distinct_labeled_components_are_not_merged(): + graph = _path_graph(3) + weights = np.array([0.1, 0.2], dtype=np.float64) + seeds = np.array([1, 0, 2], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + np.testing.assert_array_equal(labels, np.array([1, 1, 2], dtype=np.uint64)) + + +def test_components_with_same_label_stay_separate(): + # Matches nifty's Kruskal variant: two pre-labeled components with the + # same label are not merged because the condition requires at least one + # to be unlabeled. + graph = _path_graph(3) + weights = np.array([0.1, 0.2], dtype=np.float64) + seeds = np.array([1, 0, 1], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + # Edge 0-1 merges node 1 into label 1. Edge 1-2 then has both endpoints + # labeled (1 and 1), so no merge happens — but the result still gives + # node 2 label 1 from its own seed. + np.testing.assert_array_equal(labels, np.array([1, 1, 1], dtype=np.uint64)) + + +def test_unreachable_nodes_remain_zero(): + graph = bic.graph.UndirectedGraph(4) # no edges + weights = np.zeros(0, dtype=np.float64) + seeds = np.array([1, 0, 2, 0], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + np.testing.assert_array_equal(labels, seeds) + + +def test_no_seeds_returns_zero_labels(): + graph = _path_graph(4) + weights = np.array([0.1, 0.2, 0.3], dtype=np.float64) + seeds = np.zeros(4, dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + np.testing.assert_array_equal(labels, np.zeros(4, dtype=np.uint64)) + + +def test_high_weight_edge_is_processed_last(): + # Y-shape: nodes 0 and 2 are seeds, node 1 is the junction. + # The 1-3 edge has a high weight, so node 3 is grabbed by node 1's + # component last; the seed reached via the lower-weight path wins. + graph = bic.graph.UndirectedGraph.from_edges( + 4, np.array([[0, 1], [1, 2], [1, 3]], dtype=np.uint64) + ) + weights = np.array([0.1, 0.2, 0.9], dtype=np.float64) + seeds = np.array([1, 0, 2, 0], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + # Edge 0 (0-1, w=0.1): node 1 gets label 1. + # Edge 1 (1-2, w=0.2): both labeled (1 vs 2), no merge. + # Edge 2 (1-3, w=0.9): node 3 joins component of node 1 (label 1). + np.testing.assert_array_equal(labels, np.array([1, 1, 2, 1], dtype=np.uint64)) + + +def test_accepts_region_adjacency_graph(): + labels_image = np.array([[1, 1, 2], [3, 3, 2]], dtype=np.uint64) + rag = bic.graph.region_adjacency_graph(labels_image) + + weights = np.ones(rag.number_of_edges, dtype=np.float64) + seeds = np.array([0, 7, 0, 0], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(rag, weights, seeds) + + assert labels.dtype == np.uint64 + assert labels.shape == (rag.number_of_nodes,) + # node 1 is the only seed, so all reachable nodes get label 7. + np.testing.assert_array_equal(labels[1:], np.full(3, 7, dtype=np.uint64)) + + +def test_preserves_seed_label_values(): + graph = _path_graph(4) + weights = np.array([0.1, 0.2, 0.3], dtype=np.float64) + seeds = np.array([42, 0, 0, 99], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + assert set(np.unique(labels).tolist()).issubset({42, 99}) + + +def test_accepts_list_inputs(): + graph = _path_graph(3) + + labels = bic.graph.edge_weighted_watershed(graph, [0.1, 0.2], [1, 0, 2]) + + # numpy default int dtype is int64 on this platform; the seed dtype is + # preserved on the output. + np.testing.assert_array_equal(labels, np.array([1, 1, 2])) + + +@pytest.mark.parametrize("weight_dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("seed_dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_dtype_combinations_preserve_seed_dtype(weight_dtype, seed_dtype): + graph = _path_graph(5) + weights = np.array([0.1, 0.5, 0.2, 0.1], dtype=weight_dtype) + seeds = np.array([1, 0, 0, 0, 2], dtype=seed_dtype) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + assert labels.dtype == np.dtype(seed_dtype) + np.testing.assert_array_equal( + labels, np.array([1, 1, 2, 2, 2], dtype=seed_dtype) + ) + + +def test_float32_weights_are_not_copied_to_float64(): + # The float32 path should accept float32 inputs without an upcast. + graph = _path_graph(3) + weights = np.array([0.1, 0.2], dtype=np.float32) + seeds = np.array([1, 0, 2], dtype=np.uint32) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + assert labels.dtype == np.dtype("uint32") + np.testing.assert_array_equal(labels, np.array([1, 1, 2], dtype=np.uint32)) + + +def test_float16_weights_are_promoted_to_float32(): + graph = _path_graph(3) + weights = np.array([0.1, 0.2], dtype=np.float16) + seeds = np.array([1, 0, 2], dtype=np.uint64) + + labels = bic.graph.edge_weighted_watershed(graph, weights, seeds) + + np.testing.assert_array_equal(labels, np.array([1, 1, 2], dtype=np.uint64)) + + +def test_rejects_unsupported_weight_dtype(): + graph = _path_graph(3) + with pytest.raises(TypeError, match="edge_weights must have dtype"): + bic.graph.edge_weighted_watershed( + graph, + np.array([1, 2], dtype=np.int32), + np.array([1, 0, 2], dtype=np.uint64), + ) + + +def test_rejects_unsupported_seed_dtype(): + graph = _path_graph(3) + with pytest.raises(TypeError, match="seeds must have dtype"): + bic.graph.edge_weighted_watershed( + graph, + np.array([0.1, 0.2], dtype=np.float32), + np.array([1, 0, 2], dtype=np.uint8), + ) + + +def test_rejects_negative_signed_seeds(): + graph = _path_graph(3) + with pytest.raises(ValueError, match="seeds must not contain negative values"): + bic.graph.edge_weighted_watershed( + graph, + np.array([0.1, 0.2], dtype=np.float32), + np.array([1, -1, 2], dtype=np.int32), + ) + + +def test_rejects_wrong_weight_length(): + graph = _path_graph(3) + with pytest.raises(ValueError, match="edge_weights length"): + bic.graph.edge_weighted_watershed( + graph, + np.zeros(3, dtype=np.float64), + np.array([1, 0, 2], dtype=np.uint64), + ) + + +def test_rejects_wrong_seeds_length(): + graph = _path_graph(3) + with pytest.raises(ValueError, match="seeds length"): + bic.graph.edge_weighted_watershed( + graph, + np.array([0.1, 0.2], dtype=np.float64), + np.array([1, 0], dtype=np.uint64), + ) + + +def test_rejects_non_1d_inputs(): + graph = _path_graph(3) + with pytest.raises(ValueError, match="edge_weights must be a 1D array"): + bic.graph.edge_weighted_watershed( + graph, + np.zeros((2, 1), dtype=np.float64), + np.array([1, 0, 2], dtype=np.uint64), + ) + with pytest.raises(ValueError, match="seeds must be a 1D array"): + bic.graph.edge_weighted_watershed( + graph, + np.array([0.1, 0.2], dtype=np.float64), + np.zeros((3, 1), dtype=np.uint64), + ) + + +def _reference_kruskal_watershed(graph, weights, seeds): + """Pure-Python reference implementation of the Kruskal seeded watershed. + + Mirrors nifty's edgeWeightedWatershedsSegmentationKruskalImpl: walks edges + in ascending weight order, merges components iff at least one is unlabeled, + propagates the non-zero label. + """ + n = int(graph.number_of_nodes) + weights = np.asarray(weights, dtype=np.float64) + labels = np.asarray(seeds, dtype=np.uint64).copy() + uvs = graph.uv_ids() + + parent = np.arange(n, dtype=np.int64) + + def find(x): + root = x + while parent[root] != root: + root = parent[root] + while parent[x] != root: + parent[x], x = root, parent[x] + return root + + order = np.argsort(weights, kind="stable") + for edge_index in order: + u, v = int(uvs[edge_index, 0]), int(uvs[edge_index, 1]) + ru, rv = find(u), find(v) + if ru == rv: + continue + lu, lv = labels[ru], labels[rv] + if lu != 0 and lv != 0: + continue + new_label = max(int(lu), int(lv)) + parent[rv] = ru + labels[ru] = new_label + labels[rv] = new_label + + for node in range(n): + labels[node] = labels[find(node)] + return labels + + +def test_matches_python_reference_on_random_graph(): + rng = np.random.default_rng(42) + n_nodes = 50 + n_edges = 150 + + edges = set() + while len(edges) < n_edges: + u, v = rng.integers(0, n_nodes, size=2) + if u != v: + edges.add((min(int(u), int(v)), max(int(u), int(v)))) + uvs = np.array(list(edges), dtype=np.uint64) + + graph = bic.graph.UndirectedGraph.from_edges(n_nodes, uvs) + weights = rng.standard_normal(graph.number_of_edges).astype(np.float64) + seeds = np.zeros(n_nodes, dtype=np.uint64) + seed_nodes = rng.choice(n_nodes, size=5, replace=False) + for label, node in enumerate(seed_nodes, start=1): + seeds[node] = label + + actual = bic.graph.edge_weighted_watershed(graph, weights, seeds) + expected = _reference_kruskal_watershed(graph, weights, seeds) + + np.testing.assert_array_equal(actual, expected) From 397e482327e910b6b1aed593ec2dbfa04962ce6b Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Fri, 15 May 2026 10:20:34 +0200 Subject: [PATCH 2/5] Add fusion moves implementation --- MIGRATION_GUIDE.md | 71 +++++++ .../graph/detail/fusion_contract.hxx | 107 ++++++++++ .../graph/multicut/fusion_move.hxx | 188 +++++++++++++++++ .../bioimage_cpp/graph/proposal_generator.hxx | 35 ++++ .../greedy_additive_multicut.hxx | 77 +++++++ .../graph/proposal_generators/watershed.hxx | 107 ++++++++++ src/bindings/graph.cxx | 137 ++++++++++++ src/bioimage_cpp/graph/__init__.py | 167 +++++++++++++++ tests/graph/multicut/test_fusion_move.py | 195 ++++++++++++++++++ .../multicut/test_proposal_generators.py | 67 ++++++ 10 files changed, 1151 insertions(+) create mode 100644 include/bioimage_cpp/graph/detail/fusion_contract.hxx create mode 100644 include/bioimage_cpp/graph/multicut/fusion_move.hxx create mode 100644 include/bioimage_cpp/graph/proposal_generator.hxx create mode 100644 include/bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx create mode 100644 include/bioimage_cpp/graph/proposal_generators/watershed.hxx create mode 100644 tests/graph/multicut/test_fusion_move.py create mode 100644 tests/graph/multicut/test_proposal_generators.py diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index b77e306..d47813f 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -408,6 +408,77 @@ Intentional differences vs. nifty: `GreedyAdditiveMulticut` and no fallthrough is given — the greedy solver already operates on each connected component internally. +## Fusion Moves (Multicut) + +Nifty exposes the fusion-move multicut solver via the factory hierarchy with a +chosen proposal generator and sub-solver factory. + +Nifty: + +```python +import nifty.graph.opt.multicut as nmc + +objective = nmc.multicutObjective(graph, edge_costs) +pgen = nmc.watershedProposals(sigma=1.0, numberOfSeeds=0.1) +factory = nmc.fusionMoveBasedFactory( + proposalGenerator=pgen, + fusionMove=nmc.fusionMoveSettings( + mcFactory=nmc.greedyAdditiveFactory(), + ), + numberOfIterations=10, + stopIfNoImprovement=4, +) +labels = factory.create(objective).optimize() +``` + +bioimage-cpp: + +```python +import bioimage_cpp as bic + +objective = bic.graph.MulticutObjective(graph, edge_costs) +solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator( + sigma=1.0, n_seeds_fraction=0.1, seed=0, + ), + sub_solver=bic.graph.GreedyAdditiveMulticut(), + number_of_iterations=10, + stop_if_no_improvement=4, +) +labels = solver.optimize(objective) +``` + +Proposal generators: + +| nifty proposal generator | bioimage-cpp proposal generator | +| --- | --- | +| `watershedProposals(sigma=..., numberOfSeeds=...)` | `WatershedProposalGenerator(sigma=..., n_seeds_fraction=..., seed=...)` | +| `greedyAdditiveProposals(sigma=..., weightStopCond=..., nodeNumStopCond=...)` | `GreedyAdditiveProposalGenerator(sigma=..., weight_stop=..., node_num_stop=..., seed=...)` | + +Sub-solvers: any built-in multicut solver (`GreedyAdditiveMulticut`, +`GreedyFixationMulticut`, `KernighanLinMulticut`). If `sub_solver` is omitted, +the default is `GreedyAdditiveMulticut` constructed with no-noise defaults. + +Intentional differences vs. nifty: + +- Single object construction: no separate factory / solver step. +- Proposal generators are Python classes carrying their settings; the C++ + proposal-generator object is built lazily when `optimize` is called. +- The driver warm-starts from the trivial singleton labeling by running the + default greedy-additive sub-solver once before the proposal loop. +- A best-of-three safety net (current, proposal, fused) keeps the running + energy monotonically non-increasing across iterations. +- The current implementation is single-threaded and uses pairwise + (proposal, current) fuses. `number_of_threads` and + `number_of_parallel_proposals` are kept on the API for forward compatibility + but must be left at their defaults (`1` and `2` respectively). + +Notes: + +- Custom Python proposal generators are not yet supported; subclass + `ProposalGenerator` and provide your own `_build` returning a C++ + proposal-generator object if you need to extend the set. + ## Segmentation Overlaps Nifty: diff --git a/include/bioimage_cpp/graph/detail/fusion_contract.hxx b/include/bioimage_cpp/graph/detail/fusion_contract.hxx new file mode 100644 index 0000000..f226bb5 --- /dev/null +++ b/include/bioimage_cpp/graph/detail/fusion_contract.hxx @@ -0,0 +1,107 @@ +#pragma once + +#include "bioimage_cpp/detail/relabel.hxx" +#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph::detail { + +struct AgreementContraction { + // Contracted graph: one node per agreement component, dense ids in + // [0, number_of_components). + UndirectedGraph contracted_graph; + // For every original edge id: the contracted edge id it maps onto, or + // -1 if both endpoints collapsed into the same component. + std::vector contracted_edge_of_original; + // For every original node id: its dense agreement-component id. + std::vector root_of_node; +}; + +// Build the agreement-projection contracted graph: merge `(u, v)` iff every +// proposal labels `u` and `v` identically, then dense-relabel and assemble the +// contracted graph with one edge per distinct surviving (root_u, root_v) pair. +// +// `proposals` is a row-major buffer of shape (n_proposals, number_of_nodes). +// Passing 0 proposals collapses every edge (all proposals trivially agree) +// and yields a single-node contracted graph; we reject that as a usage error. +inline AgreementContraction contract_by_agreement( + const UndirectedGraph &graph, + const std::uint64_t *proposals, + const std::size_t n_proposals, + const std::size_t n_nodes_per_proposal +) { + if (n_proposals == 0) { + throw std::invalid_argument("at least one proposal is required"); + } + const auto number_of_nodes = graph.number_of_nodes(); + const auto number_of_edges = graph.number_of_edges(); + if (n_nodes_per_proposal != static_cast(number_of_nodes)) { + throw std::invalid_argument( + "proposal width must equal number_of_nodes, got " + + std::to_string(n_nodes_per_proposal) + " for number_of_nodes=" + + std::to_string(number_of_nodes) + ); + } + + bioimage_cpp::detail::UnionFind sets(static_cast(number_of_nodes)); + for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { + const auto uv = graph.uv(edge); + const auto u = static_cast(uv.first); + const auto v = static_cast(uv.second); + bool agree = true; + for (std::size_t p = 0; p < n_proposals; ++p) { + const auto *row = proposals + p * n_nodes_per_proposal; + if (row[u] != row[v]) { + agree = false; + break; + } + } + if (agree) { + sets.merge(uv.first, uv.second); + } + } + + std::vector raw_root(static_cast(number_of_nodes)); + for (std::uint64_t node = 0; node < number_of_nodes; ++node) { + raw_root[static_cast(node)] = sets.find(node); + } + auto root_of_node = bioimage_cpp::detail::dense_relabel(raw_root); + + std::uint64_t number_of_components = 0; + for (const auto root : root_of_node) { + if (root + 1 > number_of_components) { + number_of_components = root + 1; + } + } + + UndirectedGraph contracted_graph(number_of_components); + std::vector contracted_edge_of_original( + static_cast(number_of_edges), -1 + ); + + for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { + const auto uv = graph.uv(edge); + const auto ru = root_of_node[static_cast(uv.first)]; + const auto rv = root_of_node[static_cast(uv.second)]; + if (ru == rv) { + continue; + } + const auto inserted = contracted_graph.insert_edge(ru, rv); + contracted_edge_of_original[static_cast(edge)] = + static_cast(inserted); + } + + return AgreementContraction{ + std::move(contracted_graph), + std::move(contracted_edge_of_original), + std::move(root_of_node), + }; +} + +} // namespace bioimage_cpp::graph::detail diff --git a/include/bioimage_cpp/graph/multicut/fusion_move.hxx b/include/bioimage_cpp/graph/multicut/fusion_move.hxx new file mode 100644 index 0000000..f59a753 --- /dev/null +++ b/include/bioimage_cpp/graph/multicut/fusion_move.hxx @@ -0,0 +1,188 @@ +#pragma once + +#include "bioimage_cpp/graph/detail/fusion_contract.hxx" +#include "bioimage_cpp/graph/multicut/greedy_additive.hxx" +#include "bioimage_cpp/graph/multicut/objective.hxx" +#include "bioimage_cpp/graph/proposal_generator.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph::multicut { + +class FusionMoveSolver final : public SolverBase { +public: + // The proposal generator is borrowed; the caller owns its lifetime. + // The sub-solver pointer is optional: when null, the driver uses a default + // greedy-additive sub-solver internally. + // + // `number_of_threads` and `number_of_parallel_proposals` are reserved for + // future use; v1 only supports the single-threaded pairwise + // (proposal, current) fuse and rejects other values. + FusionMoveSolver( + ProposalGeneratorBase &proposal_generator, + const SolverBase *sub_solver = nullptr, + const std::size_t number_of_iterations = 10, + const std::size_t stop_if_no_improvement = 4, + const std::size_t number_of_threads = 1, + const std::size_t number_of_parallel_proposals = 2 + ) + : proposal_generator_(proposal_generator), + sub_solver_(sub_solver), + number_of_iterations_(number_of_iterations), + stop_if_no_improvement_(stop_if_no_improvement) { + if (number_of_threads != 1) { + throw std::invalid_argument( + "FusionMoveSolver currently supports number_of_threads=1 only" + ); + } + if (number_of_parallel_proposals != 2) { + throw std::invalid_argument( + "FusionMoveSolver currently supports number_of_parallel_proposals=2 only" + ); + } + } + + std::vector optimize(Objective &objective) const override { + const auto &graph = objective.graph(); + const auto &costs = objective.costs(); + const auto number_of_nodes = graph.number_of_nodes(); + + std::vector current = objective.labels(); + if (number_of_nodes == 0 || graph.number_of_edges() == 0) { + objective.set_labels(current); + return objective.labels(); + } + + const auto default_sub_solver = GreedyAdditiveSolver(); + const auto &sub_solver = sub_solver_ != nullptr + ? *sub_solver_ + : static_cast(default_sub_solver); + + // Warm start from greedy-additive if the caller passed the trivial + // singleton labeling. + if (is_singleton_labeling(current)) { + Objective warm_objective(graph, costs); + current = default_sub_solver.optimize(warm_objective); + } + + double current_energy = energy(graph, costs, current); + + std::vector proposal(static_cast(number_of_nodes)); + std::size_t iterations_without_improvement = 0; + + for (std::size_t iteration = 0; iteration < number_of_iterations_; ++iteration) { + proposal_generator_.generate(current, proposal); + + const auto fused = fuse_pair(graph, costs, current, proposal, sub_solver); + const auto fused_energy = energy(graph, costs, fused); + const auto proposal_energy = energy(graph, costs, proposal); + + // Best-of safety net across the three candidates so an iteration + // can never raise the running energy. + double best_energy = current_energy; + const std::vector *best = ¤t; + if (fused_energy < best_energy) { + best_energy = fused_energy; + best = &fused; + } + if (proposal_energy < best_energy) { + best_energy = proposal_energy; + best = &proposal; + } + + if (best_energy < current_energy) { + current = *best; + current_energy = best_energy; + iterations_without_improvement = 0; + } else { + ++iterations_without_improvement; + if (iterations_without_improvement >= stop_if_no_improvement_) { + break; + } + } + } + + objective.set_labels(current); + return objective.labels(); + } + +private: + static bool is_singleton_labeling(const std::vector &labels) { + for (std::size_t index = 0; index < labels.size(); ++index) { + if (labels[index] != static_cast(index)) { + return false; + } + } + return true; + } + + static std::vector fuse_pair( + const UndirectedGraph &graph, + const std::vector &costs, + const std::vector ¤t, + const std::vector &proposal, + const SolverBase &sub_solver + ) { + const auto number_of_nodes = static_cast(graph.number_of_nodes()); + std::vector stacked(2 * number_of_nodes); + std::copy(current.begin(), current.end(), stacked.begin()); + std::copy(proposal.begin(), proposal.end(), stacked.begin() + number_of_nodes); + + auto contraction = ::bioimage_cpp::graph::detail::contract_by_agreement( + graph, stacked.data(), 2, number_of_nodes + ); + + const auto &contracted_graph = contraction.contracted_graph; + const auto number_of_contracted_edges = contracted_graph.number_of_edges(); + + std::vector contracted_costs( + static_cast(number_of_contracted_edges), 0.0 + ); + for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { + const auto target = contraction.contracted_edge_of_original[ + static_cast(edge) + ]; + if (target < 0) { + continue; + } + contracted_costs[static_cast(target)] += + costs[static_cast(edge)]; + } + + if (number_of_contracted_edges == 0) { + std::vector result(number_of_nodes); + for (std::uint64_t node = 0; node < graph.number_of_nodes(); ++node) { + result[static_cast(node)] = contraction.root_of_node[ + static_cast(node) + ]; + } + return result; + } + + Objective sub_objective(contracted_graph, std::move(contracted_costs)); + const auto sub_labels = sub_solver.optimize(sub_objective); + + std::vector result(number_of_nodes); + for (std::uint64_t node = 0; node < graph.number_of_nodes(); ++node) { + const auto root = contraction.root_of_node[static_cast(node)]; + result[static_cast(node)] = sub_labels[ + static_cast(root) + ]; + } + return result; + } + + ProposalGeneratorBase &proposal_generator_; + const SolverBase *sub_solver_; + std::size_t number_of_iterations_; + std::size_t stop_if_no_improvement_; +}; + +} // namespace bioimage_cpp::graph::multicut diff --git a/include/bioimage_cpp/graph/proposal_generator.hxx b/include/bioimage_cpp/graph/proposal_generator.hxx new file mode 100644 index 0000000..e7c0edb --- /dev/null +++ b/include/bioimage_cpp/graph/proposal_generator.hxx @@ -0,0 +1,35 @@ +#pragma once + +#include +#include + +namespace bioimage_cpp::graph { + +// Base class for proposal generators consumed by fusion-move solvers. +// +// A proposal generator is constructed once per problem with whatever data it +// needs (graph, edge costs, hyperparameters, RNG seed). The driver then calls +// `generate(current, proposal)` once per fusion-move iteration. Each +// implementation may use `current` (e.g. for warm-start) or ignore it. +// +// Output proposals must: +// - Have the same length as the graph's node count. +// - Use label `0` reserved for "background" only if the implementation +// explicitly documents this; the fusion-move driver does not interpret +// label values, only their equality classes. +// +// Implementations should be deterministic given their constructor seed and +// the number of times `generate` has been called. +class ProposalGeneratorBase { +public: + virtual ~ProposalGeneratorBase() = default; + + virtual void generate( + const std::vector ¤t_labels, + std::vector &proposal + ) = 0; + + virtual void reset() {} +}; + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx b/include/bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx new file mode 100644 index 0000000..4408f7f --- /dev/null +++ b/include/bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx @@ -0,0 +1,77 @@ +#pragma once + +#include "bioimage_cpp/graph/multicut/greedy_additive.hxx" +#include "bioimage_cpp/graph/proposal_generator.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +// Greedy-additive multicut proposal generator. Repeatedly invokes the greedy +// additive multicut solver with additive Gaussian noise on the edge weights; +// the seed is incremented per call so consecutive proposals differ. +// +// Multicut-specific in the sense that it solves a multicut to produce a +// proposal, but the *output* is just a node labeling and is reusable from any +// fusion-move driver (e.g. lifted multicut). +class GreedyAdditiveMulticutProposalGenerator final : public ProposalGeneratorBase { +public: + GreedyAdditiveMulticutProposalGenerator( + const UndirectedGraph &graph, + std::vector edge_costs, + const double sigma = 1.0, + const double weight_stop = 0.0, + const double node_num_stop = -1.0, + const int seed = 0 + ) + : graph_(graph), + edge_costs_(std::move(edge_costs)), + sigma_(sigma), + weight_stop_(weight_stop), + node_num_stop_(node_num_stop), + seed_(seed), + call_count_(0) { + if (edge_costs_.size() != static_cast(graph.number_of_edges())) { + throw std::invalid_argument( + "edge_costs length must equal number_of_edges" + ); + } + } + + void generate( + const std::vector ¤t_labels, + std::vector &proposal + ) override { + (void)current_labels; + proposal = multicut::greedy_additive( + graph_, + edge_costs_, + weight_stop_, + node_num_stop_, + true, + seed_ + static_cast(call_count_), + sigma_ + ); + ++call_count_; + } + + void reset() override { + call_count_ = 0; + } + +private: + const UndirectedGraph &graph_; + std::vector edge_costs_; + double sigma_; + double weight_stop_; + double node_num_stop_; + int seed_; + std::size_t call_count_; +}; + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/proposal_generators/watershed.hxx b/include/bioimage_cpp/graph/proposal_generators/watershed.hxx new file mode 100644 index 0000000..ecc154f --- /dev/null +++ b/include/bioimage_cpp/graph/proposal_generators/watershed.hxx @@ -0,0 +1,107 @@ +#pragma once + +#include "bioimage_cpp/graph/edge_weighted_watershed.hxx" +#include "bioimage_cpp/graph/proposal_generator.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +// Watershed proposal generator: noisy edge weights + random seeds at endpoints +// of negative-cost edges, then `edge_weighted_watershed`. Mirrors the workhorse +// generator used by nifty's multicut fusion moves. +// +// `n_seeds_fraction` is interpreted as a fraction of `number_of_nodes` when +// <= 1.0 and as an absolute seed-pair count otherwise. +class WatershedProposalGenerator final : public ProposalGeneratorBase { +public: + WatershedProposalGenerator( + const UndirectedGraph &graph, + std::vector edge_costs, + const double sigma = 1.0, + const double n_seeds_fraction = 0.1, + const int seed = 0 + ) + : graph_(graph), + edge_costs_(std::move(edge_costs)), + sigma_(sigma), + n_seeds_fraction_(n_seeds_fraction), + seed_(seed), + generator_(static_cast(seed)), + noise_(0.0, sigma) { + if (edge_costs_.size() != static_cast(graph.number_of_edges())) { + throw std::invalid_argument( + "edge_costs length must equal number_of_edges" + ); + } + for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { + if (edge_costs_[static_cast(edge)] < 0.0) { + negative_edges_.push_back(edge); + } + } + } + + void generate( + const std::vector ¤t_labels, + std::vector &proposal + ) override { + (void)current_labels; + const auto number_of_nodes = graph_.number_of_nodes(); + proposal.assign(static_cast(number_of_nodes), 0); + + if (negative_edges_.empty()) { + return; + } + + std::vector noisy_costs(edge_costs_.size()); + for (std::size_t edge = 0; edge < edge_costs_.size(); ++edge) { + noisy_costs[edge] = static_cast(edge_costs_[edge] + noise_(generator_)); + } + + std::size_t n_seed_pairs; + if (n_seeds_fraction_ <= 1.0) { + n_seed_pairs = static_cast( + static_cast(number_of_nodes) * n_seeds_fraction_ + 0.5 + ); + } else { + n_seed_pairs = static_cast(n_seeds_fraction_ + 0.5); + } + n_seed_pairs = std::max(std::size_t{1}, n_seed_pairs); + n_seed_pairs = std::min(negative_edges_.size(), n_seed_pairs); + + std::vector seeds(static_cast(number_of_nodes), 0); + std::uniform_int_distribution edge_dist(0, negative_edges_.size() - 1); + std::uint64_t next_label = 1; + for (std::size_t i = 0; i < n_seed_pairs; ++i) { + const auto edge = negative_edges_[edge_dist(generator_)]; + const auto uv = graph_.uv(edge); + seeds[static_cast(uv.first)] = next_label++; + seeds[static_cast(uv.second)] = next_label++; + } + + proposal = edge_weighted_watershed(graph_, noisy_costs, seeds); + } + + void reset() override { + generator_.seed(static_cast(seed_)); + } + +private: + const UndirectedGraph &graph_; + std::vector edge_costs_; + double sigma_; + double n_seeds_fraction_; + int seed_; + std::vector negative_edges_; + std::mt19937 generator_; + std::normal_distribution noise_; +}; + +} // namespace bioimage_cpp::graph diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 1a85395..ac8bc35 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -5,12 +5,20 @@ #include "bioimage_cpp/graph/edge_weighted_watershed.hxx" #include "bioimage_cpp/graph/feature_accumulation.hxx" #include "bioimage_cpp/graph/multicut.hxx" +#include "bioimage_cpp/graph/multicut/fusion_move.hxx" +#include "bioimage_cpp/graph/multicut/greedy_additive.hxx" +#include "bioimage_cpp/graph/multicut/greedy_fixation.hxx" +#include "bioimage_cpp/graph/multicut/kernighan_lin.hxx" #include "bioimage_cpp/graph/node_label_projection.hxx" +#include "bioimage_cpp/graph/proposal_generator.hxx" +#include "bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx" +#include "bioimage_cpp/graph/proposal_generators/watershed.hxx" #include "bioimage_cpp/graph/region_adjacency_graph.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" #include #include +#include #include #include @@ -406,6 +414,38 @@ UInt64Array multicut_kernighan_lin( return vector_to_uint64_array(label_vector); } +UInt64Array multicut_fusion_move( + const Graph &graph, + ConstDoubleArray costs, + ConstUInt64Array initial_labels, + graph::ProposalGeneratorBase *proposal_generator, + const graph::multicut::SolverBase *sub_solver, + const std::size_t number_of_iterations, + const std::size_t stop_if_no_improvement +) { + if (proposal_generator == nullptr) { + throw std::invalid_argument("proposal_generator must not be None"); + } + auto cost_vector = double_array_to_vector(costs, "edge_costs", graph.number_of_edges()); + auto label_vector = + uint64_array_to_vector(initial_labels, "initial_labels", graph.number_of_nodes()); + + graph::multicut::FusionMoveSolver solver( + *proposal_generator, + sub_solver, + number_of_iterations, + stop_if_no_improvement + ); + + std::vector result; + { + nb::gil_scoped_release release; + graph::multicut::Objective objective(graph, std::move(cost_vector), std::move(label_vector)); + result = solver.optimize(objective); + } + return vector_to_uint64_array(result); +} + template Rag region_adjacency_graph_t( LabelArray labels, @@ -711,6 +751,103 @@ void bind_graph(nb::module_ &m) { nb::arg("epsilon") ); + // Multicut sub-solver hierarchy used by fusion moves. The classes are + // opaque to Python; constructors carry per-solver settings. + nb::class_(m, "_MulticutSolverBase"); + nb::class_( + m, "_GreedyAdditiveMulticutSubSolver" + ) + .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_( + m, "_GreedyFixationMulticutSubSolver" + ) + .def( + nb::init(), + nb::arg("weight_stop") = 0.0, + nb::arg("node_num_stop") = -1.0 + ); + nb::class_( + m, "_KernighanLinMulticutSubSolver" + ) + .def( + nb::init(), + nb::arg("number_of_outer_iterations") = 100, + nb::arg("epsilon") = 1.0e-6 + ); + + // Proposal generators used by fusion moves. + nb::class_(m, "_ProposalGeneratorBase"); + nb::class_( + m, "_WatershedProposalGenerator" + ) + .def( + "__init__", + [](graph::WatershedProposalGenerator *self, + const Graph &graph, + ConstDoubleArray edge_costs, + double sigma, + double n_seeds_fraction, + int seed) { + auto costs = double_array_to_vector( + edge_costs, "edge_costs", graph.number_of_edges() + ); + new (self) graph::WatershedProposalGenerator( + graph, std::move(costs), sigma, n_seeds_fraction, seed + ); + }, + nb::arg("graph"), + nb::arg("edge_costs"), + nb::arg("sigma") = 1.0, + nb::arg("n_seeds_fraction") = 0.1, + nb::arg("seed") = 0 + ); + nb::class_< + graph::GreedyAdditiveMulticutProposalGenerator, + graph::ProposalGeneratorBase + >(m, "_GreedyAdditiveMulticutProposalGenerator") + .def( + "__init__", + [](graph::GreedyAdditiveMulticutProposalGenerator *self, + const Graph &graph, + ConstDoubleArray edge_costs, + double sigma, + double weight_stop, + double node_num_stop, + int seed) { + auto costs = double_array_to_vector( + edge_costs, "edge_costs", graph.number_of_edges() + ); + new (self) graph::GreedyAdditiveMulticutProposalGenerator( + graph, std::move(costs), sigma, weight_stop, node_num_stop, seed + ); + }, + nb::arg("graph"), + nb::arg("edge_costs"), + nb::arg("sigma") = 1.0, + nb::arg("weight_stop") = 0.0, + nb::arg("node_num_stop") = -1.0, + nb::arg("seed") = 0 + ); + + m.def( + "_multicut_fusion_move", + &multicut_fusion_move, + nb::arg("graph"), + nb::arg("edge_costs"), + nb::arg("initial_labels"), + nb::arg("proposal_generator"), + nb::arg("sub_solver").none(), + nb::arg("number_of_iterations"), + nb::arg("stop_if_no_improvement") + ); + m.def( "_region_adjacency_graph_uint32", ®ion_adjacency_graph_t, diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index 1916663..281711b 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -376,6 +376,15 @@ def optimize(self, objective: MulticutObjective) -> np.ndarray: objective.labels = labels return objective.labels + def _build_cpp_sub_solver(self): + return _core._GreedyAdditiveMulticutSubSolver( + weight_stop=self.weight_stop, + node_num_stop=self.node_num_stop, + add_noise=self.add_noise, + seed=self.seed, + sigma=self.sigma, + ) + class GreedyFixationMulticut(MulticutSolver): def __init__(self, *, weight_stop: float = 0.0, node_num_stop: float = -1.0): @@ -392,6 +401,12 @@ def optimize(self, objective: MulticutObjective) -> np.ndarray: objective.labels = labels return objective.labels + def _build_cpp_sub_solver(self): + return _core._GreedyFixationMulticutSubSolver( + weight_stop=self.weight_stop, + node_num_stop=self.node_num_stop, + ) + class KernighanLinMulticut(MulticutSolver): def __init__( @@ -434,6 +449,154 @@ def optimize(self, objective: MulticutObjective) -> np.ndarray: objective.labels = labels return objective.labels + def _build_cpp_sub_solver(self): + return _core._KernighanLinMulticutSubSolver( + number_of_outer_iterations=self.number_of_outer_iterations, + epsilon=self.epsilon, + ) + + +class ProposalGenerator(ABC): + """Base class for fusion-move proposal generators. + + Concrete generators carry settings on the Python side and build a C++ + proposal-generator object bound to a specific problem (graph + edge costs) + when ``_build`` is called by the fusion-move driver. + """ + + @abstractmethod + def _build(self, graph: UndirectedGraph | RegionAdjacencyGraph, edge_costs: np.ndarray): + """Construct the underlying C++ proposal generator.""" + + +class WatershedProposalGenerator(ProposalGenerator): + """Watershed proposal generator (nifty's fusion-move workhorse). + + Per call: add Gaussian noise to edge costs, drop random seeds at the + endpoints of negative-cost edges, run the edge-weighted watershed. + """ + + def __init__( + self, + *, + sigma: float = 1.0, + n_seeds_fraction: float = 0.1, + seed: int = 0, + ): + self.sigma = float(sigma) + self.n_seeds_fraction = float(n_seeds_fraction) + self.seed = int(seed) + + def _build(self, graph, edge_costs): + return _core._WatershedProposalGenerator( + graph, + edge_costs, + sigma=self.sigma, + n_seeds_fraction=self.n_seeds_fraction, + seed=self.seed, + ) + + +class GreedyAdditiveProposalGenerator(ProposalGenerator): + """Greedy-additive multicut proposal generator. + + Per call: run the greedy-additive multicut solver with noisy edge weights; + the seed advances every call so successive proposals differ. + """ + + def __init__( + self, + *, + sigma: float = 1.0, + weight_stop: float = 0.0, + node_num_stop: float = -1.0, + seed: int = 0, + ): + self.sigma = float(sigma) + self.weight_stop = float(weight_stop) + self.node_num_stop = float(node_num_stop) + self.seed = int(seed) + + def _build(self, graph, edge_costs): + return _core._GreedyAdditiveMulticutProposalGenerator( + graph, + edge_costs, + sigma=self.sigma, + weight_stop=self.weight_stop, + node_num_stop=self.node_num_stop, + seed=self.seed, + ) + + +class FusionMoveMulticut(MulticutSolver): + """Fusion-move multicut solver. + + Iteratively generates proposals via ``proposal_generator``, fuses them + pairwise with the current best labeling, and accepts improvements. The + fuse step solves a contracted multicut subproblem with ``sub_solver``; + if omitted, the default sub-solver is :class:`GreedyAdditiveMulticut`. + + If the objective's current labels are the trivial singleton labeling, the + driver warm-starts with one greedy-additive pass before the proposal loop. + The best-of safety net guarantees energy never increases across iterations. + + ``number_of_threads`` and ``number_of_parallel_proposals`` are reserved for + future use; the current implementation only supports the defaults. + """ + + def __init__( + self, + *, + proposal_generator: ProposalGenerator, + sub_solver: MulticutSolver | None = None, + number_of_iterations: int = 10, + stop_if_no_improvement: int = 4, + number_of_threads: int = 1, + number_of_parallel_proposals: int = 2, + ): + if not isinstance(proposal_generator, ProposalGenerator): + raise TypeError("proposal_generator must inherit from ProposalGenerator") + if sub_solver is not None and not isinstance(sub_solver, MulticutSolver): + raise TypeError("sub_solver must inherit from MulticutSolver") + if sub_solver is not None and not hasattr(sub_solver, "_build_cpp_sub_solver"): + raise TypeError( + "sub_solver must be a built-in multicut solver " + "(custom Python solvers are not supported as fusion-move sub-solvers)" + ) + if int(number_of_threads) != 1: + raise ValueError("number_of_threads must be 1 (parallel path not yet implemented)") + if int(number_of_parallel_proposals) != 2: + raise ValueError( + "number_of_parallel_proposals must be 2 (multi-proposal fuse not yet implemented)" + ) + self.proposal_generator = proposal_generator + self.sub_solver = sub_solver + self.number_of_iterations = int(number_of_iterations) + self.stop_if_no_improvement = int(stop_if_no_improvement) + self.number_of_threads = int(number_of_threads) + self.number_of_parallel_proposals = int(number_of_parallel_proposals) + if self.number_of_iterations < 0: + raise ValueError("number_of_iterations must be non-negative") + if self.stop_if_no_improvement < 1: + raise ValueError("stop_if_no_improvement must be >= 1") + + def optimize(self, objective: MulticutObjective) -> np.ndarray: + cpp_pgen = self.proposal_generator._build(objective.graph, objective.edge_costs) + cpp_sub_solver = ( + None if self.sub_solver is None else self.sub_solver._build_cpp_sub_solver() + ) + labels = _core._multicut_fusion_move( + objective.graph, + objective.edge_costs, + objective.labels, + cpp_pgen, + cpp_sub_solver, + self.number_of_iterations, + self.stop_if_no_improvement, + ) + objective.labels = labels + return objective.labels + class ChainedMulticutSolvers(MulticutSolver): def __init__(self, solvers): @@ -747,15 +910,19 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "COMPLEX_EDGE_FEATURE_NAMES", "DEFAULT_EXTERNAL_MULTICUT_PROBLEM_PATH", "EXTERNAL_MULTICUT_PROBLEM_URL", + "FusionMoveMulticut", "GreedyAdditiveMulticut", + "GreedyAdditiveProposalGenerator", "GreedyFixationMulticut", "KernighanLinMulticut", "MulticutDecomposer", "MulticutObjective", "MulticutSolver", + "ProposalGenerator", "RegionAdjacencyGraph", "SIMPLE_EDGE_FEATURE_NAMES", "UndirectedGraph", + "WatershedProposalGenerator", "affinity_features", "affinity_features_complex", "connected_components", diff --git a/tests/graph/multicut/test_fusion_move.py b/tests/graph/multicut/test_fusion_move.py new file mode 100644 index 0000000..aa33037 --- /dev/null +++ b/tests/graph/multicut/test_fusion_move.py @@ -0,0 +1,195 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + +from ._helpers import edge_cut_labels + + +def test_fuses_a_two_cluster_problem(external_toy_problem): + graph, costs, _ = external_toy_problem + objective = bic.graph.MulticutObjective(graph, costs) + baseline = bic.graph.GreedyAdditiveMulticut().optimize(objective) + baseline_energy = objective.energy(baseline) + + objective.reset_labels() + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=1), + number_of_iterations=10, + ) + labels = solver.optimize(objective) + + assert labels.dtype == np.uint64 + assert labels.shape == (graph.number_of_nodes,) + assert objective.energy(labels) <= baseline_energy + 1e-9 + + +def test_safety_net_never_regresses(grid_problem): + graph, costs = grid_problem + objective = bic.graph.MulticutObjective(graph, costs) + baseline_energy = objective.energy( + bic.graph.GreedyAdditiveMulticut().optimize(objective) + ) + + objective.reset_labels() + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=7), + number_of_iterations=8, + ) + labels = solver.optimize(objective) + + assert objective.energy(labels) <= baseline_energy + 1e-9 + + +def test_warm_starts_from_singleton(chain_problem): + graph, costs = chain_problem + # The default objective labels are the singleton (arange) labeling; the + # driver should warm-start with greedy-additive and reach the optimum. + objective = bic.graph.MulticutObjective(graph, costs) + + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=0), + number_of_iterations=3, + ) + labels = solver.optimize(objective) + + # Optimal: cut only the negative edge between nodes 2 and 3. + expected_cut = np.array([False, False, True]) + np.testing.assert_array_equal(edge_cut_labels(graph, labels), expected_cut) + + +def test_greedy_additive_proposal_generator_runs(grid_problem): + graph, costs = grid_problem + objective = bic.graph.MulticutObjective(graph, costs) + baseline = objective.energy(bic.graph.GreedyAdditiveMulticut().optimize(objective)) + + objective.reset_labels() + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.GreedyAdditiveProposalGenerator(seed=3, sigma=1.0), + number_of_iterations=5, + ) + labels = solver.optimize(objective) + assert objective.energy(labels) <= baseline + 1e-9 + + +@pytest.mark.parametrize( + "sub_solver", + [ + bic.graph.GreedyAdditiveMulticut(), + bic.graph.GreedyFixationMulticut(), + bic.graph.KernighanLinMulticut(number_of_outer_iterations=3), + ], +) +def test_sub_solver_pluggability(grid_problem, sub_solver): + graph, costs = grid_problem + objective = bic.graph.MulticutObjective(graph, costs) + + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=2), + sub_solver=sub_solver, + number_of_iterations=4, + ) + labels = solver.optimize(objective) + + assert labels.shape == (graph.number_of_nodes,) + assert labels.dtype == np.uint64 + + +def test_stops_after_no_improvement(frustrated_triangle): + # Tiny problem with many iterations and an aggressive non-improvement + # threshold: the loop must terminate quickly without scanning all 1000 + # iterations (the test would time out otherwise). + graph, costs = frustrated_triangle + objective = bic.graph.MulticutObjective(graph, costs) + baseline = objective.energy( + bic.graph.GreedyAdditiveMulticut().optimize(objective) + ) + + objective.reset_labels() + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=0), + number_of_iterations=1000, + stop_if_no_improvement=1, + ) + labels = solver.optimize(objective) + # Best-of safety net guarantees energy never regresses past baseline. + assert objective.energy(labels) <= baseline + 1e-9 + + +def test_chains_with_kernighan_lin(external_toy_problem): + graph, costs, _ = external_toy_problem + objective = bic.graph.MulticutObjective(graph, costs) + baseline = objective.energy( + bic.graph.GreedyAdditiveMulticut().optimize(objective) + ) + + objective.reset_labels() + solver = bic.graph.ChainedMulticutSolvers([ + bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=11), + number_of_iterations=5, + ), + bic.graph.KernighanLinMulticut(number_of_outer_iterations=3), + ]) + labels = solver.optimize(objective) + assert objective.energy(labels) <= baseline + 1e-9 + + +def test_rejects_non_proposal_generator(): + with pytest.raises(TypeError, match="proposal_generator"): + bic.graph.FusionMoveMulticut(proposal_generator=object()) + + +def test_rejects_unsupported_thread_count(): + with pytest.raises(ValueError, match="number_of_threads"): + bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + number_of_threads=4, + ) + + +def test_rejects_unsupported_parallel_proposals(): + with pytest.raises(ValueError, match="number_of_parallel_proposals"): + bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + number_of_parallel_proposals=4, + ) + + +def test_rejects_invalid_iteration_settings(): + with pytest.raises(ValueError, match="number_of_iterations"): + bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + number_of_iterations=-1, + ) + with pytest.raises(ValueError, match="stop_if_no_improvement"): + bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + stop_if_no_improvement=0, + ) + + +def test_runs_on_empty_graph(): + graph = bic.graph.UndirectedGraph(0) + objective = bic.graph.MulticutObjective(graph, np.zeros(0)) + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + ) + labels = solver.optimize(objective) + assert labels.shape == (0,) + + +def test_runs_on_graph_without_negative_edges(chain_problem): + # WatershedProposalGenerator yields all-zero proposals if no negative + # edges exist; the driver must still terminate cleanly and produce the + # warm-started greedy-additive result. + graph, _ = chain_problem + costs = np.array([1.0, 2.0, 3.0], dtype=np.float64) + objective = bic.graph.MulticutObjective(graph, costs) + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=0), + number_of_iterations=3, + ) + labels = solver.optimize(objective) + # All-positive costs → no cut. + assert np.all(edge_cut_labels(graph, labels) == False) # noqa: E712 diff --git a/tests/graph/multicut/test_proposal_generators.py b/tests/graph/multicut/test_proposal_generators.py new file mode 100644 index 0000000..103c132 --- /dev/null +++ b/tests/graph/multicut/test_proposal_generators.py @@ -0,0 +1,67 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic +from bioimage_cpp import _core + + +def test_watershed_proposal_generator_is_deterministic_given_seed(): + # Two generators with the same seed must produce the same proposal series. + graph = bic.graph.UndirectedGraph.from_edges( + 6, + [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [0, 3]], + ) + costs = np.array([-1.0, 1.0, -1.0, 1.0, -1.0, 0.5], dtype=np.float64) + + objective_a = bic.graph.MulticutObjective(graph, costs) + objective_b = bic.graph.MulticutObjective(graph, costs) + + solver_a = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=42), + number_of_iterations=3, + ) + solver_b = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=42), + number_of_iterations=3, + ) + + labels_a = solver_a.optimize(objective_a) + labels_b = solver_b.optimize(objective_b) + np.testing.assert_array_equal(labels_a, labels_b) + + +def test_greedy_additive_proposal_generator_is_deterministic_given_seed(grid_problem): + graph, costs = grid_problem + objective_a = bic.graph.MulticutObjective(graph, costs) + objective_b = bic.graph.MulticutObjective(graph, costs) + + solver_a = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.GreedyAdditiveProposalGenerator(seed=5), + number_of_iterations=3, + ) + solver_b = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.GreedyAdditiveProposalGenerator(seed=5), + number_of_iterations=3, + ) + np.testing.assert_array_equal( + solver_a.optimize(objective_a), + solver_b.optimize(objective_b), + ) + + +def test_cpp_proposal_generator_validates_costs_length(): + graph = bic.graph.UndirectedGraph.from_edges(3, [[0, 1], [1, 2]]) + bad_costs = np.array([0.5], dtype=np.float64) + with pytest.raises((ValueError, RuntimeError)): + _core._WatershedProposalGenerator(graph, bad_costs) + with pytest.raises((ValueError, RuntimeError)): + _core._GreedyAdditiveMulticutProposalGenerator(graph, bad_costs) + + +def test_proposal_generator_isinstance_check(): + assert isinstance( + bic.graph.WatershedProposalGenerator(), bic.graph.ProposalGenerator + ) + assert isinstance( + bic.graph.GreedyAdditiveProposalGenerator(), bic.graph.ProposalGenerator + ) From c96541e959d2352bd33e101c80ba69bb08d799cd Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Fri, 15 May 2026 10:37:37 +0200 Subject: [PATCH 3/5] Add comparison to nifty --- .../graph/multicut/check_fusion_move.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 development/graph/multicut/check_fusion_move.py diff --git a/development/graph/multicut/check_fusion_move.py b/development/graph/multicut/check_fusion_move.py new file mode 100644 index 0000000..102d9c9 --- /dev/null +++ b/development/graph/multicut/check_fusion_move.py @@ -0,0 +1,38 @@ +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 fusion-move multicut at default settings." + ).parse_args() + # bioimage-cpp warm-starts from the trivial singleton labeling with a + # greedy-additive pass before the proposal loop. Nifty's ccFusionMoveBased + # exposes the same behaviour via the `warmStartGreedy=True` flag (which + # internally chains a greedyAdditiveFactory in front of the fusion-move + # factory). Both sides therefore enter the proposal loop from the same + # starting point. + run_comparison( + "fusion_move", + lambda: bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + ), + lambda objective: objective.ccFusionMoveBasedFactory( + proposalGenerator=objective.watershedCcProposals(), + fusionMove=objective.fusionMoveSettings( + mcFactory=objective.greedyAdditiveFactory(), + ), + numberOfIterations=10, + stopIfNoImprovement=4, + numberOfThreads=1, + warmStartGreedy=True, + ), + args, + ) + + +if __name__ == "__main__": + main() From cba92f7f9a39ee5c50dc1ac3e678b418c5292879 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Fri, 15 May 2026 11:31:32 +0200 Subject: [PATCH 4/5] Optimize fusion moves --- AGENTS.md | 19 +++ CMakeLists.txt | 5 + include/bioimage_cpp/detail/indexed_heap.hxx | 22 +++ include/bioimage_cpp/detail/profile.hxx | 82 +++++++++++ include/bioimage_cpp/detail/union_find.hxx | 10 ++ .../graph/detail/fusion_contract.hxx | 94 ++++++++++-- .../graph/edge_weighted_watershed.hxx | 138 ++++++++++++------ .../bioimage_cpp/graph/multicut/detail.hxx | 39 ++++- .../graph/multicut/fusion_move.hxx | 116 +++++++++++---- .../graph/multicut/greedy_additive.hxx | 41 +++++- .../graph/proposal_generators/watershed.hxx | 26 +++- .../bioimage_cpp/graph/undirected_graph.hxx | 51 +++++++ 12 files changed, 538 insertions(+), 105 deletions(-) create mode 100644 include/bioimage_cpp/detail/profile.hxx diff --git a/AGENTS.md b/AGENTS.md index 09b7ee7..82f3c25 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -120,6 +120,25 @@ Tests run against the installed Python package. For each public function, cover: Correct first, fast second. Benchmark before adding complexity. Prefer algorithmic improvements over build-level tweaks. Release the GIL for expensive kernels. Add threading only after the single-threaded implementation is stable; it must be portable and user-controllable. +### Profiling + +Measure before optimizing. The codebase carries a lightweight per-phase profiling utility for exactly this: + +- Header: `include/bioimage_cpp/detail/profile.hxx`. Macros: `BIOIMAGE_PROFILE_INIT(name)`, `BIOIMAGE_PROFILE_SCOPE(name, "label")`, `BIOIMAGE_PROFILE_REPORT(name)`. +- Gated behind the `BIOIMAGE_PROFILE` compile-time flag. Outside of profile builds the macros expand to no-ops and a `NullProfiler` stub so the same code compiles unchanged. +- Enable via CMake option: `pip install -e . --no-build-isolation -C cmake.define.BIOIMAGE_PROFILE=ON`. Rebuild without the flag for production work. +- Reports per-phase wall-clock totals to stderr at the end of the instrumented scope (e.g., at the end of `optimize`). Same labels accumulate across multiple invocations of the scope (e.g., per-iteration phases). + +Workflow when adding or chasing a performance issue: + +1. **Compare standalone primitives first** (`development/.../check_*.py` scripts vs. nifty). If a primitive is already fast, the gap is elsewhere — don't optimize it speculatively. +2. **Instrument the suspect function** by wrapping each logical phase in a `BIOIMAGE_PROFILE_SCOPE`. Pick labels that map to one operation each ("agreement_contract", "sub_solve", "energy_eval", ...), not full call paths. +3. **Build with `BIOIMAGE_PROFILE=ON` and run a realistic problem** — typically the external multicut instance loaded by the comparison scripts. Run with `--repeats 1` so the report isn't drowned out. +4. **Optimize the largest phase** (50% phase beats two 10% phases combined). Re-measure after each change; verify no other phase regressed. +5. **Strip the instrumentation when done** only if it adds clutter; otherwise leave it in place — it's free when the flag is off. + +Don't add `std::chrono` snippets ad hoc; use the existing macros so future profiling sessions land in a consistent format. + ## Documentation `README.md` covers: what `bioimage-cpp` is and isn't, install/build, minimal examples, design philosophy. Public functions need concise docstrings documenting input shapes, supported dtypes, output shapes/dtypes, copy behavior, background-label behavior (if relevant), and axis/coordinate conventions. diff --git a/CMakeLists.txt b/CMakeLists.txt index 1ee5a7d..e56cde5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,4 +27,9 @@ target_compile_options(_core PRIVATE $<$>:-O3> ) +option(BIOIMAGE_PROFILE "Enable per-phase profiling instrumentation (development only)" OFF) +if(BIOIMAGE_PROFILE) + target_compile_definitions(_core PRIVATE BIOIMAGE_PROFILE) +endif() + install(TARGETS _core LIBRARY DESTINATION bioimage_cpp) diff --git a/include/bioimage_cpp/detail/indexed_heap.hxx b/include/bioimage_cpp/detail/indexed_heap.hxx index 82de272..cb40745 100644 --- a/include/bioimage_cpp/detail/indexed_heap.hxx +++ b/include/bioimage_cpp/detail/indexed_heap.hxx @@ -133,6 +133,28 @@ public: sift_up(pos); } + // Bulk-load the heap from a list of entries in O(n) via Floyd's heapify. + // Replaces any current heap contents. Preconditions: + // - the heap is empty before this call (call `clear()`/`reset_capacity` + // first if needed); + // - every key in `entries` is unique. + // Compared to N successive `push` calls (O(n log n)) this matters when + // initializing a heap from a large edge set in one shot. + void build_heap(std::vector entries) { + heap_ = std::move(entries); + for (std::size_t pos = 0; pos < heap_.size(); ++pos) { + locator_.set(heap_[pos].key, pos); + } + if (heap_.size() < 2) { + return; + } + // Sift down from the last internal node back to the root. After each + // sift_down, the subtree rooted at `pos` is heap-ordered. + for (std::size_t pos = heap_.size() / 2; pos-- > 0;) { + sift_down(pos); + } + } + // Precondition: `key` is currently in the heap. void change(const KeyT &key, PriorityT priority) { const auto pos = locator_.at(key); diff --git a/include/bioimage_cpp/detail/profile.hxx b/include/bioimage_cpp/detail/profile.hxx new file mode 100644 index 0000000..ba39c54 --- /dev/null +++ b/include/bioimage_cpp/detail/profile.hxx @@ -0,0 +1,82 @@ +#pragma once + +// Lightweight per-phase profiling utility. Active only when the translation +// unit is compiled with `-DBIOIMAGE_PROFILE`; otherwise every macro is a no-op +// and adds no overhead. Use it in development/profile-mode builds to find +// hotspots; do not enable it in production wheels. + +#ifdef BIOIMAGE_PROFILE +#include +#include +#include +#include +#include + +namespace bioimage_cpp::detail { + +class Profiler { +public: + using Clock = std::chrono::steady_clock; + using Duration = std::chrono::duration; + + void accumulate(const char *name, const Duration delta) { + auto it = totals_.find(name); + if (it == totals_.end()) { + order_.push_back(name); + totals_.emplace(name, delta.count()); + } else { + it->second += delta.count(); + } + } + + void report() const { + std::fprintf(stderr, "[bioimage profile]\n"); + double total = 0.0; + for (const auto *name : order_) { + total += totals_.at(name); + } + for (const auto *name : order_) { + const auto seconds = totals_.at(name); + const auto fraction = total > 0.0 ? (100.0 * seconds / total) : 0.0; + std::fprintf(stderr, " %-22s %8.4f s (%5.1f%%)\n", name, seconds, fraction); + } + std::fprintf(stderr, " %-22s %8.4f s\n", "total", total); + } + +private: + std::vector order_; + std::unordered_map totals_; +}; + +class ProfileTimer { +public: + ProfileTimer(Profiler &profiler, const char *name) + : profiler_(profiler), name_(name), start_(Profiler::Clock::now()) {} + + ~ProfileTimer() { + profiler_.accumulate(name_, Profiler::Clock::now() - start_); + } + +private: + Profiler &profiler_; + const char *name_; + Profiler::Clock::time_point start_; +}; + +} // namespace bioimage_cpp::detail + +#define BIOIMAGE_PROFILE_INIT(var) ::bioimage_cpp::detail::Profiler var; +#define BIOIMAGE_PROFILE_SCOPE(var, name) ::bioimage_cpp::detail::ProfileTimer _bp_##__LINE__(var, name); +#define BIOIMAGE_PROFILE_REPORT(var) (var).report(); + +#else + +namespace bioimage_cpp::detail { +struct NullProfiler {}; +} // namespace bioimage_cpp::detail + +#define BIOIMAGE_PROFILE_INIT(var) ::bioimage_cpp::detail::NullProfiler var; +#define BIOIMAGE_PROFILE_SCOPE(var, name) (void)var; +#define BIOIMAGE_PROFILE_REPORT(var) (void)var; + +#endif diff --git a/include/bioimage_cpp/detail/union_find.hxx b/include/bioimage_cpp/detail/union_find.hxx index f05fa4e..404b9a2 100644 --- a/include/bioimage_cpp/detail/union_find.hxx +++ b/include/bioimage_cpp/detail/union_find.hxx @@ -70,6 +70,16 @@ public: return parents_.size(); } + // Re-initialise the union-find to `n` singletons. Reuses the existing + // vectors' capacity when possible, so the workspace pattern (reusing a + // single `UnionFind` across many solver invocations on graphs of varying + // size) avoids reallocations. + void reset(const std::size_t n) { + parents_.resize(n); + ranks_.assign(n, 0); + std::iota(parents_.begin(), parents_.end(), std::uint64_t{0}); + } + private: std::vector parents_; std::vector ranks_; diff --git a/include/bioimage_cpp/graph/detail/fusion_contract.hxx b/include/bioimage_cpp/graph/detail/fusion_contract.hxx index f226bb5..c880aa6 100644 --- a/include/bioimage_cpp/graph/detail/fusion_contract.hxx +++ b/include/bioimage_cpp/graph/detail/fusion_contract.hxx @@ -1,13 +1,15 @@ #pragma once -#include "bioimage_cpp/detail/relabel.hxx" #include "bioimage_cpp/detail/union_find.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" +#include #include #include +#include #include #include +#include #include namespace bioimage_cpp::graph::detail { @@ -30,6 +32,17 @@ struct AgreementContraction { // `proposals` is a row-major buffer of shape (n_proposals, number_of_nodes). // Passing 0 proposals collapses every edge (all proposals trivially agree) // and yields a single-node contracted graph; we reject that as a usage error. +// +// Implementation notes: +// - Dense relabeling uses a flat sentinel array indexed by root id rather +// than a hash map; roots live in [0, number_of_nodes) so this is O(N). +// - Surviving (min_root, max_root) pairs are collected and sorted by a +// packed 64-bit key, then deduped sequentially. The contracted graph is +// built in one pass via `UndirectedGraph::from_sorted_unique_edges`, +// bypassing per-edge hash insertion in `insert_edge`. +// - The contracted graph's `edge_lookup_` is left empty because the +// fusion-move sub-solver only iterates edges and adjacency, never calls +// `find_edge`. inline AgreementContraction contract_by_agreement( const UndirectedGraph &graph, const std::uint64_t *proposals, @@ -67,36 +80,87 @@ inline AgreementContraction contract_by_agreement( } } - std::vector raw_root(static_cast(number_of_nodes)); + // Dense-relabel UFD roots in one O(N) pass with a sentinel array. + constexpr std::uint64_t unset = std::numeric_limits::max(); + std::vector dense_of_raw( + static_cast(number_of_nodes), unset + ); + std::vector root_of_node(static_cast(number_of_nodes)); + std::uint64_t number_of_components = 0; for (std::uint64_t node = 0; node < number_of_nodes; ++node) { - raw_root[static_cast(node)] = sets.find(node); + const auto raw = sets.find(node); + auto dense = dense_of_raw[static_cast(raw)]; + if (dense == unset) { + dense = number_of_components++; + dense_of_raw[static_cast(raw)] = dense; + } + root_of_node[static_cast(node)] = dense; } - auto root_of_node = bioimage_cpp::detail::dense_relabel(raw_root); - std::uint64_t number_of_components = 0; - for (const auto root : root_of_node) { - if (root + 1 > number_of_components) { - number_of_components = root + 1; - } + // Sort key for surviving edges. The lower 32 bits hold `max_root` and + // the upper 32 bits hold `min_root` so a single uint64 comparison + // suffices. This requires number_of_components to fit in 32 bits, which + // is always true for graphs we can fit in memory. + if (number_of_components > (std::uint64_t{1} << 32)) { + throw std::runtime_error( + "number_of_components exceeds 2^32 — contraction packing assumption violated" + ); } - UndirectedGraph contracted_graph(number_of_components); + struct Survivor { + std::uint64_t key; // (min_root << 32) | max_root + std::uint64_t original_edge; + }; + std::vector survivors; + survivors.reserve(static_cast(number_of_edges)); + std::vector contracted_edge_of_original( static_cast(number_of_edges), -1 ); for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { const auto uv = graph.uv(edge); - const auto ru = root_of_node[static_cast(uv.first)]; - const auto rv = root_of_node[static_cast(uv.second)]; + auto ru = root_of_node[static_cast(uv.first)]; + auto rv = root_of_node[static_cast(uv.second)]; if (ru == rv) { continue; } - const auto inserted = contracted_graph.insert_edge(ru, rv); - contracted_edge_of_original[static_cast(edge)] = - static_cast(inserted); + if (ru > rv) { + std::swap(ru, rv); + } + survivors.push_back(Survivor{(ru << 32) | rv, edge}); } + std::sort( + survivors.begin(), + survivors.end(), + [](const Survivor &a, const Survivor &b) { + return a.key < b.key; + } + ); + + std::vector contracted_edges; + contracted_edges.reserve(survivors.size()); + + constexpr std::uint64_t no_key = std::numeric_limits::max(); + std::uint64_t last_key = no_key; + std::int64_t current_contracted = -1; + for (const auto &survivor : survivors) { + if (survivor.key != last_key) { + const auto ru = survivor.key >> 32; + const auto rv = survivor.key & std::uint64_t{0xFFFFFFFF}; + current_contracted = static_cast(contracted_edges.size()); + contracted_edges.push_back(UndirectedGraph::Edge{ru, rv}); + last_key = survivor.key; + } + contracted_edge_of_original[static_cast(survivor.original_edge)] = + current_contracted; + } + + auto contracted_graph = UndirectedGraph::from_sorted_unique_edges( + number_of_components, std::move(contracted_edges), /*populate_lookup=*/false + ); + return AgreementContraction{ std::move(contracted_graph), std::move(contracted_edge_of_original), diff --git a/include/bioimage_cpp/graph/edge_weighted_watershed.hxx b/include/bioimage_cpp/graph/edge_weighted_watershed.hxx index 29c0506..253a5aa 100644 --- a/include/bioimage_cpp/graph/edge_weighted_watershed.hxx +++ b/include/bioimage_cpp/graph/edge_weighted_watershed.hxx @@ -9,10 +9,75 @@ #include #include #include +#include #include namespace bioimage_cpp::graph { +namespace detail { + +template +struct EdgeWeightedWatershedScratch { + std::vector> sort_buffer; +}; + +template +inline void edge_weighted_watershed_kernel( + const UndirectedGraph &graph, + const WeightT *edge_weights, + const std::uint64_t number_of_edges, + SeedT *labels, + const std::uint64_t number_of_nodes, + EdgeWeightedWatershedScratch &scratch +) { + if (number_of_edges == 0) { + return; + } + + scratch.sort_buffer.resize(static_cast(number_of_edges)); + for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { + scratch.sort_buffer[static_cast(edge)] = { + edge_weights[static_cast(edge)], edge + }; + } + // The pair's first element is the weight, so the default less-than gives + // ascending sort order. `stable_sort` keeps the existing tie-breaking + // semantics (matches the docstring's "ascending by weight" guarantee). + std::stable_sort( + scratch.sort_buffer.begin(), + scratch.sort_buffer.end(), + [](const auto &a, const auto &b) { return a.first < b.first; } + ); + + bioimage_cpp::detail::UnionFind sets(static_cast(number_of_nodes)); + + constexpr SeedT zero{0}; + for (const auto &item : scratch.sort_buffer) { + const auto edge = item.second; + const auto uv = graph.uv(edge); + const auto ru = sets.find(uv.first); + const auto rv = sets.find(uv.second); + if (ru == rv) { + continue; + } + const auto lu = labels[static_cast(ru)]; + const auto lv = labels[static_cast(rv)]; + if (lu != zero && lv != zero) { + continue; + } + const auto new_label = std::max(lu, lv); + const auto new_root = sets.unite_roots(ru, rv); + labels[static_cast(new_root)] = new_label; + } + + for (std::uint64_t node = 0; node < number_of_nodes; ++node) { + const auto root = sets.find(node); + labels[static_cast(node)] = labels[static_cast(root)]; + } +} + +} // namespace detail + // Kruskal-style edge-weighted seeded watershed. // // Edges are visited in ascending weight order. Two distinct components are @@ -56,50 +121,39 @@ inline std::vector edge_weighted_watershed( } std::vector labels(seeds); - - if (number_of_edges == 0) { - return labels; - } - - std::vector order(static_cast(number_of_edges)); - for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { - order[static_cast(edge)] = edge; - } - std::stable_sort( - order.begin(), - order.end(), - [&edge_weights](const std::uint64_t a, const std::uint64_t b) { - return edge_weights[static_cast(a)] < - edge_weights[static_cast(b)]; - } + detail::EdgeWeightedWatershedScratch scratch; + detail::edge_weighted_watershed_kernel( + graph, + edge_weights.data(), + number_of_edges, + labels.data(), + number_of_nodes, + scratch ); - - detail::UnionFind sets(static_cast(number_of_nodes)); - - constexpr SeedT zero{0}; - for (const auto edge : order) { - const auto uv = graph.uv(edge); - const auto ru = sets.find(uv.first); - const auto rv = sets.find(uv.second); - if (ru == rv) { - continue; - } - const auto lu = labels[static_cast(ru)]; - const auto lv = labels[static_cast(rv)]; - if (lu != zero && lv != zero) { - continue; - } - const auto new_label = std::max(lu, lv); - const auto new_root = sets.unite_roots(ru, rv); - labels[static_cast(new_root)] = new_label; - } - - for (std::uint64_t node = 0; node < number_of_nodes; ++node) { - const auto root = sets.find(node); - labels[static_cast(node)] = labels[static_cast(root)]; - } - return labels; } +// Buffer-reusing variant: caller provides scratch + output buffers. Validates +// nothing; intended for tight inner loops (e.g. proposal generators that call +// the watershed once per iteration on the same graph). The caller must: +// - size `labels` to `graph.number_of_nodes()` and copy seeds into it before +// the call (the algorithm propagates labels in place on this buffer); +// - ensure `edge_weights` has length `graph.number_of_edges()`. +template +inline void edge_weighted_watershed_into( + const UndirectedGraph &graph, + const std::vector &edge_weights, + std::vector &labels, + detail::EdgeWeightedWatershedScratch &scratch +) { + detail::edge_weighted_watershed_kernel( + graph, + edge_weights.data(), + graph.number_of_edges(), + labels.data(), + graph.number_of_nodes(), + scratch + ); +} + } // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/multicut/detail.hxx b/include/bioimage_cpp/graph/multicut/detail.hxx index 76ab5c9..7d32d66 100644 --- a/include/bioimage_cpp/graph/multicut/detail.hxx +++ b/include/bioimage_cpp/graph/multicut/detail.hxx @@ -41,16 +41,33 @@ struct DynamicEdge { using EdgeHeap = bioimage_cpp::detail::DenseIndexedHeap; struct DynamicGraph { - explicit DynamicGraph(const UndirectedGraph &graph) - : adjacency(static_cast(graph.number_of_nodes())), - alive(static_cast(graph.number_of_nodes()), true), - alive_count(static_cast(graph.number_of_nodes())), - scratch_edge_id(static_cast(graph.number_of_nodes()), no_edge) { + DynamicGraph() = default; + + explicit DynamicGraph(const UndirectedGraph &graph) { + reset(graph); + } + + // Reuse the buffers of an existing DynamicGraph for a new input graph. + // Inner adjacency vectors are `clear()`-ed (keeping capacity) and the + // outer container is resized; degree-based reserves prevent any per-edge + // adjacency growth during initialise. + void reset(const UndirectedGraph &graph) { + const auto n_nodes = static_cast(graph.number_of_nodes()); + const auto n_edges = static_cast(graph.number_of_edges()); + + for (auto &adj : adjacency) { + adj.clear(); + } + adjacency.resize(n_nodes); for (std::uint64_t node = 0; node < graph.number_of_nodes(); ++node) { const auto degree = graph.node_adjacency(node).size(); adjacency[static_cast(node)].reserve(degree); } - edges.resize(static_cast(graph.number_of_edges())); + + alive.assign(n_nodes, true); + alive_count = n_nodes; + scratch_edge_id.assign(n_nodes, no_edge); + edges.resize(n_edges); } // O(degree(u)). Returns no_edge when (u, v) is not an edge. @@ -102,6 +119,9 @@ inline void initialize_dynamic_graph( const auto n_edges = static_cast(graph.number_of_edges()); heap.reset_capacity(n_edges); + std::vector heap_entries; + heap_entries.reserve(n_edges); + std::mt19937 generator(seed); std::normal_distribution noise(0.0, sigma); for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { @@ -120,8 +140,13 @@ inline void initialize_dynamic_graph( e.is_constraint = 0; dynamic_graph.adjacency[u].push_back({v, edge_id}); dynamic_graph.adjacency[v].push_back({u, edge_id}); - heap.push(edge_id, priority_for(weight, absolute_priority)); + heap_entries.push_back({edge_id, priority_for(weight, absolute_priority)}); } + + // Floyd's heapify is O(n_edges) — meaningfully faster than n_edges + // successive `push` calls when n_edges is in the hundreds of thousands + // (the common case for graphs we run the multicut on). + heap.build_heap(std::move(heap_entries)); } inline std::vector labels_from_sets( diff --git a/include/bioimage_cpp/graph/multicut/fusion_move.hxx b/include/bioimage_cpp/graph/multicut/fusion_move.hxx index f59a753..51ca220 100644 --- a/include/bioimage_cpp/graph/multicut/fusion_move.hxx +++ b/include/bioimage_cpp/graph/multicut/fusion_move.hxx @@ -1,5 +1,6 @@ #pragma once +#include "bioimage_cpp/detail/profile.hxx" #include "bioimage_cpp/graph/detail/fusion_contract.hxx" #include "bioimage_cpp/graph/multicut/greedy_additive.hxx" #include "bioimage_cpp/graph/multicut/objective.hxx" @@ -50,6 +51,7 @@ public: } std::vector optimize(Objective &objective) const override { + BIOIMAGE_PROFILE_INIT(profile); const auto &graph = objective.graph(); const auto &costs = objective.costs(); const auto number_of_nodes = graph.number_of_nodes(); @@ -60,29 +62,49 @@ public: return objective.labels(); } - const auto default_sub_solver = GreedyAdditiveSolver(); - const auto &sub_solver = sub_solver_ != nullptr - ? *sub_solver_ - : static_cast(default_sub_solver); + // Workspace reused across the warm-start, every fuse iteration's + // sub-solve, and (transitively) any callers chaining additional + // greedy-additive runs. + GreedyAdditiveWorkspace greedy_workspace; // Warm start from greedy-additive if the caller passed the trivial // singleton labeling. if (is_singleton_labeling(current)) { - Objective warm_objective(graph, costs); - current = default_sub_solver.optimize(warm_objective); + BIOIMAGE_PROFILE_SCOPE(profile, "warm_start"); + current = greedy_additive( + graph, costs, 0.0, -1.0, false, 42, 1.0, greedy_workspace + ); } - double current_energy = energy(graph, costs, current); + double current_energy; + { + BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); + current_energy = energy(graph, costs, current); + } std::vector proposal(static_cast(number_of_nodes)); std::size_t iterations_without_improvement = 0; for (std::size_t iteration = 0; iteration < number_of_iterations_; ++iteration) { - proposal_generator_.generate(current, proposal); + { + BIOIMAGE_PROFILE_SCOPE(profile, "proposal"); + proposal_generator_.generate(current, proposal); + } - const auto fused = fuse_pair(graph, costs, current, proposal, sub_solver); - const auto fused_energy = energy(graph, costs, fused); - const auto proposal_energy = energy(graph, costs, proposal); + std::vector fused = fuse_pair( + graph, costs, current, proposal, sub_solver_, greedy_workspace, profile + ); + + double fused_energy; + { + BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); + fused_energy = energy(graph, costs, fused); + } + double proposal_energy; + { + BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); + proposal_energy = energy(graph, costs, proposal); + } // Best-of safety net across the three candidates so an iteration // can never raise the running energy. @@ -110,6 +132,7 @@ public: } objective.set_labels(current); + BIOIMAGE_PROFILE_REPORT(profile); return objective.labels(); } @@ -123,21 +146,28 @@ private: return true; } + template static std::vector fuse_pair( const UndirectedGraph &graph, const std::vector &costs, const std::vector ¤t, const std::vector &proposal, - const SolverBase &sub_solver + const SolverBase *sub_solver, + GreedyAdditiveWorkspace &greedy_workspace, + [[maybe_unused]] ProfilerT &profile ) { const auto number_of_nodes = static_cast(graph.number_of_nodes()); std::vector stacked(2 * number_of_nodes); std::copy(current.begin(), current.end(), stacked.begin()); std::copy(proposal.begin(), proposal.end(), stacked.begin() + number_of_nodes); - auto contraction = ::bioimage_cpp::graph::detail::contract_by_agreement( - graph, stacked.data(), 2, number_of_nodes - ); + ::bioimage_cpp::graph::detail::AgreementContraction contraction; + { + BIOIMAGE_PROFILE_SCOPE(profile, "agreement_contract"); + contraction = ::bioimage_cpp::graph::detail::contract_by_agreement( + graph, stacked.data(), 2, number_of_nodes + ); + } const auto &contracted_graph = contraction.contracted_graph; const auto number_of_contracted_edges = contracted_graph.number_of_edges(); @@ -145,15 +175,18 @@ private: std::vector contracted_costs( static_cast(number_of_contracted_edges), 0.0 ); - for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { - const auto target = contraction.contracted_edge_of_original[ - static_cast(edge) - ]; - if (target < 0) { - continue; + { + BIOIMAGE_PROFILE_SCOPE(profile, "cost_aggregate"); + for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { + const auto target = contraction.contracted_edge_of_original[ + static_cast(edge) + ]; + if (target < 0) { + continue; + } + contracted_costs[static_cast(target)] += + costs[static_cast(edge)]; } - contracted_costs[static_cast(target)] += - costs[static_cast(edge)]; } if (number_of_contracted_edges == 0) { @@ -166,15 +199,38 @@ private: return result; } - Objective sub_objective(contracted_graph, std::move(contracted_costs)); - const auto sub_labels = sub_solver.optimize(sub_objective); + std::vector sub_labels; + { + BIOIMAGE_PROFILE_SCOPE(profile, "sub_solve"); + if (sub_solver == nullptr) { + // Fast path: call greedy-additive directly with the shared + // workspace, bypassing Objective construction and the + // dense-relabel that `optimize(Objective&)` does internally. + sub_labels = greedy_additive( + contracted_graph, + contracted_costs, + 0.0, + -1.0, + false, + 42, + 1.0, + greedy_workspace + ); + } else { + Objective sub_objective(contracted_graph, std::move(contracted_costs)); + sub_labels = sub_solver->optimize(sub_objective); + } + } std::vector result(number_of_nodes); - for (std::uint64_t node = 0; node < graph.number_of_nodes(); ++node) { - const auto root = contraction.root_of_node[static_cast(node)]; - result[static_cast(node)] = sub_labels[ - static_cast(root) - ]; + { + BIOIMAGE_PROFILE_SCOPE(profile, "lift"); + for (std::uint64_t node = 0; node < graph.number_of_nodes(); ++node) { + const auto root = contraction.root_of_node[static_cast(node)]; + result[static_cast(node)] = sub_labels[ + static_cast(root) + ]; + } } return result; } diff --git a/include/bioimage_cpp/graph/multicut/greedy_additive.hxx b/include/bioimage_cpp/graph/multicut/greedy_additive.hxx index 82f1b4f..ccf50ec 100644 --- a/include/bioimage_cpp/graph/multicut/greedy_additive.hxx +++ b/include/bioimage_cpp/graph/multicut/greedy_additive.hxx @@ -9,6 +9,22 @@ namespace bioimage_cpp::graph::multicut { +// Reusable scratch state for `greedy_additive`. Construct once and call +// `greedy_additive(..., workspace)` repeatedly to avoid per-call allocation +// of the DynamicGraph, UnionFind, and EdgeHeap. Capacities only grow; the +// internal vectors are reset (not freed) between calls. +struct GreedyAdditiveWorkspace { + detail::DynamicGraph dynamic_graph; + bioimage_cpp::detail::UnionFind union_find{0}; + detail::EdgeHeap heap; + + void reset(const UndirectedGraph &graph) { + dynamic_graph.reset(graph); + union_find.reset(static_cast(graph.number_of_nodes())); + heap.reset_capacity(static_cast(graph.number_of_edges())); + } +}; + inline std::vector greedy_additive( const UndirectedGraph &graph, const std::vector &costs, @@ -16,12 +32,14 @@ inline std::vector greedy_additive( const double node_num_stop, const bool add_noise, const int seed, - const double sigma + const double sigma, + GreedyAdditiveWorkspace &workspace ) { validate_costs(graph, costs); - detail::DynamicGraph dynamic_graph(graph); - bioimage_cpp::detail::UnionFind sets(static_cast(graph.number_of_nodes())); - detail::EdgeHeap heap; + workspace.reset(graph); + auto &dynamic_graph = workspace.dynamic_graph; + auto &sets = workspace.union_find; + auto &heap = workspace.heap; detail::initialize_dynamic_graph(graph, costs, dynamic_graph, heap, false, add_noise, seed, sigma); while (!heap.empty() && dynamic_graph.alive_count > 1) { @@ -39,6 +57,21 @@ inline std::vector greedy_additive( return detail::labels_from_sets(sets, graph); } +inline std::vector greedy_additive( + const UndirectedGraph &graph, + const std::vector &costs, + const double weight_stop, + const double node_num_stop, + const bool add_noise, + const int seed, + const double sigma +) { + GreedyAdditiveWorkspace workspace; + return greedy_additive( + graph, costs, weight_stop, node_num_stop, add_noise, seed, sigma, workspace + ); +} + class GreedyAdditiveSolver final : public SolverBase { public: GreedyAdditiveSolver( diff --git a/include/bioimage_cpp/graph/proposal_generators/watershed.hxx b/include/bioimage_cpp/graph/proposal_generators/watershed.hxx index ecc154f..40a5e2f 100644 --- a/include/bioimage_cpp/graph/proposal_generators/watershed.hxx +++ b/include/bioimage_cpp/graph/proposal_generators/watershed.hxx @@ -20,6 +20,9 @@ namespace bioimage_cpp::graph { // // `n_seeds_fraction` is interpreted as a fraction of `number_of_nodes` when // <= 1.0 and as an absolute seed-pair count otherwise. +// +// Scratch buffers (noisy costs, seeds, watershed sort buffer) are held as +// members and reused across `generate` calls. class WatershedProposalGenerator final : public ProposalGeneratorBase { public: WatershedProposalGenerator( @@ -35,7 +38,9 @@ public: n_seeds_fraction_(n_seeds_fraction), seed_(seed), generator_(static_cast(seed)), - noise_(0.0, sigma) { + noise_(0.0, sigma), + noisy_costs_(edge_costs_.size()), + seeds_buffer_(static_cast(graph.number_of_nodes())) { if (edge_costs_.size() != static_cast(graph.number_of_edges())) { throw std::invalid_argument( "edge_costs length must equal number_of_edges" @@ -60,9 +65,8 @@ public: return; } - std::vector noisy_costs(edge_costs_.size()); for (std::size_t edge = 0; edge < edge_costs_.size(); ++edge) { - noisy_costs[edge] = static_cast(edge_costs_[edge] + noise_(generator_)); + noisy_costs_[edge] = static_cast(edge_costs_[edge] + noise_(generator_)); } std::size_t n_seed_pairs; @@ -76,17 +80,22 @@ public: n_seed_pairs = std::max(std::size_t{1}, n_seed_pairs); n_seed_pairs = std::min(negative_edges_.size(), n_seed_pairs); - std::vector seeds(static_cast(number_of_nodes), 0); + std::fill(seeds_buffer_.begin(), seeds_buffer_.end(), std::uint64_t{0}); std::uniform_int_distribution edge_dist(0, negative_edges_.size() - 1); std::uint64_t next_label = 1; for (std::size_t i = 0; i < n_seed_pairs; ++i) { const auto edge = negative_edges_[edge_dist(generator_)]; const auto uv = graph_.uv(edge); - seeds[static_cast(uv.first)] = next_label++; - seeds[static_cast(uv.second)] = next_label++; + seeds_buffer_[static_cast(uv.first)] = next_label++; + seeds_buffer_[static_cast(uv.second)] = next_label++; } - proposal = edge_weighted_watershed(graph_, noisy_costs, seeds); + // Copy seeds into the output buffer; the kernel propagates labels in + // place on `proposal`. + proposal.assign(seeds_buffer_.begin(), seeds_buffer_.end()); + edge_weighted_watershed_into( + graph_, noisy_costs_, proposal, watershed_scratch_ + ); } void reset() override { @@ -102,6 +111,9 @@ private: std::vector negative_edges_; std::mt19937 generator_; std::normal_distribution noise_; + std::vector noisy_costs_; + std::vector seeds_buffer_; + detail::EdgeWeightedWatershedScratch watershed_scratch_; }; } // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/undirected_graph.hxx b/include/bioimage_cpp/graph/undirected_graph.hxx index dc34605..7aa5b20 100644 --- a/include/bioimage_cpp/graph/undirected_graph.hxx +++ b/include/bioimage_cpp/graph/undirected_graph.hxx @@ -137,6 +137,57 @@ public: return 2 + 2 * number_of_edges(); } + // Fast construction from a pre-sorted, deduplicated edge list. + // + // Precondition: `edges` is sorted ascending by `(u, v)` with `u < v` in + // every entry, and contains no duplicates. No node id may equal or exceed + // `number_of_nodes`. The call takes ownership of `edges` and uses it as + // the graph's edge storage, bypassing the per-edge hash dedup that + // `insert_edge` performs — useful when bulk-building a contracted graph + // whose unique edges are already known. + // + // When `populate_lookup` is false, the edge-lookup hash map is left empty + // and `find_edge`/`insert_edge` are not available on the returned graph. + // Used by the fusion-move contraction primitive, whose sub-solver only + // walks edges and adjacency lists. + static UndirectedGraph from_sorted_unique_edges( + const NodeId number_of_nodes, + std::vector edges, + const bool populate_lookup = true + ) { + UndirectedGraph graph(number_of_nodes, static_cast(edges.size())); + + std::vector degree(static_cast(number_of_nodes), 0); + for (const auto &edge : edges) { + ++degree[static_cast(edge.first)]; + ++degree[static_cast(edge.second)]; + } + for (NodeId node = 0; node < number_of_nodes; ++node) { + graph.adjacency_[static_cast(node)].reserve( + degree[static_cast(node)] + ); + } + + if (populate_lookup) { + graph.edge_lookup_.reserve(edges.size()); + } + for (std::size_t index = 0; index < edges.size(); ++index) { + const auto &edge = edges[index]; + const auto edge_id = static_cast(index); + graph.adjacency_[static_cast(edge.first)].push_back( + Adjacency{edge.second, edge_id} + ); + graph.adjacency_[static_cast(edge.second)].push_back( + Adjacency{edge.first, edge_id} + ); + if (populate_lookup) { + graph.edge_lookup_.emplace(edge, edge_id); + } + } + graph.edges_ = std::move(edges); + return graph; + } + [[nodiscard]] std::pair, std::vector> extract_subgraph_from_nodes(const std::vector &nodes) const { std::unordered_set node_set; From 796e53adc6182dd5a2faf6befad1018f1f1b7eed Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Fri, 15 May 2026 14:56:41 +0200 Subject: [PATCH 5/5] Implement parallelization for fusion movs --- AGENTS.md | 51 +++++ MIGRATION_GUIDE.md | 19 +- .../graph/multicut/check_fusion_move.py | 21 +- include/bioimage_cpp/detail/profile.hxx | 29 ++- .../graph/multicut/fusion_move.hxx | 216 +++++++++++++----- src/bindings/graph.cxx | 22 +- src/bioimage_cpp/graph/__init__.py | 80 +++++-- tests/graph/multicut/test_fusion_move.py | 80 ++++++- 8 files changed, 408 insertions(+), 110 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 82f3c25..9efa05c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,6 +33,57 @@ Before introducing union-finds, priority queues, edge hashing, stride math, thre If a needed helper does not exist but is generally useful, add it to `detail/` as a header-only utility with a focused API rather than inlining it into one algorithm. Keep the contract small and well-documented; over time `detail/` is what lets multiple modules stay small and consistent. +## Reusable algorithmic infrastructure + +Some larger pieces of infrastructure sit above `detail/` but are still +intended to be reused across objective types. Check these before starting a +new fusion-move / proposal-based / contraction-based solver: + +- `include/bioimage_cpp/graph/detail/fusion_contract.hxx` — objective-agnostic + agreement-projection primitive. `contract_by_agreement(graph, proposals, + n_proposals, ...)` returns `{contracted_graph, contracted_edge_of_original, + root_of_node}`. Already supports N ≥ 1 proposals; reused by both pairwise + and joint multi-proposal fuses. +- `include/bioimage_cpp/graph/proposal_generator.hxx` — `ProposalGeneratorBase` + abstract class. Concrete generators (Watershed, GreedyAdditiveMulticut) live + in `proposal_generators/` and depend only on `(graph, edge_costs)` plus an + RNG seed. They emit `std::vector` node labelings and are + therefore reusable across multicut, lifted multicut, mincut, etc. +- `include/bioimage_cpp/graph/multicut/greedy_additive.hxx` — exposes + `GreedyAdditiveWorkspace` so multiple invocations on different graphs + share scratch buffers. Use this pattern (workspace + `reset(graph)`) when + a fusion-move driver calls a sub-solver inside its iteration loop. +- `UndirectedGraph::from_sorted_unique_edges(N, edges, populate_lookup=false)` + — bulk graph construction without the per-edge hash insertion in + `insert_edge`. Pair with `populate_lookup=false` when the consumer only + walks edges / adjacency (the multicut sub-solvers do). +- `detail/threading.hxx::parallel_for_chunks` — the only threading primitive + we use. New parallel solvers should not introduce alternatives. + +When porting fusion moves to a new objective (e.g. lifted multicut): + +1. The driver loop in `multicut/fusion_move.hxx::FusionMoveSolver::optimize` + is short and dense. Duplicate it for the new objective rather than + abstracting it via a template/CRTP base — the moving parts (cost + aggregation, energy evaluator, sub-solver type) are objective-specific + and template gymnastics buy little. +2. Reuse `contract_by_agreement` unchanged; it operates on the *base* graph + only. +3. Write a new `fuse_multi(...)` that aggregates *both* base and lifted (or + other auxiliary) weights through `contraction.contracted_edge_of_original` + and `contraction.root_of_node`, calls the new objective's sub-solver, and + lifts labels back via `root_of_node`. +4. Reuse the existing `WatershedProposalGenerator` and + `GreedyAdditiveMulticutProposalGenerator` verbatim; they only depend on + the base graph + base costs and emit node labelings. Add objective-specific + generators (e.g. `GreedyAdditiveLiftedMulticutProposalGenerator`) only if a + meaningful new proposal strategy emerges. +5. Reuse the per-thread parallel pattern from `optimize`: stage-1 parallel + proposal generation + parallel pairwise fuse, stage-2 sequential joint + multi-fuse on leftover candidates. Per-thread `GreedyAdditiveWorkspace` + becomes per-thread `Workspace` if the new sub-solver follows + the same pattern. + ## Dependencies **Allowed**: C++20 stdlib, `nanobind`, `numpy`, `scikit-build-core`, `cmake`, `pytest`, `cibuildwheel` (CI only). diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index d47813f..904f637 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -466,12 +466,19 @@ Intentional differences vs. nifty: proposal-generator object is built lazily when `optimize` is called. - The driver warm-starts from the trivial singleton labeling by running the default greedy-additive sub-solver once before the proposal loop. -- A best-of-three safety net (current, proposal, fused) keeps the running - energy monotonically non-increasing across iterations. -- The current implementation is single-threaded and uses pairwise - (proposal, current) fuses. `number_of_threads` and - `number_of_parallel_proposals` are kept on the API for forward compatibility - but must be left at their defaults (`1` and `2` respectively). +- A best-of safety net keeps the running energy monotonically non-increasing + across iterations (compared against current, proposals, fused, and the + stage-2 joint fuse). +- Parallel proposal generation and a multi-proposal joint fuse are supported: + `number_of_threads=T` runs `number_of_parallel_proposals=P` proposal + generators in parallel within each iteration. By default `P=2` when `T=1` + and `P=T` when `T>1`; pass an explicit `number_of_parallel_proposals` to + override. Each parallel slot uses an independent proposal generator with + seed `proposal_generator.seed + slot_index` so the result is deterministic + for a given `(seed, T, P)`. When at least two parallel pairwise fuses fail + to improve on the current best, a joint multi-proposal fuse runs over the + surviving fused candidates (matches nifty's `ccFusionMoveBased` stage-2 + behaviour). Notes: diff --git a/development/graph/multicut/check_fusion_move.py b/development/graph/multicut/check_fusion_move.py index 102d9c9..a779762 100644 --- a/development/graph/multicut/check_fusion_move.py +++ b/development/graph/multicut/check_fusion_move.py @@ -7,18 +7,23 @@ def main() -> None: args = parser( - "Compare bioimage-cpp and nifty fusion-move multicut at default settings." + "Compare bioimage-cpp and nifty fusion-move multicut at matched settings." ).parse_args() - # bioimage-cpp warm-starts from the trivial singleton labeling with a - # greedy-additive pass before the proposal loop. Nifty's ccFusionMoveBased - # exposes the same behaviour via the `warmStartGreedy=True` flag (which - # internally chains a greedyAdditiveFactory in front of the fusion-move - # factory). Both sides therefore enter the proposal loop from the same - # starting point. + # Match settings on both sides: + # - threads / parallel-proposals = `--threads N` (default 1). We set + # P = T explicitly on both sides so the comparison is apples-to-apples + # regardless of either library's API default for P. + # - greedy-additive warm-start (we do it automatically when initial + # labels are the trivial singleton; nifty exposes it as + # `warmStartGreedy=True`). + # - greedy-additive sub-solver. + threads = int(args.threads) run_comparison( "fusion_move", lambda: bic.graph.FusionMoveMulticut( proposal_generator=bic.graph.WatershedProposalGenerator(), + number_of_threads=threads, + number_of_parallel_proposals=threads, ), lambda objective: objective.ccFusionMoveBasedFactory( proposalGenerator=objective.watershedCcProposals(), @@ -27,7 +32,7 @@ def main() -> None: ), numberOfIterations=10, stopIfNoImprovement=4, - numberOfThreads=1, + numberOfThreads=threads, warmStartGreedy=True, ), args, diff --git a/include/bioimage_cpp/detail/profile.hxx b/include/bioimage_cpp/detail/profile.hxx index ba39c54..124ab4c 100644 --- a/include/bioimage_cpp/detail/profile.hxx +++ b/include/bioimage_cpp/detail/profile.hxx @@ -5,6 +5,12 @@ // and adds no overhead. Use it in development/profile-mode builds to find // hotspots; do not enable it in production wheels. +// `NullProfiler` is always available; helper templates can request "no +// profiling here" via the same type regardless of build mode. +namespace bioimage_cpp::detail { +struct NullProfiler {}; +} // namespace bioimage_cpp::detail + #ifdef BIOIMAGE_PROFILE #include #include @@ -63,18 +69,31 @@ private: Profiler::Clock::time_point start_; }; +// Overload that accepts a NullProfiler and does nothing. Lets the same +// macros expand cleanly when a profiled translation unit calls into a helper +// that explicitly does not want to participate in profiling (e.g. work inside +// parallel workers, where we measure wall-clock at the dispatch level). +class ProfileTimerNull { +public: + ProfileTimerNull(NullProfiler &, const char *) {} +}; + +inline ProfileTimer make_profile_timer(Profiler &profiler, const char *name) { + return ProfileTimer(profiler, name); +} + +inline ProfileTimerNull make_profile_timer(NullProfiler &profiler, const char *name) { + return ProfileTimerNull(profiler, name); +} + } // namespace bioimage_cpp::detail #define BIOIMAGE_PROFILE_INIT(var) ::bioimage_cpp::detail::Profiler var; -#define BIOIMAGE_PROFILE_SCOPE(var, name) ::bioimage_cpp::detail::ProfileTimer _bp_##__LINE__(var, name); +#define BIOIMAGE_PROFILE_SCOPE(var, name) auto _bp_##__LINE__ = ::bioimage_cpp::detail::make_profile_timer(var, name); #define BIOIMAGE_PROFILE_REPORT(var) (var).report(); #else -namespace bioimage_cpp::detail { -struct NullProfiler {}; -} // namespace bioimage_cpp::detail - #define BIOIMAGE_PROFILE_INIT(var) ::bioimage_cpp::detail::NullProfiler var; #define BIOIMAGE_PROFILE_SCOPE(var, name) (void)var; #define BIOIMAGE_PROFILE_REPORT(var) (void)var; diff --git a/include/bioimage_cpp/graph/multicut/fusion_move.hxx b/include/bioimage_cpp/graph/multicut/fusion_move.hxx index 51ca220..1cdd3e4 100644 --- a/include/bioimage_cpp/graph/multicut/fusion_move.hxx +++ b/include/bioimage_cpp/graph/multicut/fusion_move.hxx @@ -1,6 +1,7 @@ #pragma once #include "bioimage_cpp/detail/profile.hxx" +#include "bioimage_cpp/detail/threading.hxx" #include "bioimage_cpp/graph/detail/fusion_contract.hxx" #include "bioimage_cpp/graph/multicut/greedy_additive.hxx" #include "bioimage_cpp/graph/multicut/objective.hxx" @@ -8,9 +9,12 @@ #include "bioimage_cpp/graph/undirected_graph.hxx" #include +#include #include #include +#include #include +#include #include #include #include @@ -19,35 +23,46 @@ namespace bioimage_cpp::graph::multicut { class FusionMoveSolver final : public SolverBase { public: - // The proposal generator is borrowed; the caller owns its lifetime. - // The sub-solver pointer is optional: when null, the driver uses a default - // greedy-additive sub-solver internally. + // Each entry in `proposal_generators` is one parallel-proposal source. + // The container must have exactly `number_of_parallel_proposals` entries. + // When `number_of_threads > 1` the workers index the generators by + // proposal slot; each generator must have independent state (own RNG, + // own scratch). The pointers are borrowed; the caller owns lifetimes. // - // `number_of_threads` and `number_of_parallel_proposals` are reserved for - // future use; v1 only supports the single-threaded pairwise - // (proposal, current) fuse and rejects other values. + // `sub_solver` is optional: nullptr uses an internal greedy-additive + // sub-solver with a shared workspace per worker. FusionMoveSolver( - ProposalGeneratorBase &proposal_generator, + std::vector proposal_generators, const SolverBase *sub_solver = nullptr, const std::size_t number_of_iterations = 10, const std::size_t stop_if_no_improvement = 4, const std::size_t number_of_threads = 1, const std::size_t number_of_parallel_proposals = 2 ) - : proposal_generator_(proposal_generator), + : proposal_generators_(std::move(proposal_generators)), sub_solver_(sub_solver), number_of_iterations_(number_of_iterations), - stop_if_no_improvement_(stop_if_no_improvement) { - if (number_of_threads != 1) { + stop_if_no_improvement_(stop_if_no_improvement), + number_of_threads_(number_of_threads), + number_of_parallel_proposals_(number_of_parallel_proposals) { + if (number_of_parallel_proposals < 1) { throw std::invalid_argument( - "FusionMoveSolver currently supports number_of_threads=1 only" + "number_of_parallel_proposals must be >= 1" ); } - if (number_of_parallel_proposals != 2) { + if (number_of_threads < 1) { + throw std::invalid_argument("number_of_threads must be >= 1"); + } + if (proposal_generators_.size() != number_of_parallel_proposals) { throw std::invalid_argument( - "FusionMoveSolver currently supports number_of_parallel_proposals=2 only" + "proposal_generators length must equal number_of_parallel_proposals" ); } + for (const auto *pgen : proposal_generators_) { + if (pgen == nullptr) { + throw std::invalid_argument("proposal_generators must not contain null"); + } + } } std::vector optimize(Objective &objective) const override { @@ -62,17 +77,19 @@ public: return objective.labels(); } - // Workspace reused across the warm-start, every fuse iteration's - // sub-solve, and (transitively) any callers chaining additional - // greedy-additive runs. - GreedyAdditiveWorkspace greedy_workspace; + // One workspace per worker thread; reused across the warm-start, every + // pairwise fuse, and the stage-2 joint fuse. + const auto effective_threads = ::bioimage_cpp::detail::normalize_thread_count( + number_of_threads_, number_of_parallel_proposals_ + ); + std::vector workspaces(effective_threads); - // Warm start from greedy-additive if the caller passed the trivial + // Warm-start from greedy-additive if the caller passed the trivial // singleton labeling. if (is_singleton_labeling(current)) { BIOIMAGE_PROFILE_SCOPE(profile, "warm_start"); current = greedy_additive( - graph, costs, 0.0, -1.0, false, 42, 1.0, greedy_workspace + graph, costs, 0.0, -1.0, false, 42, 1.0, workspaces[0] ); } @@ -82,44 +99,106 @@ public: current_energy = energy(graph, costs, current); } - std::vector proposal(static_cast(number_of_nodes)); + // Per-proposal-slot buffers. The proposal generator writes into + // `proposal_buffers[p]`; the pairwise-fuse writes into + // `fused_buffers[p]`. Both are reused across iterations. + const std::size_t P = number_of_parallel_proposals_; + std::vector> proposal_buffers(P); + std::vector> fused_buffers(P); + std::vector proposal_energies(P); + std::vector fused_energies(P); + std::vector is_leftover(P); + + constexpr double kEnergyEps = 1e-7; + std::size_t iterations_without_improvement = 0; for (std::size_t iteration = 0; iteration < number_of_iterations_; ++iteration) { - { - BIOIMAGE_PROFILE_SCOPE(profile, "proposal"); - proposal_generator_.generate(current, proposal); - } + // === Stage 1: parallel proposal generation + parallel pairwise fuse === - std::vector fused = fuse_pair( - graph, costs, current, proposal, sub_solver_, greedy_workspace, profile - ); + // Snapshot current under no mutation (only the calling thread writes + // to `current` between iterations, so workers can read it freely). + const auto ¤t_snapshot = current; + + std::fill(is_leftover.begin(), is_leftover.end(), 0); - double fused_energy; - { - BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); - fused_energy = energy(graph, costs, fused); - } - double proposal_energy; { - BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); - proposal_energy = energy(graph, costs, proposal); + BIOIMAGE_PROFILE_SCOPE(profile, "proposal_and_pairwise_fuse"); + ::bioimage_cpp::detail::parallel_for_chunks( + effective_threads, + P, + [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { + auto &workspace = workspaces[thread_id]; + for (std::size_t p = begin; p < end; ++p) { + proposal_generators_[p]->generate( + current_snapshot, proposal_buffers[p] + ); + proposal_energies[p] = energy(graph, costs, proposal_buffers[p]); + + fuse_pair_into( + graph, + costs, + current_snapshot, + proposal_buffers[p], + sub_solver_, + workspace, + fused_buffers[p] + ); + fused_energies[p] = energy(graph, costs, fused_buffers[p]); + } + } + ); } - // Best-of safety net across the three candidates so an iteration - // can never raise the running energy. + // === Aggregate stage-1 results sequentially (small loop) === + // Track the best candidate across {current, proposals, fused}. + // Collect "leftovers" — fuse results that did not improve on + // `current_energy` and were not effectively equal — for stage 2. double best_energy = current_energy; const std::vector *best = ¤t; - if (fused_energy < best_energy) { - best_energy = fused_energy; - best = &fused; + std::size_t leftover_count = 0; + for (std::size_t p = 0; p < P; ++p) { + if (proposal_energies[p] < best_energy) { + best_energy = proposal_energies[p]; + best = &proposal_buffers[p]; + } + if (fused_energies[p] < best_energy) { + best_energy = fused_energies[p]; + best = &fused_buffers[p]; + } + if (fused_energies[p] > current_energy + kEnergyEps) { + is_leftover[p] = 1; + ++leftover_count; + } } - if (proposal_energy < best_energy) { - best_energy = proposal_energy; - best = &proposal; + + // === Stage 2: joint multi-proposal fuse on leftovers === + std::vector joint_result; + double joint_energy = std::numeric_limits::infinity(); + if (leftover_count >= 2) { + BIOIMAGE_PROFILE_SCOPE(profile, "joint_fuse"); + std::vector *> leftovers; + leftovers.reserve(leftover_count); + for (std::size_t p = 0; p < P; ++p) { + if (is_leftover[p]) { + leftovers.push_back(&fused_buffers[p]); + } + } + joint_result = fuse_multi( + graph, costs, leftovers, sub_solver_, workspaces[0], profile + ); + { + BIOIMAGE_PROFILE_SCOPE(profile, "energy_eval"); + joint_energy = energy(graph, costs, joint_result); + } + if (joint_energy < best_energy) { + best_energy = joint_energy; + best = &joint_result; + } } - if (best_energy < current_energy) { + // === Update current under the best-of safety net === + if (best_energy + kEnergyEps < current_energy) { current = *best; current_energy = best_energy; iterations_without_improvement = 0; @@ -146,26 +225,56 @@ private: return true; } - template - static std::vector fuse_pair( + // Pairwise fuse that writes into a caller-provided output buffer (used by + // the parallel stage-1 loop so workers don't allocate). + static void fuse_pair_into( const UndirectedGraph &graph, const std::vector &costs, const std::vector ¤t, const std::vector &proposal, const SolverBase *sub_solver, + GreedyAdditiveWorkspace &workspace, + std::vector &output + ) { + const std::array *, 2> proposals{ + ¤t, &proposal + }; + std::vector *> proposal_list( + proposals.begin(), proposals.end() + ); + ::bioimage_cpp::detail::NullProfiler null_profile; + output = fuse_multi(graph, costs, proposal_list, sub_solver, workspace, null_profile); + } + + // Multi-input fuse: contract by agreement over all N proposals, sum costs + // onto the contracted edges, sub-solve, lift labels back. N=2 is the + // pairwise case; N>2 is the stage-2 joint fuse on leftovers. + template + static std::vector fuse_multi( + const UndirectedGraph &graph, + const std::vector &costs, + const std::vector *> &proposals, + const SolverBase *sub_solver, GreedyAdditiveWorkspace &greedy_workspace, [[maybe_unused]] ProfilerT &profile ) { const auto number_of_nodes = static_cast(graph.number_of_nodes()); - std::vector stacked(2 * number_of_nodes); - std::copy(current.begin(), current.end(), stacked.begin()); - std::copy(proposal.begin(), proposal.end(), stacked.begin() + number_of_nodes); + const auto n_proposals = proposals.size(); + + std::vector stacked(n_proposals * number_of_nodes); + for (std::size_t p = 0; p < n_proposals; ++p) { + std::copy( + proposals[p]->begin(), + proposals[p]->end(), + stacked.begin() + static_cast(p * number_of_nodes) + ); + } ::bioimage_cpp::graph::detail::AgreementContraction contraction; { BIOIMAGE_PROFILE_SCOPE(profile, "agreement_contract"); contraction = ::bioimage_cpp::graph::detail::contract_by_agreement( - graph, stacked.data(), 2, number_of_nodes + graph, stacked.data(), n_proposals, number_of_nodes ); } @@ -203,9 +312,6 @@ private: { BIOIMAGE_PROFILE_SCOPE(profile, "sub_solve"); if (sub_solver == nullptr) { - // Fast path: call greedy-additive directly with the shared - // workspace, bypassing Objective construction and the - // dense-relabel that `optimize(Objective&)` does internally. sub_labels = greedy_additive( contracted_graph, contracted_costs, @@ -235,10 +341,12 @@ private: return result; } - ProposalGeneratorBase &proposal_generator_; + std::vector proposal_generators_; const SolverBase *sub_solver_; std::size_t number_of_iterations_; std::size_t stop_if_no_improvement_; + std::size_t number_of_threads_; + std::size_t number_of_parallel_proposals_; }; } // namespace bioimage_cpp::graph::multicut diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index ac8bc35..35dca71 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -418,23 +418,27 @@ UInt64Array multicut_fusion_move( const Graph &graph, ConstDoubleArray costs, ConstUInt64Array initial_labels, - graph::ProposalGeneratorBase *proposal_generator, + std::vector proposal_generators, const graph::multicut::SolverBase *sub_solver, const std::size_t number_of_iterations, - const std::size_t stop_if_no_improvement + const std::size_t stop_if_no_improvement, + const std::size_t number_of_threads, + const std::size_t number_of_parallel_proposals ) { - if (proposal_generator == nullptr) { - throw std::invalid_argument("proposal_generator must not be None"); + if (proposal_generators.empty()) { + throw std::invalid_argument("proposal_generators must not be empty"); } auto cost_vector = double_array_to_vector(costs, "edge_costs", graph.number_of_edges()); auto label_vector = uint64_array_to_vector(initial_labels, "initial_labels", graph.number_of_nodes()); graph::multicut::FusionMoveSolver solver( - *proposal_generator, + std::move(proposal_generators), sub_solver, number_of_iterations, - stop_if_no_improvement + stop_if_no_improvement, + number_of_threads, + number_of_parallel_proposals ); std::vector result; @@ -842,10 +846,12 @@ void bind_graph(nb::module_ &m) { nb::arg("graph"), nb::arg("edge_costs"), nb::arg("initial_labels"), - nb::arg("proposal_generator"), + nb::arg("proposal_generators"), nb::arg("sub_solver").none(), nb::arg("number_of_iterations"), - nb::arg("stop_if_no_improvement") + nb::arg("stop_if_no_improvement"), + nb::arg("number_of_threads"), + nb::arg("number_of_parallel_proposals") ); m.def( diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index 281711b..83496b9 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -459,14 +459,20 @@ def _build_cpp_sub_solver(self): class ProposalGenerator(ABC): """Base class for fusion-move proposal generators. - Concrete generators carry settings on the Python side and build a C++ - proposal-generator object bound to a specific problem (graph + edge costs) - when ``_build`` is called by the fusion-move driver. + Concrete generators carry settings on the Python side. ``_build_for_thread`` + constructs an independent underlying C++ proposal-generator object whose + seed is offset by ``seed_offset`` so that parallel proposal slots produce + distinct, reproducible streams. """ @abstractmethod - def _build(self, graph: UndirectedGraph | RegionAdjacencyGraph, edge_costs: np.ndarray): - """Construct the underlying C++ proposal generator.""" + def _build_for_thread( + self, + graph: UndirectedGraph | RegionAdjacencyGraph, + edge_costs: np.ndarray, + seed_offset: int, + ): + """Construct the underlying C++ proposal generator with a seed offset.""" class WatershedProposalGenerator(ProposalGenerator): @@ -487,13 +493,13 @@ def __init__( self.n_seeds_fraction = float(n_seeds_fraction) self.seed = int(seed) - def _build(self, graph, edge_costs): + def _build_for_thread(self, graph, edge_costs, seed_offset): return _core._WatershedProposalGenerator( graph, edge_costs, sigma=self.sigma, n_seeds_fraction=self.n_seeds_fraction, - seed=self.seed, + seed=self.seed + int(seed_offset), ) @@ -517,14 +523,14 @@ def __init__( self.node_num_stop = float(node_num_stop) self.seed = int(seed) - def _build(self, graph, edge_costs): + def _build_for_thread(self, graph, edge_costs, seed_offset): return _core._GreedyAdditiveMulticutProposalGenerator( graph, edge_costs, sigma=self.sigma, weight_stop=self.weight_stop, node_num_stop=self.node_num_stop, - seed=self.seed, + seed=self.seed + int(seed_offset), ) @@ -532,16 +538,26 @@ class FusionMoveMulticut(MulticutSolver): """Fusion-move multicut solver. Iteratively generates proposals via ``proposal_generator``, fuses them - pairwise with the current best labeling, and accepts improvements. The - fuse step solves a contracted multicut subproblem with ``sub_solver``; - if omitted, the default sub-solver is :class:`GreedyAdditiveMulticut`. + with the current best labeling, and accepts improvements. The fuse step + solves a contracted multicut subproblem with ``sub_solver``; if omitted, + the default sub-solver is :class:`GreedyAdditiveMulticut`. If the objective's current labels are the trivial singleton labeling, the driver warm-starts with one greedy-additive pass before the proposal loop. The best-of safety net guarantees energy never increases across iterations. - ``number_of_threads`` and ``number_of_parallel_proposals`` are reserved for - future use; the current implementation only supports the defaults. + Threading: ``number_of_threads > 1`` runs ``number_of_parallel_proposals`` + proposal generators in parallel within each iteration. Each parallel slot + uses an independent proposal generator with seed ``proposal_generator.seed + + slot_index``. By default ``number_of_parallel_proposals`` is ``2`` when + ``number_of_threads == 1`` and ``number_of_threads`` otherwise; pass it + explicitly to override. + + Multi-proposal fuse: when at least two parallel pairwise fuses fail to + improve on the current best, a joint multi-proposal fuse is run over the + surviving fused candidates (matches nifty's ``ccFusionMoveBased`` stage-2 + behaviour). With ``number_of_parallel_proposals == 2`` this stage rarely + triggers; it becomes useful as ``number_of_parallel_proposals`` grows. """ def __init__( @@ -552,7 +568,7 @@ def __init__( number_of_iterations: int = 10, stop_if_no_improvement: int = 4, number_of_threads: int = 1, - number_of_parallel_proposals: int = 2, + number_of_parallel_proposals: int | None = None, ): if not isinstance(proposal_generator, ProposalGenerator): raise TypeError("proposal_generator must inherit from ProposalGenerator") @@ -563,25 +579,37 @@ def __init__( "sub_solver must be a built-in multicut solver " "(custom Python solvers are not supported as fusion-move sub-solvers)" ) - if int(number_of_threads) != 1: - raise ValueError("number_of_threads must be 1 (parallel path not yet implemented)") - if int(number_of_parallel_proposals) != 2: - raise ValueError( - "number_of_parallel_proposals must be 2 (multi-proposal fuse not yet implemented)" - ) + n_threads = int(number_of_threads) + if n_threads < 1: + raise ValueError("number_of_threads must be >= 1") + if number_of_parallel_proposals is None: + n_parallel = 2 if n_threads == 1 else n_threads + else: + n_parallel = int(number_of_parallel_proposals) + if n_parallel < 1: + raise ValueError("number_of_parallel_proposals must be >= 1") + self.proposal_generator = proposal_generator self.sub_solver = sub_solver self.number_of_iterations = int(number_of_iterations) self.stop_if_no_improvement = int(stop_if_no_improvement) - self.number_of_threads = int(number_of_threads) - self.number_of_parallel_proposals = int(number_of_parallel_proposals) + self.number_of_threads = n_threads + self.number_of_parallel_proposals = n_parallel if self.number_of_iterations < 0: raise ValueError("number_of_iterations must be non-negative") if self.stop_if_no_improvement < 1: raise ValueError("stop_if_no_improvement must be >= 1") def optimize(self, objective: MulticutObjective) -> np.ndarray: - cpp_pgen = self.proposal_generator._build(objective.graph, objective.edge_costs) + # Build one C++ proposal generator per parallel slot, each with a + # distinct seed offset, so parallel streams are independent and + # reproducible. + cpp_pgens = [ + self.proposal_generator._build_for_thread( + objective.graph, objective.edge_costs, slot + ) + for slot in range(self.number_of_parallel_proposals) + ] cpp_sub_solver = ( None if self.sub_solver is None else self.sub_solver._build_cpp_sub_solver() ) @@ -589,10 +617,12 @@ def optimize(self, objective: MulticutObjective) -> np.ndarray: objective.graph, objective.edge_costs, objective.labels, - cpp_pgen, + cpp_pgens, cpp_sub_solver, self.number_of_iterations, self.stop_if_no_improvement, + self.number_of_threads, + self.number_of_parallel_proposals, ) objective.labels = labels return objective.labels diff --git a/tests/graph/multicut/test_fusion_move.py b/tests/graph/multicut/test_fusion_move.py index aa33037..c586089 100644 --- a/tests/graph/multicut/test_fusion_move.py +++ b/tests/graph/multicut/test_fusion_move.py @@ -140,19 +140,19 @@ def test_rejects_non_proposal_generator(): bic.graph.FusionMoveMulticut(proposal_generator=object()) -def test_rejects_unsupported_thread_count(): +def test_rejects_zero_thread_count(): with pytest.raises(ValueError, match="number_of_threads"): bic.graph.FusionMoveMulticut( proposal_generator=bic.graph.WatershedProposalGenerator(), - number_of_threads=4, + number_of_threads=0, ) -def test_rejects_unsupported_parallel_proposals(): +def test_rejects_zero_parallel_proposals(): with pytest.raises(ValueError, match="number_of_parallel_proposals"): bic.graph.FusionMoveMulticut( proposal_generator=bic.graph.WatershedProposalGenerator(), - number_of_parallel_proposals=4, + number_of_parallel_proposals=0, ) @@ -179,6 +179,78 @@ def test_runs_on_empty_graph(): assert labels.shape == (0,) +def test_parallel_threads_match_single_threaded_safety_net(grid_problem): + # With T>1 the loop should never regress past the greedy-additive baseline + # (best-of safety net is enforced in C++). + graph, costs = grid_problem + baseline = bic.graph.MulticutObjective(graph, costs).energy( + bic.graph.GreedyAdditiveMulticut().optimize(bic.graph.MulticutObjective(graph, costs)) + ) + + objective = bic.graph.MulticutObjective(graph, costs) + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=11), + number_of_threads=4, + number_of_iterations=5, + ) + labels = solver.optimize(objective) + assert objective.energy(labels) <= baseline + 1e-9 + + +def test_multi_proposal_runs(grid_problem): + # number_of_parallel_proposals > 2 exercises the stage-2 joint fuse path + # (no determinism guarantee against T=1 here — the algorithm shape changes). + graph, costs = grid_problem + baseline = bic.graph.MulticutObjective(graph, costs).energy( + bic.graph.GreedyAdditiveMulticut().optimize(bic.graph.MulticutObjective(graph, costs)) + ) + + objective = bic.graph.MulticutObjective(graph, costs) + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=3), + number_of_threads=2, + number_of_parallel_proposals=4, + number_of_iterations=5, + ) + labels = solver.optimize(objective) + assert objective.energy(labels) <= baseline + 1e-9 + + +def test_parallel_is_deterministic_given_settings(grid_problem): + # Same seed + same T + same P + same iterations → same result. + graph, costs = grid_problem + + def run(): + objective = bic.graph.MulticutObjective(graph, costs) + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(seed=7), + number_of_threads=4, + number_of_iterations=5, + ) + return solver.optimize(objective) + + np.testing.assert_array_equal(run(), run()) + + +def test_default_parallel_proposals_tracks_threads(): + pgen = bic.graph.WatershedProposalGenerator() + one_thread = bic.graph.FusionMoveMulticut(proposal_generator=pgen) + four_threads = bic.graph.FusionMoveMulticut( + proposal_generator=pgen, number_of_threads=4 + ) + assert one_thread.number_of_parallel_proposals == 2 + assert four_threads.number_of_parallel_proposals == 4 + + +def test_explicit_parallel_proposals_overrides_default(): + solver = bic.graph.FusionMoveMulticut( + proposal_generator=bic.graph.WatershedProposalGenerator(), + number_of_threads=4, + number_of_parallel_proposals=8, + ) + assert solver.number_of_parallel_proposals == 8 + + def test_runs_on_graph_without_negative_edges(chain_problem): # WatershedProposalGenerator yields all-zero proposals if no negative # edges exist; the driver must still terminate cleanly and produce the