Multicut performance improvements

Author: constantinpape

include/bioimage_cpp/graph/multicut/detail.hxx

@@ -1,115 +1,107 @@
 #pragma once
 
+#include "bioimage_cpp/detail/indexed_heap.hxx"
 #include "bioimage_cpp/detail/union_find.hxx"
 #include "bioimage_cpp/graph/connected_components.hxx"
 #include "bioimage_cpp/graph/multicut/objective.hxx"
 #include "bioimage_cpp/graph/undirected_graph.hxx"
 
-#include <algorithm>
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
-#include <queue>
+#include <limits>
 #include <random>
-#include <unordered_map>
-#include <unordered_set>
 #include <utility>
 #include <vector>
 
 namespace bioimage_cpp::graph::multicut::detail {
 
+inline constexpr std::size_t no_edge = std::numeric_limits<std::size_t>::max();
+
+// Per-super-node adjacency entry.
+struct NeighborEntry {
+    std::size_t neighbor;
+    std::size_t edge_id;
+};
+
+// Per-edge data stored in a flat vector indexed by stable edge_id. Edge ids
+// never grow past the original number of input edges — when two edges fold
+// into one during a merge, one id survives and the other becomes orphaned
+// (no adjacency entry or heap entry references it again).
 struct DynamicEdge {
+    std::size_t u = 0;
+    std::size_t v = 0;
     double weight = 0.0;
-    std::uint64_t edition = 0;
+    unsigned char is_constraint = 0;
 };
 
-struct DynamicGraph {
-    explicit DynamicGraph(const std::size_t size) : adjacency(size), constraints(size), alive(size, true), alive_count(size) {
-    }
+// Heap keyed by dense stable edge_id. The vector-backed DenseLocator updates
+// in O(1) per sift step, which was the missing ingredient for greedy_additive
+// to beat the std::priority_queue baseline.
+using EdgeHeap = bioimage_cpp::detail::DenseIndexedHeap<double>;
 
+struct DynamicGraph {
     explicit DynamicGraph(const UndirectedGraph &graph)
         : adjacency(static_cast<std::size_t>(graph.number_of_nodes())),
-          constraints(static_cast<std::size_t>(graph.number_of_nodes())),
           alive(static_cast<std::size_t>(graph.number_of_nodes()), true),
-          alive_count(static_cast<std::size_t>(graph.number_of_nodes())) {
+          alive_count(static_cast<std::size_t>(graph.number_of_nodes())),
+          scratch_edge_id(static_cast<std::size_t>(graph.number_of_nodes()), no_edge) {
         for (std::uint64_t node = 0; node < graph.number_of_nodes(); ++node) {
             const auto degree = graph.node_adjacency(node).size();
             adjacency[static_cast<std::size_t>(node)].reserve(degree);
         }
+        edges.resize(static_cast<std::size_t>(graph.number_of_edges()));
     }
 
-    [[nodiscard]] bool edge_exists(const std::size_t u, const std::size_t v) const {
-        return alive[u] && alive[v] && adjacency[u].find(v) != adjacency[u].end();
-    }
-
-    [[nodiscard]] const DynamicEdge *edge(const std::size_t u, const std::size_t v) const {
-        if (!alive[u] || !alive[v]) {
-            return nullptr;
+    // O(degree(u)). Returns no_edge when (u, v) is not an edge.
+    [[nodiscard]] std::size_t find_edge(const std::size_t u, const std::size_t v) const {
+        if (u >= adjacency.size() || v >= adjacency.size()) {
+            return no_edge;
         }
-        const auto found = adjacency[u].find(v);
-        return found == adjacency[u].end() ? nullptr : &found->second;
-    }
-
-    [[nodiscard]] bool has_constraint(const std::size_t u, const std::size_t v) const {
-        return constraints[u].find(v) != constraints[u].end();
-    }
-
-    void add_constraint(const std::size_t u, const std::size_t v) {
-        constraints[u].insert(v);
-        constraints[v].insert(u);
-    }
-
-    void set_initial_edge(const std::size_t u, const std::size_t v, const double weight) {
-        adjacency[u][v] = DynamicEdge{weight, 0};
-        adjacency[v][u] = DynamicEdge{weight, 0};
+        for (const auto &entry : adjacency[u]) {
+            if (entry.neighbor == v) {
+                return entry.edge_id;
+            }
+        }
+        return no_edge;
     }
 
-    std::uint64_t update_edge(const std::size_t u, const std::size_t v, const double weight) {
-        const auto current = adjacency[u].find(v);
-        const auto edition = current == adjacency[u].end() ? 0 : current->second.edition + 1;
-        adjacency[u][v] = DynamicEdge{weight, edition};
-        adjacency[v][u] = DynamicEdge{weight, edition};
-        return edition;
+    [[nodiscard]] bool edge_exists(const std::size_t u, const std::size_t v) const {
+        return find_edge(u, v) != no_edge;
     }
 
-    void remove_node(const std::size_t node) {
-        for (const auto &entry : adjacency[node]) {
-            adjacency[entry.first].erase(node);
-            constraints[entry.first].erase(node);
-        }
-        adjacency[node].clear();
-        constraints[node].clear();
-        alive[node] = false;
-        --alive_count;
+    [[nodiscard]] bool has_constraint(const std::size_t u, const std::size_t v) const {
+        const auto id = find_edge(u, v);
+        return id != no_edge && edges[id].is_constraint != 0;
     }
 
-    std::vector<std::unordered_map<std::size_t, DynamicEdge>> adjacency;
-    std::vector<std::unordered_set<std::size_t>> constraints;
+    std::vector<std::vector<NeighborEntry>> adjacency;
+    std::vector<DynamicEdge> edges;
     std::vector<bool> alive;
     std::size_t alive_count;
+    // Sized to #super-nodes, all entries == no_edge between merges. Used by
+    // merge_dynamic_nodes to find the existing stable-side edge for each of
+    // removed's neighbors in O(1) without hashing.
+    std::vector<std::size_t> scratch_edge_id;
 };
 
-struct QueueEdge {
-    std::size_t u = 0;
-    std::size_t v = 0;
-    double priority = 0.0;
-    std::uint64_t edition = 0;
-
-    bool operator<(const QueueEdge &other) const {
-        return priority < other.priority;
-    }
-};
+inline double priority_for(const double weight, const bool absolute_priority) {
+    return absolute_priority ? std::abs(weight) : weight;
+}
 
 inline void initialize_dynamic_graph(
     const UndirectedGraph &graph,
     const std::vector<double> &costs,
     DynamicGraph &dynamic_graph,
-    std::priority_queue<QueueEdge> &queue,
+    EdgeHeap &heap,
     const bool absolute_priority,
     const bool add_noise = false,
     const int seed = 42,
     const double sigma = 1.0
 ) {
+    const auto n_edges = static_cast<std::size_t>(graph.number_of_edges());
+    heap.reset_capacity(n_edges);
+
     std::mt19937 generator(seed);
     std::normal_distribution<double> noise(0.0, sigma);
     for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) {
@@ -120,8 +112,15 @@ inline void initialize_dynamic_graph(
         if (add_noise) {
             weight += noise(generator);
         }
-        dynamic_graph.set_initial_edge(u, v, weight);
-        queue.push(QueueEdge{u, v, absolute_priority ? std::abs(weight) : weight, 0});
+        const auto edge_id = static_cast<std::size_t>(edge);
+        auto &e = dynamic_graph.edges[edge_id];
+        e.u = u;
+        e.v = v;
+        e.weight = weight;
+        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));
     }
 }
 
@@ -138,10 +137,45 @@ inline std::size_t stop_node_count(const UndirectedGraph &graph, const double no
         : static_cast<std::size_t>(double(graph.number_of_nodes()) * node_num_stop + 0.5);
 }
 
+namespace internal {
+
+// Erase the first entry with neighbor == target from `list` (swap-with-back).
+// Returns whether anything was erased.
+inline bool erase_by_neighbor(std::vector<NeighborEntry> &list, const std::size_t target) {
+    for (std::size_t i = 0; i < list.size(); ++i) {
+        if (list[i].neighbor == target) {
+            list[i] = list.back();
+            list.pop_back();
+            return true;
+        }
+    }
+    return false;
+}
+
+// Rename the first entry whose neighbor == from_node to point to to_node.
+inline void rename_neighbor(
+    std::vector<NeighborEntry> &list,
+    const std::size_t from_node,
+    const std::size_t to_node
+) {
+    for (auto &entry : list) {
+        if (entry.neighbor == from_node) {
+            entry.neighbor = to_node;
+            return;
+        }
+    }
+}
+
+} // namespace internal
+
+// Contract the edge between u and v in the dynamic graph. The smaller-degree
+// super-node is folded into the larger-degree one to keep the per-super-node
+// adjacency growth amortized. The heap is kept in sync without staleness: each
+// edge id appears at most once.
 inline std::size_t merge_dynamic_nodes(
     DynamicGraph &dynamic_graph,
     bioimage_cpp::detail::UnionFind &sets,
-    std::priority_queue<QueueEdge> &queue,
+    EdgeHeap &heap,
     std::size_t u,
     std::size_t v,
     const bool absolute_priority
@@ -151,36 +185,79 @@ inline std::size_t merge_dynamic_nodes(
     if (u == v) {
         return u;
     }
+
     auto stable = u;
     auto removed = v;
     if (dynamic_graph.adjacency[stable].size() < dynamic_graph.adjacency[removed].size()) {
         std::swap(stable, removed);
     }
     sets.merge_to(stable, removed);
-    std::vector<std::pair<std::size_t, DynamicEdge>> neighbors(
-        dynamic_graph.adjacency[removed].begin(),
-        dynamic_graph.adjacency[removed].end()
-    );
-    dynamic_graph.adjacency[stable].erase(removed);
-    dynamic_graph.constraints[stable].erase(removed);
-    for (const auto &[neighbor, removed_edge] : neighbors) {
+
+    // Stamp stable's neighbors so each removed-neighbor lookup is O(1).
+    for (const auto &entry : dynamic_graph.adjacency[stable]) {
+        dynamic_graph.scratch_edge_id[entry.neighbor] = entry.edge_id;
+    }
+
+    // The contracted edge (stable, removed) must be in stable's adjacency.
+    const auto contracted_edge_id = dynamic_graph.scratch_edge_id[removed];
+    heap.erase(contracted_edge_id);
+    dynamic_graph.scratch_edge_id[removed] = no_edge;
+    internal::erase_by_neighbor(dynamic_graph.adjacency[stable], removed);
+
+    // Snapshot removed's neighbors before mutating its adjacency.
+    const auto removed_neighbors = dynamic_graph.adjacency[removed];
+
+    for (const auto &entry : removed_neighbors) {
+        const auto neighbor = entry.neighbor;
+        const auto removed_edge_id = entry.edge_id;
         if (neighbor == stable) {
             continue;
         }
-        const auto current = dynamic_graph.edge(stable, neighbor);
-        const auto merged_weight = (current == nullptr ? 0.0 : current->weight) + removed_edge.weight;
-        if (dynamic_graph.has_constraint(removed, neighbor)) {
-            dynamic_graph.add_constraint(stable, neighbor);
+
+        const auto existing_id = dynamic_graph.scratch_edge_id[neighbor];
+        if (existing_id == no_edge) {
+            // No existing stable-side edge to `neighbor`. Re-key the
+            // removed-side edge by replacing `removed` with `stable` on both
+            // sides. The edge_id, weight and constraint flag are preserved,
+            // so the heap entry remains valid without any priority change.
+            dynamic_graph.adjacency[stable].push_back({neighbor, removed_edge_id});
+            dynamic_graph.scratch_edge_id[neighbor] = removed_edge_id;
+            internal::rename_neighbor(dynamic_graph.adjacency[neighbor], removed, stable);
+            auto &e = dynamic_graph.edges[removed_edge_id];
+            if (e.u == removed) {
+                e.u = stable;
+            } else {
+                e.v = stable;
+            }
+        } else {
+            // Stable already has an edge to `neighbor`; fold removed's weight
+            // (and constraint flag) into it. Then drop the removed-side edge.
+            auto &keep = dynamic_graph.edges[existing_id];
+            const auto &fold = dynamic_graph.edges[removed_edge_id];
+            keep.weight += fold.weight;
+            const bool propagated_constraint =
+                keep.is_constraint != 0 || fold.is_constraint != 0;
+            keep.is_constraint = propagated_constraint ? 1 : 0;
+
+            heap.erase(removed_edge_id);
+            internal::erase_by_neighbor(dynamic_graph.adjacency[neighbor], removed);
+
+            if (propagated_constraint) {
+                heap.erase(existing_id);
+            } else {
+                heap.change(existing_id, priority_for(keep.weight, absolute_priority));
+            }
         }
-        const auto edition = dynamic_graph.update_edge(stable, neighbor, merged_weight);
-        queue.push(QueueEdge{
-            std::min(stable, neighbor),
-            std::max(stable, neighbor),
-            absolute_priority ? std::abs(merged_weight) : merged_weight,
-            edition
-        });
     }
-    dynamic_graph.remove_node(removed);
+
+    // Clear scratch via the updated stable adjacency (includes appended entries).
+    for (const auto &entry : dynamic_graph.adjacency[stable]) {
+        dynamic_graph.scratch_edge_id[entry.neighbor] = no_edge;
+    }
+
+    dynamic_graph.adjacency[removed].clear();
+    dynamic_graph.alive[removed] = false;
+    --dynamic_graph.alive_count;
     return stable;
 }
 

include/bioimage_cpp/graph/multicut/greedy_additive.hxx

@@ -3,11 +3,8 @@
 #include "bioimage_cpp/graph/multicut/detail.hxx"
 #include "bioimage_cpp/graph/multicut/objective.hxx"
 
-#include <algorithm>
 #include <cstddef>
 #include <cstdint>
-#include <queue>
-#include <unordered_map>
 #include <vector>
 
 namespace bioimage_cpp::graph::multicut {
@@ -24,26 +21,20 @@ inline std::vector<std::uint64_t> greedy_additive(
     validate_costs(graph, costs);
     detail::DynamicGraph dynamic_graph(graph);
     bioimage_cpp::detail::UnionFind sets(static_cast<std::size_t>(graph.number_of_nodes()));
-    std::priority_queue<detail::QueueEdge> queue;
-    detail::initialize_dynamic_graph(graph, costs, dynamic_graph, queue, false, add_noise, seed, sigma);
+    detail::EdgeHeap heap;
+    detail::initialize_dynamic_graph(graph, costs, dynamic_graph, heap, false, add_noise, seed, sigma);
 
-    while (!queue.empty() && dynamic_graph.alive_count > 1) {
-        auto edge = queue.top();
-        queue.pop();
-        const auto u = std::min(edge.u, edge.v);
-        const auto v = std::max(edge.u, edge.v);
-        const auto *dynamic_edge = dynamic_graph.edge(u, v);
-        if (dynamic_edge == nullptr || edge.edition < dynamic_edge->edition) {
-            continue;
-        }
-        const auto weight = dynamic_edge->weight;
-        if (weight <= weight_stop) {
+    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(graph, node_num_stop)) {
+        if (node_num_stop > 0.0
+            && dynamic_graph.alive_count <= detail::stop_node_count(graph, node_num_stop)) {
             break;
         }
-        detail::merge_dynamic_nodes(dynamic_graph, sets, queue, u, v, false);
+        const auto &edge = dynamic_graph.edges[top.key];
+        detail::merge_dynamic_nodes(dynamic_graph, sets, heap, edge.u, edge.v, false);
     }
     return detail::labels_from_sets(sets, graph);
 }