Skip to main content

tensorlogic_quantrs_hooks/loopy_bp/
energy.rs

1//! Bethe free energy calculation for converged LBP beliefs.
2
3use scirs2_core::ndarray::{Array1, ArrayD};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7use crate::graph::FactorGraph;
8
9/// Bethe free energy components.
10///
11/// The Bethe approximation decomposes the global free energy into single-variable
12/// and factor contributions.  At the fixed point of LBP, the Bethe free energy
13/// equals the variational free energy under the Bethe approximation to the
14/// belief propagation.
15#[derive(Clone, Debug, Serialize, Deserialize)]
16pub struct BetheFreeEnergy {
17    /// Factor-node energy term: ∑_a ∑_{x_a} b_a(x_a) ln[φ_a(x_a) / b_a(x_a)]
18    pub factor_energy: f64,
19    /// Variable-node entropy term: (1 - d_i) ∑_i ∑_{x_i} b_i(x_i) ln b_i(x_i)
20    pub variable_entropy: f64,
21    /// Total Bethe free energy = factor_energy + variable_entropy.
22    pub total: f64,
23    /// Approximate log-partition function: -F_Bethe.
24    pub log_z: f64,
25}
26
27/// Compute the Bethe free energy from converged LBP beliefs.
28///
29/// `beliefs_var` is a map from variable name → marginal belief vector.
30/// `beliefs_fac` is a map from factor id → joint belief tensor (over factor scope).
31pub fn bethe_free_energy(
32    graph: &FactorGraph,
33    beliefs_var: &HashMap<String, Array1<f64>>,
34    beliefs_fac: &HashMap<String, ArrayD<f64>>,
35) -> BetheFreeEnergy {
36    let eps = 1e-300_f64;
37
38    // Factor-node contribution: ∑_a ∑_{x_a} b_a(x_a) [ln φ_a(x_a) - ln b_a(x_a)]
39    let mut factor_energy = 0.0_f64;
40    for (fac_id, fac_belief) in beliefs_fac {
41        if let Some(factor) = graph.get_factor(fac_id) {
42            for (b, phi) in fac_belief.iter().zip(factor.values.iter()) {
43                if *b > eps {
44                    let log_phi = if *phi > eps { phi.ln() } else { -700.0 };
45                    factor_energy += b * (log_phi - b.ln());
46                }
47            }
48        }
49    }
50
51    // Variable-node contribution: ∑_i (1 - d_i) ∑_{x_i} b_i(x_i) ln b_i(x_i)
52    let mut variable_entropy = 0.0_f64;
53    for (var_name, belief) in beliefs_var {
54        // degree d_i = number of factors containing variable i
55        let degree = graph
56            .get_adjacent_factors(var_name)
57            .map(|v| v.len())
58            .unwrap_or(0) as f64;
59        let entropy_i: f64 = belief
60            .iter()
61            .filter(|&&b| b > eps)
62            .map(|&b| b * b.ln())
63            .sum::<f64>();
64        // variable term = -(1 - d_i) * H_i  where H_i = -∑ b ln b
65        // ↔ (1 - d_i) * ∑ b ln b
66        variable_entropy += (1.0 - degree) * entropy_i;
67    }
68
69    let total = -(factor_energy + variable_entropy);
70    BetheFreeEnergy {
71        factor_energy,
72        variable_entropy,
73        total,
74        log_z: -total,
75    }
76}