qualia_core_db/solvers/learning/sampling/
mcmc.rs

1//! Markov-Chain Monte Carlo (PRML ch 11) — random-walk Metropolis-Hastings over an
2//! arbitrary (unnormalized) log-density. This is the general inference engine the
3//! Bayesian methods lean on: give it `log p(x)` (up to a constant) and it returns
4//! samples from `p`. Kernel-class `Divergent` (the accept/reject branch).
5//!
6//! The target need only be specified up to a normalizing constant — the Metropolis
7//! acceptance ratio `exp(log p(x') − log p(x))` cancels it.
8
9/// Result of an MCMC run.
10#[derive(Debug, Clone)]
11pub struct McmcResult {
12    /// Post-burn-in samples, row-major `n_samples × dim`.
13    pub samples: Vec<f64>,
14    pub n_samples: usize,
15    pub dim: usize,
16    /// Fraction of proposals accepted (tune `proposal_std` toward ~0.234–0.5).
17    pub acceptance_rate: f64,
18}
19
20impl McmcResult {
21    /// Per-dimension sample mean.
22    pub fn mean(&self) -> Vec<f64> {
23        let mut m = vec![0.0; self.dim];
24        for i in 0..self.n_samples {
25            for j in 0..self.dim {
26                m[j] += self.samples[i * self.dim + j];
27            }
28        }
29        for v in m.iter_mut() {
30            *v /= self.n_samples.max(1) as f64;
31        }
32        m
33    }
34
35    /// Per-dimension sample variance (population).
36    pub fn variance(&self) -> Vec<f64> {
37        let m = self.mean();
38        let mut v = vec![0.0; self.dim];
39        for i in 0..self.n_samples {
40            for j in 0..self.dim {
41                let d = self.samples[i * self.dim + j] - m[j];
42                v[j] += d * d;
43            }
44        }
45        for x in v.iter_mut() {
46            *x /= self.n_samples.max(1) as f64;
47        }
48        v
49    }
50}
51
52struct Rng(u64);
53impl Rng {
54    fn unit(&mut self) -> f64 {
55        self.0 = self
56            .0
57            .wrapping_mul(6364136223846793005)
58            .wrapping_add(1442695040888963407);
59        ((self.0 >> 11) as f64) / ((1u64 << 53) as f64)
60    }
61    fn gaussian(&mut self) -> f64 {
62        let u1 = self.unit().max(1e-12);
63        let u2 = self.unit();
64        (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
65    }
66}
67
68/// Random-walk Metropolis-Hastings. `log_density(x)` returns the target's log-density
69/// (up to a constant) at a `dim`-vector; `initial` is the starting point;
70/// `proposal_std` the per-dimension Gaussian step; `n_samples` post-burn-in draws;
71/// `burn_in` discarded warm-up steps. Deterministic given `seed`.
72pub fn metropolis_hastings<F>(
73    log_density: F,
74    initial: &[f64],
75    proposal_std: f64,
76    n_samples: usize,
77    burn_in: usize,
78    seed: u64,
79) -> McmcResult
80where
81    F: Fn(&[f64]) -> f64,
82{
83    let dim = initial.len();
84    let mut rng = Rng(seed ^ 0x9E3779B97F4A7C15);
85    let mut x = initial.to_vec();
86    let mut lp = log_density(&x);
87    let mut prop = vec![0.0; dim];
88    let mut samples = Vec::with_capacity(n_samples * dim);
89    let mut accepted = 0u64;
90    let total = burn_in + n_samples;
91
92    for step in 0..total {
93        // Propose x' = x + N(0, proposal_std²).
94        for d in 0..dim {
95            prop[d] = x[d] + proposal_std * rng.gaussian();
96        }
97        let lp_prop = log_density(&prop);
98        // Accept with probability min(1, exp(lp' − lp)).
99        let accept = lp_prop >= lp || rng.unit() < (lp_prop - lp).exp();
100        if accept {
101            x.copy_from_slice(&prop);
102            lp = lp_prop;
103            if step >= burn_in {
104                accepted += 1;
105            }
106        }
107        if step >= burn_in {
108            samples.extend_from_slice(&x);
109        }
110    }
111
112    McmcResult {
113        samples,
114        n_samples,
115        dim,
116        acceptance_rate: if n_samples > 0 {
117            accepted as f64 / n_samples as f64
118        } else {
119            0.0
120        },
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn samples_a_1d_gaussian() {
130        // Target N(3, 2²): log p(x) = −(x−3)²/(2·4).
131        let target = |x: &[f64]| -(x[0] - 3.0) * (x[0] - 3.0) / 8.0;
132        let r = metropolis_hastings(target, &[0.0], 2.0, 40_000, 5_000, 1);
133        let m = r.mean();
134        let v = r.variance();
135        assert!((m[0] - 3.0).abs() < 0.15, "mean {}", m[0]);
136        assert!((v[0] - 4.0).abs() < 0.6, "var {}", v[0]);
137        assert!(r.acceptance_rate > 0.1 && r.acceptance_rate < 0.95);
138    }
139
140    #[test]
141    fn samples_a_2d_gaussian_mean() {
142        // Independent N(1,1)×N(-2,1).
143        let target = |x: &[f64]| -(x[0] - 1.0).powi(2) / 2.0 - (x[1] + 2.0).powi(2) / 2.0;
144        let r = metropolis_hastings(target, &[0.0, 0.0], 1.0, 40_000, 5_000, 7);
145        let m = r.mean();
146        assert!(
147            (m[0] - 1.0).abs() < 0.15 && (m[1] + 2.0).abs() < 0.15,
148            "mean {m:?}"
149        );
150    }
151
152    #[test]
153    fn unnormalized_target_is_fine() {
154        // log p need not be normalized — a constant offset must not change samples.
155        let target = |x: &[f64]| 12345.0 - (x[0]).powi(2) / 2.0; // N(0,1) + constant
156        let r = metropolis_hastings(target, &[0.0], 1.5, 20_000, 3_000, 2);
157        assert!(r.mean()[0].abs() < 0.15);
158        assert!((r.variance()[0] - 1.0).abs() < 0.3);
159    }
160}
qualia_core_db/solvers/learning/sampling/mcmc.rs

qualia_core_db/solvers/learning/sampling/
mcmc.rs
