EnergyModel.hpp

Go to the documentation of this file.

#pragma once
 
#include <vector>
#include <functional>
#include <unordered_map>
#include <cstdint>
 
template <typename T>
class EnergyModel {
public:
    using UnaryCostFn = std::function<T(int node, int label)>;
    using PairwiseCostFn = std::function<T(int node1, int node2, int label1, int label2)>;
 
    EnergyModel(int num_nodes, int num_labels)
        : num_nodes_(num_nodes), num_labels_(num_labels), labels_(num_nodes, 0), neighbors_(num_nodes) {}
 
    [[nodiscard]] int num_nodes() const { return num_nodes_; }
 
    [[nodiscard]] int num_labels() const { return num_labels_; }
 
    [[nodiscard]] int get_label(int node) const { return labels_[node]; }
 
    void set_label(int node, int label) { labels_[node] = label; }
 
    [[nodiscard]] const std::vector<int>& get_labels() const { return labels_; }
 
    void set_labels(const std::vector<int>& labels) { labels_ = labels; }
 
    void set_unary_cost_fn(UnaryCostFn fn) { unary_cost_fn_ = fn; }
 
    void set_pairwise_cost_fn(PairwiseCostFn fn) { pairwise_cost_fn_ = fn; }
 
    [[nodiscard]] T get_unary_cost(int node, int label) const {
        if (!unary_costs_.empty()) {
            return unary_costs_[node * num_labels_ + label];
        }
        return unary_cost_fn_ ? unary_cost_fn_(node, label) : 0;
    }
 
    [[nodiscard]] T get_pairwise_cost(int node1, int node2, int label1, int label2) const {
        if (!edge_weights_.empty()) {
            auto it = edge_weights_.find(edge_key(node1, node2));
            if (it != edge_weights_.end()) {
                return label1 == label2 ? T{0} : it->second;
            }
        }
        if (!pairwise_costs_.empty()) {
            return pairwise_costs_[label1 * num_labels_ + label2];
        }
        return pairwise_cost_fn_ ? pairwise_cost_fn_(node1, node2, label1, label2) : 0;
    }
 
    void set_unary_costs(const std::vector<T>& costs) {
        if (costs.size() != static_cast<size_t>(num_nodes_ * num_labels_)) {
            throw std::invalid_argument("Unary costs array must have size num_nodes * num_labels");
        }
        unary_costs_ = costs;
    }
 
    void set_pairwise_costs(const std::vector<T>& costs) {
        if (costs.size() != static_cast<size_t>(num_labels_ * num_labels_)) {
            throw std::invalid_argument("Pairwise costs array must have size num_labels * num_labels");
        }
        pairwise_costs_ = costs;
    }
 
    void set_edge_weights(const std::vector<int>& n1s, const std::vector<int>& n2s, const std::vector<T>& weights) {
        if (n1s.size() != n2s.size() || n1s.size() != weights.size()) {
            throw std::invalid_argument("n1s, n2s, and weights must have the same size");
        }
        for (size_t i = 0; i < n1s.size(); ++i) {
            edge_weights_[edge_key(n1s[i], n2s[i])] = weights[i];
        }
    }
 
    void add_neighbor(int node1, int node2) {
        neighbors_[node1].push_back(node2);
        neighbors_[node2].push_back(node1);
        neighbor_csr_valid_ = false;
    }
 
    void add_grid_edges(int width, int height) {
        if (width * height != num_nodes_) {
            throw std::invalid_argument("Grid dimensions do not match the number of nodes");
        }
        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                int node = y * width + x;
                if (x + 1 < width)  add_neighbor(node, node + 1);
                if (y + 1 < height) add_neighbor(node, node + width);
            }
        }
    }
 
    [[nodiscard]] const std::vector<int>& get_neighbors(int node) const {
        return neighbors_[node];
    }
 
    class NeighborRange {
    public:
        NeighborRange(const int* begin, const int* end) : begin_(begin), end_(end) {}
        [[nodiscard]] const int* begin() const { return begin_; }
        [[nodiscard]] const int* end() const { return end_; }
        [[nodiscard]] int size() const { return static_cast<int>(end_ - begin_); }
    private:
        const int* begin_;
        const int* end_;
    };
 
    [[nodiscard]] NeighborRange get_neighbors_range(int node) const {
        if (!neighbor_csr_valid_) build_neighbor_csr();
        const int begin = neighbor_offsets_[node];
        const int end = neighbor_offsets_[node + 1];
        const int* base = neighbor_data_.data();
        return {base + begin, base + end};
    }
 
    [[nodiscard]] std::vector<int> get_active_nodes(int alpha_label) const {
        std::vector<int> active_nodes;
        active_nodes.reserve(num_nodes_);
        for (int i = 0; i < num_nodes_; ++i) {
            if (labels_[i] != alpha_label) active_nodes.push_back(i);
        }
        return active_nodes;
    }
 
    [[nodiscard]] T evaluate_total_energy(const std::vector<int> &eval_labels) const {
        T total = 0;
        for (int i = 0; i < num_nodes_; ++i) {
            total += get_unary_cost(i, eval_labels[i]);
            for (int neighbor : neighbors_[i]) {
                if (i < neighbor) {
                    total += get_pairwise_cost(i, neighbor, eval_labels[i], eval_labels[neighbor]);
                }
            }
        }
        return total;
    }
 
    [[nodiscard]] T evaluate_total_energy() const {
        return evaluate_total_energy(labels_);
    }
 
private:
    int64_t edge_key(int n1, int n2) const {
        int a = std::min(n1, n2), b = std::max(n1, n2);
        return (int64_t)a * num_nodes_ + b;
    }
 
    void build_neighbor_csr() const {
        neighbor_offsets_.assign(num_nodes_ + 1, 0);
        size_t total = 0;
        for (int i = 0; i < num_nodes_; ++i) total += neighbors_[i].size();
        neighbor_data_.resize(total);
        int pos = 0;
        for (int i = 0; i < num_nodes_; ++i) {
            neighbor_offsets_[i] = pos;
            for (int j : neighbors_[i]) neighbor_data_[pos++] = j;
        }
        neighbor_offsets_[num_nodes_] = pos;
        neighbor_csr_valid_ = true;
    }
 
    int num_nodes_;
    int num_labels_;
    std::vector<int> labels_;
    std::vector<std::vector<int>> neighbors_;
    UnaryCostFn unary_cost_fn_;
    PairwiseCostFn pairwise_cost_fn_;
    std::vector<T> unary_costs_;
    std::vector<T> pairwise_costs_;
    std::unordered_map<int64_t, T> edge_weights_;
 
    mutable std::vector<int> neighbor_data_;
    mutable std::vector<int> neighbor_offsets_;
    mutable bool neighbor_csr_valid_ = false;
};