Optimize greedy additive solver

Author: constantinpape

include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "bioimage_cpp/detail/profile.hxx"
 #include "bioimage_cpp/graph/lifted_multicut/detail.hxx"
 #include "bioimage_cpp/graph/lifted_multicut/objective.hxx"
 
@@ -40,28 +41,44 @@ inline std::vector<std::uint64_t> greedy_additive(
     const double sigma,
     GreedyAdditiveWorkspace &workspace
 ) {
+    BIOIMAGE_PROFILE_INIT(profile);
     validate_weights(lifted_graph, weights);
-    workspace.reset(lifted_graph);
+    {
+        BIOIMAGE_PROFILE_SCOPE(profile, "workspace_reset");
+        workspace.reset(lifted_graph);
+    }
     auto &dynamic_graph = workspace.dynamic_graph;
     auto &sets = workspace.union_find;
     auto &heap = workspace.heap;
-    detail::initialize_dynamic_graph(
-        lifted_graph, weights, n_base_edges, dynamic_graph, heap, add_noise, seed, sigma
-    );
+    {
+        BIOIMAGE_PROFILE_SCOPE(profile, "initialize");
+        detail::initialize_dynamic_graph(
+            lifted_graph, weights, n_base_edges, dynamic_graph, heap, add_noise, seed, sigma
+        );
+    }
 
-    while (!heap.empty() && dynamic_graph.alive_count > 1) {
-        const auto top = heap.top();
-        if (top.priority <= weight_stop) {
-            break;
-        }
-        if (node_num_stop > 0.0
-            && dynamic_graph.alive_count <= detail::stop_node_count(lifted_graph, node_num_stop)) {
-            break;
+    {
+        BIOIMAGE_PROFILE_SCOPE(profile, "contraction_loop");
+        while (!heap.empty() && dynamic_graph.alive_count > 1) {
+            const auto top = heap.top();
+            if (top.priority <= weight_stop) {
+                break;
+            }
+            if (node_num_stop > 0.0
+                && dynamic_graph.alive_count <= detail::stop_node_count(lifted_graph, node_num_stop)) {
+                break;
+            }
+            const auto &edge = dynamic_graph.edges[top.key];
+            detail::merge_dynamic_nodes(dynamic_graph, sets, heap, edge.u, edge.v);
         }
-        const auto &edge = dynamic_graph.edges[top.key];
-        detail::merge_dynamic_nodes(dynamic_graph, sets, heap, edge.u, edge.v);
     }
-    return detail::labels_from_sets(sets, lifted_graph);
+    std::vector<std::uint64_t> labels;
+    {
+        BIOIMAGE_PROFILE_SCOPE(profile, "labels_from_sets");
+        labels = detail::labels_from_sets(sets, lifted_graph);
+    }
+    BIOIMAGE_PROFILE_REPORT(profile);
+    return labels;
 }
 
 inline std::vector<std::uint64_t> greedy_additive(

src/bindings/graph.cxx

@@ -219,6 +219,30 @@ Graph graph_from_edges(const std::uint64_t number_of_nodes, ConstUInt64Array uvs
     return graph;
 }
 
+// Bulk-construct a graph from a pre-deduplicated (u, v) array, bypassing the
+// per-edge hash dedup that ``insert_edge`` performs. The caller asserts that
+// no (u, v) pair appears twice in ``uvs`` and that ``u != v`` in every row.
+// Edges receive ids matching their position in ``uvs``. This is the fast path
+// for copying an existing graph (its ``uv_ids()`` are unique by construction).
+Graph graph_from_unique_edges(const std::uint64_t number_of_nodes, ConstUInt64Array uvs) {
+    require_uv_array(uvs, "uvs");
+    std::vector<Graph::Edge> edges;
+    edges.reserve(static_cast<std::size_t>(uvs.shape(0)));
+    const auto *in = uvs.data();
+    for (std::size_t index = 0; index < uvs.shape(0); ++index) {
+        const auto u = in[2 * index];
+        const auto v = in[2 * index + 1];
+        if (u == v) {
+            throw std::invalid_argument("self edges are not supported");
+        }
+        if (u >= number_of_nodes || v >= number_of_nodes) {
+            throw std::out_of_range("edge endpoint exceeds number_of_nodes");
+        }
+        edges.emplace_back(u, v);
+    }
+    return Graph::from_sorted_unique_edges(number_of_nodes, std::move(edges));
+}
+
 Graph graph_deserialize(ConstUInt64Array serialization) {
     if (serialization.ndim() != 1) {
         throw std::invalid_argument("serialization must be a 1D uint64 array");
@@ -855,6 +879,12 @@ void bind_graph(nb::module_ &m) {
             nb::arg("number_of_nodes"),
             nb::arg("uvs")
         )
+        .def_static(
+            "from_unique_edges",
+            &graph_from_unique_edges,
+            nb::arg("number_of_nodes"),
+            nb::arg("uvs")
+        )
         .def_static(
             "deserialize",
             &graph_deserialize,

src/bioimage_cpp/graph/__init__.py

@@ -165,11 +165,17 @@ def _as_serialization_array(serialization) -> np.ndarray:
     return np.ascontiguousarray(array)
 
 
-def _copy_graph(graph: UndirectedGraph | RegionAdjacencyGraph) -> UndirectedGraph:
-    copied = UndirectedGraph(int(graph.number_of_nodes), int(graph.number_of_edges))
-    if graph.number_of_edges:
-        copied.insert_edges(graph.uv_ids())
-    return copied
+def _copy_graph(graph: UndirectedGraph | RegionAdjacencyGraph) -> _core.UndirectedGraph:
+    # `uv_ids()` always returns a unique list (graphs deduplicate on insert),
+    # so we can use the bulk constructor that skips per-edge hash dedup —
+    # significantly faster than `insert_edges` for large graphs. The result
+    # is a ``_core.UndirectedGraph``; downstream code (objectives, solvers,
+    # validators) uses base-class methods that work identically.
+    if graph.number_of_edges == 0:
+        return _core.UndirectedGraph(int(graph.number_of_nodes))
+    return _core.UndirectedGraph.from_unique_edges(
+        int(graph.number_of_nodes), graph.uv_ids()
+    )
 
 
 def _as_edge_costs(edge_costs, graph: UndirectedGraph | RegionAdjacencyGraph) -> np.ndarray:
@@ -787,14 +793,22 @@ def __init__(
         overwrite_existing: bool = False,
         initial_labels=None,
     ):
-        base_graph = _copy_graph(graph)
+        # The objective holds a reference to the user's base graph (no
+        # defensive copy — the C++ ``Objective`` already keeps a const
+        # reference). The user is expected to treat the input graph as
+        # immutable while the objective is alive; mutations are visible to
+        # the objective and may produce undefined behaviour.
+        base_graph = graph
         base_costs = _as_edge_costs(edge_costs, base_graph)
 
-        lifted_graph = UndirectedGraph(
-            int(base_graph.number_of_nodes), int(base_graph.number_of_edges)
-        )
+        # Use the bulk constructor for the lifted graph's base portion to
+        # bypass the per-edge hash dedup that ``insert_edges`` performs.
         if int(base_graph.number_of_edges) > 0:
-            lifted_graph.insert_edges(base_graph.uv_ids())
+            lifted_graph = _core.UndirectedGraph.from_unique_edges(
+                int(base_graph.number_of_nodes), base_graph.uv_ids()
+            )
+        else:
+            lifted_graph = _core.UndirectedGraph(int(base_graph.number_of_nodes))
 
         weights_list = [base_costs.copy()]
 
@@ -953,6 +967,55 @@ def _add_lifted_edges(
     # In-place updates require a single flat working buffer; coalesce first.
     if len(weights_list) > 1:
         weights_list[:] = [np.ascontiguousarray(np.concatenate(weights_list))]
+
+    # Fast path: bulk-insert and detect uniqueness from the row count delta.
+    # For the typical case — ``lifted_uvs`` produced by
+    # ``lifted_edges_from_affinities`` or by the BFS constructor — every row
+    # is a brand-new edge, so the delta equals the input length and we can
+    # append the weights array directly. No ``find_edges`` calls are needed.
+    if not overwrite_existing:
+        pre_count = int(lifted_graph.number_of_edges)
+        lifted_graph.insert_edges(lifted_uvs)
+        post_count = int(lifted_graph.number_of_edges)
+        n_new = post_count - pre_count
+
+        if n_new == lifted_uvs.shape[0]:
+            weights_list.append(
+                np.ascontiguousarray(lifted_costs.astype(np.float64, copy=False))
+            )
+            return
+
+        # Some rows collided with existing edges or with each other. Use
+        # find_edges to recover the per-row edge id (insertion is already
+        # done; this is just a lookup).
+        edge_ids = np.asarray(lifted_graph.find_edges(lifted_uvs))
+        lifted_costs_f64 = lifted_costs.astype(np.float64, copy=False)
+        working = weights_list[0]
+
+        collision_mask = edge_ids < pre_count
+        if collision_mask.any():
+            np.add.at(
+                working,
+                edge_ids[collision_mask].astype(np.intp, copy=False),
+                lifted_costs_f64[collision_mask],
+            )
+
+        new_mask = ~collision_mask
+        if new_mask.any():
+            slot = (edge_ids[new_mask] - pre_count).astype(np.int64, copy=False)
+            new_weights = np.bincount(
+                slot, weights=lifted_costs_f64[new_mask], minlength=n_new
+            ).astype(np.float64, copy=False)
+        else:
+            new_weights = np.zeros(n_new, dtype=np.float64)
+
+        weights_list[0] = working
+        if n_new > 0:
+            weights_list.append(new_weights)
+        return
+
+    # Slow path: per-row Python loop for ``overwrite_existing=True`` (rare).
+    # Order-sensitive last-write-wins semantics on collisions.
     working = weights_list[0]
     new_costs: list[float] = []
     for index in range(lifted_uvs.shape[0]):
@@ -964,10 +1027,7 @@ def _add_lifted_edges(
         if int(lifted_graph.number_of_edges) > pre:
             new_costs.append(weight)
         else:
-            if overwrite_existing:
-                working[edge] = weight
-            else:
-                working[edge] = working[edge] + weight
+            working[edge] = weight
     if new_costs:
         weights_list[0] = working
         weights_list.append(np.asarray(new_costs, dtype=np.float64))