ferrotherm 0.32.0

//! Chromatic block-Gibbs sampling.
//!
//! One sweep updates every node exactly once, color class by color class. Within a class, node
//! updates are conditionally independent (no two adjacent) — the parallelism a TSU exploits in
//! physics and a GPU exploits in threads; here the classes are simple loops, kept in the same
//! order so CPU, WebGPU, and device runs are cross-checkable draw for draw.

use crate::graph::Graph;
use crate::ledger::Ledger;
use crate::rng::Pcg;

// The logistic the module docs write as sigma(2 beta f_i). Only the conditional test
// exercises it, so it is dead in a non-test build — kept because it is the executable
// statement of the formula graph.rs and kernel.rs both cite in prose.
#[cfg_attr(not(test), allow(dead_code))]
#[inline]
fn sigma(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

pub struct Sampler<'g> {
    pub g: &'g Graph,
    pub beta: f64,
    pub s: Vec<i8>,
    pub rng: Pcg,
    /// How many threads the last parallel sweep actually used. See [`Sampler::threads_used`].
    threads_used: usize,
    /// Nodes whose value is held fixed (conditioning / "clamping"); sweeps skip them.
    pub clamped: Vec<bool>,
    /// Base seed for the parallel path's per-(sweep, class, chunk) RNG streams.
    ///
    /// Kept on every target, including wasm32 where [`Sampler::sweep_par`] is the serial twin and
    /// nothing reads it. Making the FIELD conditional would make the struct's shape depend on the
    /// target, and the two things that then differ are exactly the things this crate promises do
    /// not: a `Sampler` built the same way would not be the same object across a build boundary.
    /// The dead-code allowance is the cheaper of the two, and it is scoped to the one target where
    /// the field is genuinely unread.
    #[cfg_attr(target_arch = "wasm32", allow(dead_code))]
    par_seed: u64,
    /// Sweeps completed via the parallel path (advances its stream derivation).
    par_sweeps: u64,
}

impl<'g> Sampler<'g> {
    pub fn new(g: &'g Graph, beta: f64, seed: u64) -> Self {
        let mut rng = Pcg::new(seed, 0x5EED);
        let s = (0..g.n).map(|_| rng.spin(0.5)).collect();
        Sampler {
            g,
            beta,
            s,
            rng,
            clamped: vec![false; g.n],
            par_seed: seed,
            par_sweeps: 0,
            threads_used: 1,
        }
    }

    /// Clamp node i to value v (observation / conditioning input).
    pub fn clamp(&mut self, i: usize, v: i8) {
        debug_assert!(v == 1 || v == -1);
        self.s[i] = v;
        self.clamped[i] = true;
    }

    pub fn unclamp(&mut self, i: usize) {
        self.clamped[i] = false;
    }

    /// One full chromatic sweep (every free node updated once). If a ledger is given, it is
    /// charged one Gibbs cycle per free node — the device-side price of this sweep.
    pub fn sweep(&mut self, ledger: Option<&mut Ledger>) {
        let mut updated = 0u64;
        for class in &self.g.classes {
            for &iu in class {
                let i = iu as usize;
                if self.clamped[i] {
                    continue;
                }
                let f = self.g.field(i, &self.s);
                let p_up = crate::kernel::p_up(f, self.beta);
                self.s[i] = self.rng.spin(p_up);
                updated += 1;
            }
        }
        if let Some(l) = ledger {
            l.samples += updated;
        }
    }

    /// Run `n` sweeps.
    pub fn sweeps(&mut self, n: usize, mut ledger: Option<&mut Ledger>) {
        for _ in 0..n {
            self.sweep(ledger.as_deref_mut());
        }
    }

    /// One full chromatic sweep across `threads` OS threads — the performance core.
    ///
    /// Within a color class every node's conditional is independent (no two adjacent), so the
    /// class is split into contiguous chunks, each updated by its own thread reading the shared
    /// spin field and writing only its own chunk's nodes. Reads touch only OTHER-color nodes,
    /// which no thread writes during this phase, so the access pattern is race-free by
    /// construction of the coloring.
    ///
    /// Determinism: each (sweep, class, chunk) gets its own counter-derived RNG stream, so the
    /// result is bit-reproducible for a fixed (seed, threads). A different thread count is a
    /// different, equally valid sample path (document the thread count next to the seed).
    #[cfg(not(target_arch = "wasm32"))]
    pub fn sweep_par(&mut self, threads: usize, ledger: Option<&mut Ledger>) {
        assert!(threads >= 1);
        let beta = self.beta;
        let g = self.g;
        let sweep_idx = self.par_sweeps;
        let base = self.par_seed;
        let mut updated = 0u64;
        for (ci, class) in g.classes.iter().enumerate() {
            let chunk = class.len().div_ceil(threads);
            if chunk == 0 {
                continue;
            }
            // SAFETY: chunks are disjoint index sets within one color class; every write target
            // is unique to one thread, and every read is either a bias, an other-color neighbour
            // (not written this phase), or the thread's own not-yet-updated node.
            let sp = self.s.as_mut_ptr() as usize;
            let clamped = &self.clamped;
            std::thread::scope(|scope| {
                for (ti, part) in class.chunks(chunk).enumerate() {
                    let part: &[u32] = part;
                    scope.spawn(move || {
                        let mut rng = Pcg::new(
                            base ^ sweep_idx.wrapping_mul(0x9E3779B97F4A7C15) ^ (ci as u64) << 32,
                            0xC0DE ^ ti as u64,
                        );
                        let s_ptr = sp as *mut i8;
                        for &iu in part {
                            let i = iu as usize;
                            if clamped[i] {
                                continue;
                            }
                            let mut f = g.h[i];
                            for k in g.offset[i]..g.offset[i + 1] {
                                f += g.w[k] * unsafe { *s_ptr.add(g.nbr[k] as usize) } as f64;
                            }
                            let p_up = crate::kernel::p_up(f, beta);
                            unsafe {
                                *s_ptr.add(i) = rng.spin(p_up);
                            }
                        }
                    });
                }
            });
            updated += class.iter().filter(|&&iu| !self.clamped[iu as usize]).count() as u64;
        }
        self.par_sweeps += 1;
        // The PEAK number of chunks any one class was actually split into -- not `threads`, and not
        // `min(threads, biggest class)` either, which is what the first version of this computed
        // and which over-reports: five nodes across four threads is a chunk of two, so three
        // threads run and not four. A caller dividing throughput by this number needs the count
        // that ran or its per-thread figure is quietly wrong.
        self.threads_used = g
            .classes
            .iter()
            .map(|c| {
                let chunk = c.len().div_ceil(threads);
                if chunk == 0 { 0 } else { c.len().div_ceil(chunk) }
            })
            .max()
            .unwrap_or(1)
            .max(1);
        if let Some(l) = ledger {
            l.samples += updated;
        }
    }

    /// The wasm twin: serial, because there are no threads to spread across.
    ///
    /// `wasm32-unknown-unknown` has a std whose `thread::spawn` COMPILES and then panics at
    /// runtime with "operation not supported on this platform" — the worst of both, since nothing
    /// stops you shipping it and the failure arrives in a browser rather than in a build. So the
    /// browser gets the serial sweep, and [`Self::threads_used`] reports 1 so a caller learns what
    /// actually ran instead of trusting the number it asked for.
    ///
    /// This is not a silent downgrade. Running serially and SAYING one thread is a different thing
    /// from running serially and reporting eight.
    #[cfg(target_arch = "wasm32")]
    pub fn sweep_par(&mut self, threads: usize, ledger: Option<&mut Ledger>) {
        assert!(threads >= 1);
        self.sweep(ledger);
        self.par_sweeps += 1;
        self.threads_used = 1;
    }

    /// How many threads the last [`Self::sweep_par`] actually used, at its peak.
    ///
    /// The largest number of chunks any single colour class was split into. Not the number asked
    /// for: a class of five across four threads is a chunk of two, so three threads run. Always 1
    /// in a browser, whatever was requested. A caller that reports throughput per thread needs
    /// this number rather than the one it passed in.
    pub fn threads_used(&self) -> usize {
        self.threads_used
    }

    /// Run `n` parallel sweeps.
    pub fn sweeps_par(&mut self, n: usize, threads: usize, mut ledger: Option<&mut Ledger>) {
        for _ in 0..n {
            self.sweep_par(threads, ledger.as_deref_mut());
        }
    }

    /// Read the full state (device price: one read per node). Prefer [`Self::read_subset`]:
    /// full-state readback is the crossings-tax regime.
    pub fn read_all(&self, ledger: Option<&mut Ledger>) -> Vec<i8> {
        if let Some(l) = ledger {
            l.reads += self.g.n as u64;
        }
        self.s.clone()
    }

    /// Read only the named nodes (e.g. action bits).
    pub fn read_subset(&self, idx: &[usize], ledger: Option<&mut Ledger>) -> Vec<i8> {
        if let Some(l) = ledger {
            l.reads += idx.len() as u64;
        }
        idx.iter().map(|&i| self.s[i]).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::sigma;
    use super::*;
    use crate::graph::GraphBuilder;

    /// The sampler's stationary distribution must match the exact Boltzmann distribution on an
    /// enumerable system. 4-node cycle, mixed couplings and biases, TV < 0.02.
    #[test]
    fn matches_exact_boltzmann() {
        let mut gb = GraphBuilder::new(4);
        gb.couple(0, 1, 0.7);
        gb.couple(1, 2, -0.4);
        gb.couple(2, 3, 0.55);
        gb.couple(3, 0, 0.3);
        gb.bias(0, 0.2);
        gb.bias(2, -0.35);
        let g = gb.build();
        let beta = 0.9;

        // exact
        let mut z = 0.0;
        let mut p_exact = [0.0f64; 16];
        for m in 0..16u32 {
            let s: Vec<i8> = (0..4).map(|b| if m >> b & 1 == 1 { 1 } else { -1 }).collect();
            let w = (-beta * g.energy(&s)).exp();
            p_exact[m as usize] = w;
            z += w;
        }
        for p in p_exact.iter_mut() {
            *p /= z;
        }

        // sampled
        let mut smp = Sampler::new(&g, beta, 0xC0FFEE);
        smp.sweeps(200, None); // burn-in
        let mut counts = [0u64; 16];
        let n_samples = 200_000;
        for _ in 0..n_samples {
            smp.sweep(None);
            let mut m = 0usize;
            for b in 0..4 {
                if smp.s[b] == 1 {
                    m |= 1 << b;
                }
            }
            counts[m] += 1;
        }
        let tv: f64 = (0..16)
            .map(|m| (counts[m] as f64 / n_samples as f64 - p_exact[m]).abs())
            .sum::<f64>()
            / 2.0;
        assert!(tv < 0.02, "TV distance to exact Boltzmann = {tv}");
    }

    /// The parallel path must satisfy the same physics standard as the sequential one: Onsager's
    /// exact magnetization on the 2D lattice, and bit-reproducibility for fixed (seed, threads).
    #[test]
    fn parallel_sweep_physics_and_determinism() {
        let g = crate::ising::lattice2d(48, 1.0);
        let beta = 0.6;
        let mut smp = Sampler::new(&g, beta, 0x9A7);
        for s in smp.s.iter_mut() {
            *s = 1;
        }
        smp.sweeps_par(2000, 8, None);
        let mut acc = 0.0;
        let reads = 2000;
        for _ in 0..reads {
            smp.sweep_par(8, None);
            let m: i64 = smp.s.iter().map(|&v| v as i64).sum();
            acc += (m as f64 / g.n as f64).abs();
        }
        let m = acc / reads as f64;
        let exact = crate::ising::onsager_m(beta);
        assert!((m - exact).abs() < 0.01, "parallel |M| {m:.4} vs Onsager {exact:.4}");
        // determinism for fixed (seed, threads)
        let mut a = Sampler::new(&g, beta, 0x1234);
        let mut b = Sampler::new(&g, beta, 0x1234);
        a.sweeps_par(50, 4, None);
        b.sweeps_par(50, 4, None);
        assert_eq!(a.s, b.s, "same (seed, threads) must reproduce bit-identically");
    }

    /// Clamped nodes must never change and must steer the conditional distribution.
    #[test]
    fn clamping_conditions() {
        let mut gb = GraphBuilder::new(2);
        gb.couple(0, 1, 1.5);
        let g = gb.build();
        let mut smp = Sampler::new(&g, 1.0, 7);
        smp.clamp(0, 1);
        let mut up = 0u64;
        let n = 20_000;
        for _ in 0..n {
            smp.sweep(None);
            assert_eq!(smp.s[0], 1);
            if smp.s[1] == 1 {
                up += 1;
            }
        }
        // exact: P(s1=+1 | s0=+1) = sigma(2*beta*J) = sigma(3.0)
        let want = sigma(3.0);
        let got = up as f64 / n as f64;
        assert!((got - want).abs() < 0.01, "got {got}, want {want}");
    }
}