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}