bioimage-cpp/include/bioimage_cpp/graph/multicut/greedy_additive.hxx at 349408e7cb2dd111d3d34b192c422bc251e85c3e · computational-cell-analytics/bioimage-cpp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
#pragma once

#include "bioimage_cpp/graph/multicut/detail.hxx"
#include "bioimage_cpp/graph/multicut/objective.hxx"

#include <cstddef>
#include <cstdint>
#include <vector>

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::util::UnionFind union_find{0};
    detail::EdgeHeap heap;

    void reset(const UndirectedGraph &graph) {
        dynamic_graph.reset(graph);
        union_find.reset(static_cast<std::size_t>(graph.number_of_nodes()));
        heap.reset_capacity(static_cast<std::size_t>(graph.number_of_edges()));
    }
};

inline std::vector<std::uint64_t> greedy_additive(
    const UndirectedGraph &graph,
    const std::vector<double> &costs,
    const double weight_stop,
    const double node_num_stop,
    const bool add_noise,
    const int seed,
    const double sigma,
    GreedyAdditiveWorkspace &workspace
) {
    validate_costs(graph, costs);
    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) {
        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)) {
            break;
        }
        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);
}

inline std::vector<std::uint64_t> greedy_additive(
    const UndirectedGraph &graph,
    const std::vector<double> &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(
        const double weight_stop = 0.0,
        const double node_num_stop = -1.0,
        const bool add_noise = false,
        const int seed = 42,
        const double sigma = 1.0
    )
        : weight_stop_(weight_stop),
          node_num_stop_(node_num_stop),
          add_noise_(add_noise),
          seed_(seed),
          sigma_(sigma) {
    }

    std::vector<std::uint64_t> optimize(Objective &objective) const override {
        auto labels = greedy_additive(
            objective.graph(),
            objective.costs(),
            weight_stop_,
            node_num_stop_,
            add_noise_,
            seed_,
            sigma_
        );
        objective.set_labels(labels);
        return labels;
    }

private:
    double weight_stop_;
    double node_num_stop_;
    bool add_noise_;
    int seed_;
    double sigma_;
};

} // namespace bioimage_cpp::graph::multicut