Optimization round 1

Author: constantinpape

CMakeLists.txt

@@ -21,5 +21,8 @@ nanobind_add_module(_core
 
 target_include_directories(_core PRIVATE include)
 target_compile_features(_core PRIVATE cxx_std_20)
+target_compile_options(_core PRIVATE
+    $<$<NOT:$<CONFIG:Debug>>:-O3>
+)
 
 install(TARGETS _core LIBRARY DESTINATION bioimage_cpp)

development/graph/multicut/_compatibility.py

@@ -67,6 +67,11 @@ def time_call(function: Callable[[], np.ndarray], repeats: int):
     return timings, result
 
 
+def optimize_bic_solver(make_bic_solver, objective):
+    objective.reset_labels()
+    return make_bic_solver().optimize(objective)
+
+
 def bic_energy(graph, costs: np.ndarray, labels: np.ndarray) -> float:
     import bioimage_cpp as bic
 
@@ -89,10 +94,11 @@ def run_comparison(
     import nifty.graph.opt.multicut as nmc
 
     bic_graph, nifty_graph, costs = load_problem(args.path, timeout=args.timeout)
+    bic_objective = bic.graph.MulticutObjective(bic_graph, costs)
     nifty_objective = nmc.multicutObjective(nifty_graph, costs)
 
     bic_timings, bic_labels = time_call(
-        lambda: make_bic_solver().optimize(bic.graph.MulticutObjective(bic_graph, costs)),
+        lambda: optimize_bic_solver(make_bic_solver, bic_objective),
         args.repeats,
     )
     nifty_timings, nifty_labels = time_call(

include/bioimage_cpp/graph/connected_components.hxx

@@ -39,6 +39,18 @@ public:
         }
     }
 
+    void merge_to(std::uint64_t stable, std::uint64_t removed) {
+        stable = find(stable);
+        removed = find(removed);
+        if (stable == removed) {
+            return;
+        }
+        parents_[static_cast<std::size_t>(removed)] = stable;
+        if (ranks_[static_cast<std::size_t>(stable)] <= ranks_[static_cast<std::size_t>(removed)]) {
+            ranks_[static_cast<std::size_t>(stable)] = ranks_[static_cast<std::size_t>(removed)] + 1;
+        }
+    }
+
 private:
     std::vector<std::uint64_t> parents_;
     std::vector<std::uint64_t> ranks_;

include/bioimage_cpp/graph/multicut/detail.hxx

@@ -17,14 +17,38 @@
 
 namespace bioimage_cpp::graph::multicut::detail {
 
+struct DynamicEdge {
+    double weight = 0.0;
+    std::uint64_t edition = 0;
+};
+
 struct DynamicGraph {
     explicit DynamicGraph(const std::size_t size) : adjacency(size), constraints(size), alive(size, true), alive_count(size) {
     }
 
+    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())) {
+        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);
+        }
+    }
+
     [[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;
+        }
+        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();
     }
@@ -34,9 +58,17 @@ struct DynamicGraph {
         constraints[v].insert(u);
     }
 
-    void set_edge(const std::size_t u, const std::size_t v, const double weight) {
-        adjacency[u][v] = weight;
-        adjacency[v][u] = weight;
+    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};
+    }
+
+    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;
     }
 
     void remove_node(const std::size_t node) {
@@ -50,7 +82,7 @@ struct DynamicGraph {
         --alive_count;
     }
 
-    std::vector<std::unordered_map<std::size_t, double>> adjacency;
+    std::vector<std::unordered_map<std::size_t, DynamicEdge>> adjacency;
     std::vector<std::unordered_set<std::size_t>> constraints;
     std::vector<bool> alive;
     std::size_t alive_count;
@@ -71,7 +103,6 @@ inline void initialize_dynamic_graph(
     const UndirectedGraph &graph,
     const std::vector<double> &costs,
     DynamicGraph &dynamic_graph,
-    std::vector<std::unordered_map<std::size_t, std::uint64_t>> &editions,
     std::priority_queue<QueueEdge> &queue,
     const bool absolute_priority,
     const bool add_noise = false,
@@ -88,9 +119,7 @@ inline void initialize_dynamic_graph(
         if (add_noise) {
             weight += noise(generator);
         }
-        dynamic_graph.set_edge(u, v, weight);
-        editions[u][v] = 0;
-        editions[v][u] = 0;
+        dynamic_graph.set_initial_edge(u, v, weight);
         queue.push(QueueEdge{u, v, absolute_priority ? std::abs(weight) : weight, 0});
     }
 }
@@ -108,7 +137,6 @@ inline std::size_t stop_node_count(const UndirectedGraph &graph, const double no
 inline std::size_t merge_dynamic_nodes(
     DynamicGraph &dynamic_graph,
     UnionFind &sets,
-    std::vector<std::unordered_map<std::size_t, std::uint64_t>> &editions,
     std::priority_queue<QueueEdge> &queue,
     std::size_t u,
     std::size_t v,
@@ -119,28 +147,28 @@ inline std::size_t merge_dynamic_nodes(
     if (u == v) {
         return u;
     }
-    sets.merge(u, v);
-    const auto stable = static_cast<std::size_t>(sets.find(u));
-    const auto removed = (stable == u) ? v : u;
-    std::vector<std::pair<std::size_t, double>> neighbors(
+    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_weight] : neighbors) {
+    for (const auto &[neighbor, removed_edge] : neighbors) {
         if (neighbor == stable) {
             continue;
         }
-        const auto current = dynamic_graph.edge_exists(stable, neighbor)
-            ? dynamic_graph.adjacency[stable][neighbor]
-            : 0.0;
-        const auto merged_weight = current + removed_weight;
+        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);
         }
-        dynamic_graph.set_edge(stable, neighbor, merged_weight);
-        const auto edition = ++editions[std::min(stable, neighbor)][std::max(stable, neighbor)];
+        const auto edition = dynamic_graph.update_edge(stable, neighbor, merged_weight);
         queue.push(QueueEdge{
             std::min(stable, neighbor),
             std::max(stable, neighbor),

include/bioimage_cpp/graph/multicut/greedy_additive.hxx

@@ -22,28 +22,28 @@ inline std::vector<std::uint64_t> greedy_additive(
     const double sigma
 ) {
     validate_costs(graph, costs);
-    detail::DynamicGraph dynamic_graph(static_cast<std::size_t>(graph.number_of_nodes()));
+    detail::DynamicGraph dynamic_graph(graph);
     UnionFind sets(static_cast<std::size_t>(graph.number_of_nodes()));
-    std::vector<std::unordered_map<std::size_t, std::uint64_t>> editions(static_cast<std::size_t>(graph.number_of_nodes()));
     std::priority_queue<detail::QueueEdge> queue;
-    detail::initialize_dynamic_graph(graph, costs, dynamic_graph, editions, queue, false, add_noise, seed, sigma);
+    detail::initialize_dynamic_graph(graph, costs, dynamic_graph, queue, 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);
-        if (!dynamic_graph.edge_exists(u, v) || edge.edition < editions[u][v]) {
+        const auto *dynamic_edge = dynamic_graph.edge(u, v);
+        if (dynamic_edge == nullptr || edge.edition < dynamic_edge->edition) {
             continue;
         }
-        const auto weight = dynamic_graph.adjacency[u][v];
+        const auto weight = dynamic_edge->weight;
         if (weight <= weight_stop) {
             break;
         }
         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, editions, queue, u, v, false);
+        detail::merge_dynamic_nodes(dynamic_graph, sets, queue, u, v, false);
     }
     return detail::labels_from_sets(sets, graph);
 }

include/bioimage_cpp/graph/multicut/greedy_fixation.hxx

@@ -20,21 +20,21 @@ inline std::vector<std::uint64_t> greedy_fixation(
     const double node_num_stop
 ) {
     validate_costs(graph, costs);
-    detail::DynamicGraph dynamic_graph(static_cast<std::size_t>(graph.number_of_nodes()));
+    detail::DynamicGraph dynamic_graph(graph);
     UnionFind sets(static_cast<std::size_t>(graph.number_of_nodes()));
-    std::vector<std::unordered_map<std::size_t, std::uint64_t>> editions(static_cast<std::size_t>(graph.number_of_nodes()));
     std::priority_queue<detail::QueueEdge> queue;
-    detail::initialize_dynamic_graph(graph, costs, dynamic_graph, editions, queue, true);
+    detail::initialize_dynamic_graph(graph, costs, dynamic_graph, queue, true);
 
     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);
-        if (!dynamic_graph.edge_exists(u, v) || edge.edition < editions[u][v]) {
+        const auto *dynamic_edge = dynamic_graph.edge(u, v);
+        if (dynamic_edge == nullptr || edge.edition < dynamic_edge->edition) {
             continue;
         }
-        const auto weight = dynamic_graph.adjacency[u][v];
+        const auto weight = dynamic_edge->weight;
         if (std::abs(weight) <= weight_stop) {
             break;
         }
@@ -45,7 +45,7 @@ inline std::vector<std::uint64_t> greedy_fixation(
             continue;
         }
         if (weight > 0.0) {
-            detail::merge_dynamic_nodes(dynamic_graph, sets, editions, queue, u, v, true);
+            detail::merge_dynamic_nodes(dynamic_graph, sets, queue, u, v, true);
         } else if (weight < 0.0) {
             dynamic_graph.add_constraint(u, v);
         }

src/bioimage_cpp/graph/__init__.py

@@ -371,6 +371,9 @@ def __init__(
         self.number_of_threads = _normalize_number_of_threads(number_of_threads)
 
     def optimize(self, objective: MulticutObjective) -> np.ndarray:
+        if self.fallthrough_solver is None and isinstance(self.sub_solver, GreedyAdditiveMulticut):
+            return self.sub_solver.optimize(objective)
+
         component_labels = connected_components(
             objective.graph,
             edge_mask=objective.edge_costs > 0.0,