wai-quantum 0.3.36

//! Neural quantum states — `wai.quantum.nqs`.
//!
//! A neural network can stand in for a many-body wavefunction: it gives the
//! amplitude of each configuration, and variational Monte Carlo trains it
//! toward the ground state. Results from such networks are offered as
//! evidence about systems no exact method reaches. This module is the method
//! in its checkable form, deterministic and held to exact answers.
//!
//! - **Models.** [`SpinModel`] is a spin-½ Hamiltonian on any graph, with
//!   exchange, Ising, and field terms. Presets are Heisenberg and
//!   transverse-field Ising, on chains and square lattices.
//! - **The network.** [`Rbm`] is a restricted Boltzmann machine with `α·n`
//!   hidden units. On a bipartite lattice it can carry Marshall's sign, which
//!   makes the Heisenberg ground state positive in what remains.
//! - **Variational Monte Carlo.** [`Vmc`] samples `|ψ|²` by Metropolis
//!   chains: single flips, or swaps of antiparallel spins at fixed
//!   magnetisation. Or it enumerates every configuration (at most 24 spins)
//!   for expectations with no Monte Carlo noise. Each step is stochastic
//!   reconfiguration, `θ ← θ − η (S + λ)⁻¹ F`, solved in parameter space or
//!   in sample space, whichever is smaller; the two give the same step.
//! - **Referees.** [`SpinModel::exact_ground_energy`] is Lanczos on all `2ⁿ`
//!   amplitudes (at most 22 spins). [`SpinModel::opsum`] gives the model to
//!   DMRG on its exact MPO, for lattices past exact diagonalisation.
//!
//! Chains draw from their own seeded generators, and every sum runs in a
//! fixed order, so a run is the same bits at any thread count.
//!
//! # Checked
//!
//! - **The parts.** The model's exact diagonalisation equals the chain
//!   referees. Its operator sum's MPO equals its dense matrix on a 3 × 3
//!   lattice with every kind of term. The network's amplitude ratios equal
//!   ratios of its amplitudes to `10⁻¹²`, Marshall sign included, and its
//!   log-derivatives agree with central differences. Both spaces take the
//!   same reconfiguration step.
//! - **Exact expectations.** On an 8-site Heisenberg chain, 200 steps with
//!   every configuration enumerated reach the exact energy to `10⁻⁷`.
//! - **Sampling.** Sampled energies agree with enumerated ones within their
//!   error bars, and steps repeat bit for bit at any thread count.
//! - **Against exact answers** (`nqs_ground` example; relative error of the
//!   final estimate, with its statistical error):
//!
//!   | model             | lattice      | reference               | α | steps | relative error          |
//!   |-------------------|--------------|-------------------------|---|-------|-------------------------|
//!   | Heisenberg        | chain of 12  | exact diagonalisation   | 2 | 300   | 7.7·10⁻⁵ ± 7·10⁻⁵       |
//!   | Heisenberg        | 4 × 4 open   | exact diagonalisation   | 2 | 600   | 4.6·10⁻⁴ ± 2·10⁻⁴       |
//!   | Ising, `h = 3.044`| 4 × 4 open   | exact diagonalisation   | 1 | 400   | −3·10⁻⁶ ± 1·10⁻⁵        |
//!   | Heisenberg        | 6 × 6 open   | DMRG, bond 256          | 2 | 800   | 3.6·10⁻³ ± 3·10⁻⁴       |
//!
//!   Each run used 1000 samples a step. The 6 × 6 DMRG reference is itself
//!   an upper bound (discarded weight `2·10⁻⁵`). The network was still
//!   descending when it stopped there.
//!
//! # Honest boundaries
//!
//! - **A real network.** The sign structure must be known, as Marshall's is
//!   on bipartite lattices. Frustrated models, whose signs are not known,
//!   need complex parameters or a learned sign, which are not here.
//! - **Error bars ignore autocorrelation.** The reported error is the
//!   standard error of uncorrelated samples. Chains thinned by two sweeps are
//!   correlated, so the true error is larger.
//! - **Dense reconfiguration.** `S` (or its sample-space twin) is formed and
//!   factored in full each step. That bounds the network and sample sizes
//!   long before iterative solvers would.

use crate::quantum_dmrg::{Op, OpSum, default_threads, lanczos};
use crate::repro::{exp, ln};

// ---------------------------------------------------------------------------
// Spin models
// ---------------------------------------------------------------------------

/// A spin-½ model on any graph, in Pauli matrices:
/// `H = Σ J⊥ (σˣσˣ + σʸσʸ) + Σ J_z σᶻσᶻ + Σ hₓ σˣ + Σ h_z σᶻ`.
/// Spin `+1` is `|↑⟩`, the first basis state.
#[derive(Clone, Debug, PartialEq)]
pub struct SpinModel {
    pub n: usize,
    /// `(i, j, J⊥)`: `J⊥ (σˣᵢσˣⱼ + σʸᵢσʸⱼ)`, which swaps antiparallel spins with
    /// amplitude `2J⊥`.
    pub exchange: Vec<(usize, usize, f64)>,
    /// `(i, j, J_z)`: `J_z σᶻᵢσᶻⱼ`.
    pub zz: Vec<(usize, usize, f64)>,
    /// `(i, hₓ)`: `hₓ σˣᵢ`.
    pub x: Vec<(usize, f64)>,
    /// `(i, h_z)`: `h_z σᶻᵢ`.
    pub z: Vec<(usize, f64)>,
}

/// The bonds of an open (or periodic) chain.
pub fn chain_edges(n: usize, periodic: bool) -> Vec<(usize, usize)> {
    let mut e: Vec<(usize, usize)> = (0..n.saturating_sub(1)).map(|i| (i, i + 1)).collect();
    if periodic && n > 2 {
        e.push((n - 1, 0));
    }
    e
}

/// The nearest-neighbour bonds of an `lx × ly` square lattice, site
/// `y·lx + x`, open or periodic in both directions.
pub fn square_edges(lx: usize, ly: usize, periodic: bool) -> Vec<(usize, usize)> {
    let mut e = Vec::new();
    for y in 0..ly {
        for x in 0..lx {
            let i = y * lx + x;
            if x + 1 < lx {
                e.push((i, i + 1));
            } else if periodic && lx > 2 {
                e.push((i, y * lx));
            }
            if y + 1 < ly {
                e.push((i, i + lx));
            } else if periodic && ly > 2 {
                e.push((i, x));
            }
        }
    }
    e
}

/// The checkerboard sublattice of an `lx × ly` lattice (a chain is `ly = 1`).
pub fn checkerboard(lx: usize, ly: usize) -> Vec<bool> {
    (0..lx * ly).map(|i| (i % lx + i / lx).is_multiple_of(2)).collect()
}

impl SpinModel {
    /// The Heisenberg model `J Σ Sᵢ·Sⱼ` (with `S = σ/2`) on `edges`.
    pub fn heisenberg(n: usize, edges: &[(usize, usize)], j: f64) -> SpinModel {
        SpinModel { n, exchange: edges.iter().map(|&(a, b)| (a, b, j / 4.0)).collect(), zz: edges.iter().map(|&(a, b)| (a, b, j / 4.0)).collect(), x: Vec::new(), z: Vec::new() }
    }

    /// The transverse-field Ising model `−J Σ σᶻσᶻ − h Σ σˣ` on `edges`.
    pub fn ising(n: usize, edges: &[(usize, usize)], j: f64, h: f64) -> SpinModel {
        SpinModel { n, exchange: Vec::new(), zz: edges.iter().map(|&(a, b)| (a, b, -j)).collect(), x: (0..n).map(|i| (i, -h)).collect(), z: Vec::new() }
    }

    /// Whether `Σ σᶻ` is conserved (no transverse field).
    pub fn conserves_magnetisation(&self) -> bool {
        self.x.is_empty()
    }

    /// The model as an operator sum, for DMRG on its MPO.
    pub fn opsum(&self) -> OpSum {
        let mut s = OpSum::new(self.n);
        for &(i, j, c) in &self.exchange {
            // σˣσˣ + σʸσʸ = 2(σ⁺σ⁻ + σ⁻σ⁺).
            s.add(2.0 * c, &[(i, Op::plus()), (j, Op::minus())]);
            s.add(2.0 * c, &[(i, Op::minus()), (j, Op::plus())]);
        }
        for &(i, j, c) in &self.zz {
            s.add(c, &[(i, Op::z()), (j, Op::z())]);
        }
        for &(i, c) in &self.x {
            s.add(c, &[(i, Op::x())]);
        }
        for &(i, c) in &self.z {
            s.add(c, &[(i, Op::z())]);
        }
        s
    }

    /// `out = H inp` on all `2ⁿ` amplitudes, site 0 the most significant bit
    /// (a set bit is `|↓⟩`).
    fn apply_dense(&self, inp: &[f64], out: &mut [f64]) {
        let n = self.n;
        let bit = |i: usize| 1usize << (n - 1 - i);
        let spin = |x: usize, i: usize| if x & bit(i) == 0 { 1.0 } else { -1.0 };
        out.iter_mut().for_each(|o| *o = 0.0);
        for (x, &amp) in inp.iter().enumerate() {
            if amp == 0.0 {
                continue;
            }
            let mut diag = 0.0;
            for &(i, j, c) in &self.zz {
                diag += c * spin(x, i) * spin(x, j);
            }
            for &(i, c) in &self.z {
                diag += c * spin(x, i);
            }
            out[x] += diag * amp;
            for &(i, j, c) in &self.exchange {
                if spin(x, i) != spin(x, j) {
                    out[x ^ bit(i) ^ bit(j)] += 2.0 * c * amp;
                }
            }
            for &(i, c) in &self.x {
                out[x ^ bit(i)] += c * amp;
            }
        }
    }

    /// The exact ground energy by Lanczos on all `2ⁿ` amplitudes (at most 22
    /// sites), restricted to total `Σ σᶻ = magnetisation` when given (the
    /// model must conserve it).
    pub fn exact_ground_energy(&self, magnetisation: Option<i32>) -> Option<f64> {
        let n = self.n;
        if n == 0 || n > 22 {
            return None;
        }
        if magnetisation.is_some() && !self.conserves_magnetisation() {
            return None;
        }
        let dim = 1usize << n;
        let in_sector = |x: usize| magnetisation.is_none_or(|m| n as i32 - 2 * x.count_ones() as i32 == m);
        let start: Vec<f64> = (0..dim as u64).map(|x| if in_sector(x as usize) { 1.0 + (x.wrapping_mul(2_654_435_761) % 1000) as f64 / 1000.0 } else { 0.0 }).collect();
        if start.iter().all(|&v| v == 0.0) {
            return None;
        }
        let apply = |inp: &[f64], out: &mut [f64]| self.apply_dense(inp, out);
        Some(lanczos(&apply, &start, 60, 30, 1e-9).0)
    }
}

// ---------------------------------------------------------------------------
// The restricted Boltzmann machine
// ---------------------------------------------------------------------------

/// `ln cosh x`, stable for large `|x|`.
fn lncosh(x: f64) -> f64 {
    let a = x.abs();
    a + ln(1.0 + exp(-2.0 * a)) - core::f64::consts::LN_2
}

fn tanh(x: f64) -> f64 {
    let e = exp(-2.0 * x.abs());
    let t = (1.0 - e) / (1.0 + e);
    if x < 0.0 { -t } else { t }
}

/// A real restricted Boltzmann machine: `ψ(s) = sign(s) · exp(Σ aᵢsᵢ) ·
/// Π_j cosh(b_j + Σ W_jᵢ sᵢ)`, with `m = α·n` hidden units. On a bipartite
/// lattice the optional sign is Marshall's, `(−1)` to the number of up spins
/// on sublattice A, which makes the Heisenberg ground state positive in the
/// rest.
#[derive(Clone, Debug, PartialEq)]
pub struct Rbm {
    pub n: usize,
    pub m: usize,
    pub a: Vec<f64>,
    pub b: Vec<f64>,
    /// `W[j·n + i]`.
    pub w: Vec<f64>,
    pub sublattice: Option<Vec<bool>>,
}

/// A deterministic generator (splitmix64).
#[derive(Clone, Debug, PartialEq)]
struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        z ^ (z >> 31)
    }

    /// Uniform on `[0, 1)`.
    fn uniform(&mut self) -> f64 {
        (self.next() >> 11) as f64 / 9_007_199_254_740_992.0
    }

    fn below(&mut self, k: usize) -> usize {
        (self.next() % k as u64) as usize
    }
}

impl Rbm {
    /// `α·n` hidden units, weights uniform in `±scale`, biases zero.
    pub fn new(n: usize, alpha: usize, scale: f64, seed: u64) -> Rbm {
        let m = alpha.max(1) * n;
        let mut rng = Rng(seed);
        let w = (0..m * n).map(|_| scale * (2.0 * rng.uniform() - 1.0)).collect();
        Rbm { n, m, a: vec![0.0; n], b: vec![0.0; m], w, sublattice: None }
    }

    /// Multiply by Marshall's sign for this sublattice.
    pub fn with_marshall_sign(mut self, sublattice: Vec<bool>) -> Rbm {
        assert_eq!(sublattice.len(), self.n);
        self.sublattice = Some(sublattice);
        self
    }

    /// The number of variational parameters.
    pub fn params(&self) -> usize {
        self.n + self.m + self.n * self.m
    }

    fn theta(&self, s: &[i8]) -> Vec<f64> {
        (0..self.m)
            .map(|j| {
                let row = &self.w[j * self.n..(j + 1) * self.n];
                let mut t = self.b[j];
                for (wi, &si) in row.iter().zip(s) {
                    t += wi * f64::from(si);
                }
                t
            })
            .collect()
    }

    /// `ln |ψ(s)|` up to a constant.
    pub fn log_amplitude(&self, s: &[i8]) -> f64 {
        let mut l: f64 = self.a.iter().zip(s).map(|(a, &si)| a * f64::from(si)).sum();
        for t in self.theta(s) {
            l += lncosh(t);
        }
        l
    }

    /// The sign of `ψ(s)`.
    pub fn sign(&self, s: &[i8]) -> f64 {
        match &self.sublattice {
            Some(sub) => {
                let ups = s.iter().zip(sub).filter(|&(&si, &a)| a && si > 0).count();
                if ups.is_multiple_of(2) { 1.0 } else { -1.0 }
            }
            None => 1.0,
        }
    }

    /// `ψ(s′)/ψ(s)` where `s′` is `s` with `flips` flipped, given `θ(s)`.
    fn ratio(&self, s: &[i8], theta: &[f64], flips: &[usize]) -> f64 {
        let mut l = 0.0;
        for &i in flips {
            l -= 2.0 * self.a[i] * f64::from(s[i]);
        }
        for (j, &t) in theta.iter().enumerate() {
            let mut d = 0.0;
            for &i in flips {
                d -= 2.0 * self.w[j * self.n + i] * f64::from(s[i]);
            }
            l += lncosh(t + d) - lncosh(t);
        }
        let mut r = exp(l);
        if let Some(sub) = &self.sublattice
            && flips.iter().filter(|&&i| sub[i]).count() % 2 == 1
        {
            r = -r;
        }
        r
    }

    /// `∂ ln ψ / ∂(a, b, W)` at `s`, given `θ(s)`.
    fn log_derivatives(&self, s: &[i8], theta: &[f64], out: &mut [f64]) {
        let (n, m) = (self.n, self.m);
        for (o, &si) in out[..n].iter_mut().zip(s) {
            *o = f64::from(si);
        }
        for (j, &th) in theta.iter().enumerate() {
            let t = tanh(th);
            out[n + j] = t;
            for (o, &si) in out[n + m + j * n..n + m + (j + 1) * n].iter_mut().zip(s) {
                *o = t * f64::from(si);
            }
        }
    }

    fn shift(&mut self, delta: &[f64], step: f64) {
        let (n, m) = (self.n, self.m);
        for (x, d) in self.a.iter_mut().zip(&delta[..n]) {
            *x -= step * d;
        }
        for (x, d) in self.b.iter_mut().zip(&delta[n..n + m]) {
            *x -= step * d;
        }
        for (x, d) in self.w.iter_mut().zip(&delta[n + m..]) {
            *x -= step * d;
        }
    }
}

/// `E_loc(s) = Σ_{s′} H_{ss′} ψ(s′)/ψ(s)`.
fn local_energy(model: &SpinModel, rbm: &Rbm, s: &[i8], theta: &[f64]) -> f64 {
    let mut e = 0.0;
    for &(i, j, c) in &model.zz {
        e += c * f64::from(s[i]) * f64::from(s[j]);
    }
    for &(i, c) in &model.z {
        e += c * f64::from(s[i]);
    }
    for &(i, j, c) in &model.exchange {
        if s[i] != s[j] {
            e += 2.0 * c * rbm.ratio(s, theta, &[i, j]);
        }
    }
    for &(i, c) in &model.x {
        e += c * rbm.ratio(s, theta, &[i]);
    }
    e
}

// ---------------------------------------------------------------------------
// Variational Monte Carlo
// ---------------------------------------------------------------------------

/// Settings for [`Vmc`].
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct VmcConfig {
    /// Samples per iteration, over all chains.
    pub samples: usize,
    pub chains: usize,
    /// Sweeps (`n` moves each) discarded when a chain starts.
    pub burn: usize,
    /// Sweeps between kept samples.
    pub thin: usize,
    /// Enumerate every configuration (at most 24 spins) instead of sampling:
    /// exact expectations, no Monte Carlo noise.
    pub exact: bool,
    /// Keep `Σ σᶻ` at this value: moves swap antiparallel spins. The model
    /// must conserve it.
    pub magnetisation: Option<i32>,
    /// Learning rate of stochastic reconfiguration.
    pub lr: f64,
    /// Diagonal shift added to the metric `S`.
    pub shift: f64,
    pub seed: u64,
    /// Threads over chains and configurations; the result does not depend on
    /// it.
    pub threads: usize,
}

impl Default for VmcConfig {
    fn default() -> VmcConfig {
        VmcConfig { samples: 2000, chains: 8, burn: 50, thin: 2, exact: false, magnetisation: None, lr: 0.05, shift: 1e-3, seed: 1, threads: default_threads() }
    }
}

/// One iteration's estimate.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Estimate {
    /// `⟨E_loc⟩`.
    pub energy: f64,
    /// `⟨(E_loc − E)²⟩`: zero exactly at an eigenstate.
    pub variance: f64,
    /// The statistical error of `energy`, from the variance and the sample
    /// count, ignoring autocorrelation (zero when exact).
    pub error: f64,
    /// Fraction of Metropolis moves accepted (one when exact).
    pub acceptance: f64,
}

/// Variational Monte Carlo with stochastic reconfiguration on an RBM.
#[derive(Clone, Debug, PartialEq)]
pub struct Vmc {
    pub model: SpinModel,
    pub rbm: Rbm,
    pub cfg: VmcConfig,
    chains: Vec<(Vec<i8>, Rng)>,
}

/// Samples: configurations, their weights (summing to one), local energies
/// and log-derivatives (`params` per row).
struct Batch {
    weights: Vec<f64>,
    energies: Vec<f64>,
    derivs: Vec<f64>,
    acceptance: f64,
}

impl Vmc {
    pub fn new(model: SpinModel, rbm: Rbm, cfg: VmcConfig) -> Vmc {
        assert_eq!(model.n, rbm.n, "the model and the machine differ in size");
        if cfg.magnetisation.is_some() {
            assert!(model.conserves_magnetisation(), "the model does not conserve the magnetisation");
        }
        let n = model.n;
        let chains = (0..cfg.chains.max(1))
            .map(|c| {
                let mut rng = Rng(cfg.seed ^ (c as u64).wrapping_mul(0xa076_1d64_78bd_642f));
                let s = start_configuration(n, cfg.magnetisation, &mut rng);
                (s, rng)
            })
            .collect();
        let mut v = Vmc { model, rbm, cfg, chains };
        if !cfg.exact {
            for c in 0..v.chains.len() {
                let (mut s, mut rng) = v.chains[c].clone();
                metropolis(&v.rbm, &mut s, &mut rng, cfg.burn * n, cfg.magnetisation.is_some());
                v.chains[c] = (s, rng);
            }
        }
        v
    }

    fn batch(&mut self) -> Batch {
        if self.cfg.exact { self.enumerate() } else { self.sample() }
    }

    fn enumerate(&self) -> Batch {
        let n = self.model.n;
        assert!(n <= 24, "exact enumeration takes at most 24 spins");
        let configs: Vec<usize> = (0..1usize << n).filter(|&x| self.cfg.magnetisation.is_none_or(|m| n as i32 - 2 * x.count_ones() as i32 == m)).collect();
        let p = self.rbm.params();
        let k = configs.len();
        let mut logs = vec![0.0; k];
        let mut energies = vec![0.0; k];
        let mut derivs = vec![0.0; k * p];
        let threads = self.cfg.threads.max(1).min(k.div_ceil(64)).max(1);
        let per = k.div_ceil(threads);
        let work = |range: core::ops::Range<usize>, logs: &mut [f64], energies: &mut [f64], derivs: &mut [f64]| {
            let start = range.start;
            for idx in range {
                let x = configs[idx];
                let s: Vec<i8> = (0..n).map(|i| if x >> (n - 1 - i) & 1 == 0 { 1 } else { -1 }).collect();
                let theta = self.rbm.theta(&s);
                let r = idx - start;
                logs[r] = self.rbm.log_amplitude(&s);
                energies[r] = local_energy(&self.model, &self.rbm, &s, &theta);
                self.rbm.log_derivatives(&s, &theta, &mut derivs[r * p..(r + 1) * p]);
            }
        };
        if threads <= 1 || cfg!(target_arch = "wasm32") {
            work(0..k, &mut logs, &mut energies, &mut derivs);
        } else {
            std::thread::scope(|scope| {
                let mut lrest: &mut [f64] = &mut logs;
                let mut erest: &mut [f64] = &mut energies;
                let mut drest: &mut [f64] = &mut derivs;
                for t in 0..threads {
                    let lo = t * per;
                    let hi = ((t + 1) * per).min(k);
                    if lo >= hi {
                        break;
                    }
                    let (l, lr) = lrest.split_at_mut(hi - lo);
                    let (e, er) = erest.split_at_mut(hi - lo);
                    let (d, dr) = drest.split_at_mut((hi - lo) * p);
                    lrest = lr;
                    erest = er;
                    drest = dr;
                    let work = &work;
                    scope.spawn(move || work(lo..hi, l, e, d));
                }
            });
        }
        // |ψ|², normalised, in a fixed order.
        let top = logs.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        let raw: Vec<f64> = logs.iter().map(|l| exp(2.0 * (l - top))).collect();
        let z: f64 = raw.iter().sum();
        Batch { weights: raw.iter().map(|r| r / z).collect(), energies, derivs, acceptance: 1.0 }
    }

    fn sample(&mut self) -> Batch {
        let n = self.model.n;
        let p = self.rbm.params();
        let nc = self.chains.len();
        let per_chain = self.cfg.samples.div_ceil(nc).max(1);
        let exchange = self.cfg.magnetisation.is_some();
        let thin = self.cfg.thin.max(1) * n;
        let rbm = &self.rbm;
        let model = &self.model;
        let run = |chain: &mut (Vec<i8>, Rng)| -> (Vec<f64>, Vec<f64>, f64) {
            let (s, rng) = chain;
            let mut energies = Vec::with_capacity(per_chain);
            let mut derivs = vec![0.0; per_chain * p];
            let mut accepted = 0usize;
            for k in 0..per_chain {
                accepted += metropolis(rbm, s, rng, thin, exchange);
                let theta = rbm.theta(s);
                energies.push(local_energy(model, rbm, s, &theta));
                rbm.log_derivatives(s, &theta, &mut derivs[k * p..(k + 1) * p]);
            }
            (energies, derivs, accepted as f64 / (per_chain * thin) as f64)
        };
        let threads = self.cfg.threads.max(1).min(nc);
        let results: Vec<(Vec<f64>, Vec<f64>, f64)> = if threads <= 1 || cfg!(target_arch = "wasm32") {
            self.chains.iter_mut().map(run).collect()
        } else {
            let per = nc.div_ceil(threads);
            std::thread::scope(|scope| {
                let handles: Vec<_> = self
                    .chains
                    .chunks_mut(per)
                    .map(|group| {
                        let run = &run;
                        scope.spawn(move || group.iter_mut().map(run).collect::<Vec<_>>())
                    })
                    .collect();
                handles.into_iter().flat_map(|h| h.join().expect("a sampling thread panicked")).collect()
            })
        };
        let total = per_chain * nc;
        let mut energies = Vec::with_capacity(total);
        let mut derivs = Vec::with_capacity(total * p);
        let mut acc = 0.0;
        for (e, d, a) in results {
            energies.extend(e);
            derivs.extend(d);
            acc += a;
        }
        Batch { weights: vec![1.0 / total as f64; total], energies, derivs, acceptance: acc / nc as f64 }
    }

    /// Estimate the energy without changing the machine.
    pub fn estimate(&mut self) -> Estimate {
        let b = self.batch();
        summarise(&b, self.cfg.exact)
    }

    /// One step of stochastic reconfiguration: `θ ← θ − η (S + λ)⁻¹ F`, with
    /// `S` the covariance of the log-derivatives and `F` their covariance
    /// with the local energy. It is solved in parameter space or in sample
    /// space, whichever is smaller; the two give the same step.
    pub fn step(&mut self) -> Estimate {
        let b = self.batch();
        let est = summarise(&b, self.cfg.exact);
        let delta = sr_direction(&b, self.rbm.params(), self.cfg.shift);
        self.rbm.shift(&delta, self.cfg.lr);
        est
    }
}

fn summarise(b: &Batch, exact: bool) -> Estimate {
    let energy: f64 = b.weights.iter().zip(&b.energies).map(|(w, e)| w * e).sum();
    let variance: f64 = b.weights.iter().zip(&b.energies).map(|(w, e)| w * (e - energy) * (e - energy)).sum();
    let error = if exact { 0.0 } else { (variance / b.energies.len() as f64).sqrt() };
    Estimate { energy, variance, error, acceptance: b.acceptance }
}

/// `(S + λ)⁻¹ F` from a batch.
fn sr_direction(b: &Batch, p: usize, shift: f64) -> Vec<f64> {
    let k = b.weights.len();
    let energy: f64 = b.weights.iter().zip(&b.energies).map(|(w, e)| w * e).sum();
    let mut mean = vec![0.0; p];
    for (r, &w) in b.weights.iter().enumerate() {
        for (m, &d) in mean.iter_mut().zip(&b.derivs[r * p..(r + 1) * p]) {
            *m += w * d;
        }
    }
    // Rows √w (O − Ō) and √w (E − Ē).
    let mut o = vec![0.0; k * p];
    let mut eps = vec![0.0; k];
    for r in 0..k {
        let sw = b.weights[r].sqrt();
        for c in 0..p {
            o[r * p + c] = sw * (b.derivs[r * p + c] - mean[c]);
        }
        eps[r] = sw * (b.energies[r] - energy);
    }
    if p <= k { sr_parameter_space(&o, &eps, k, p, shift) } else { sr_sample_space(&o, eps, k, p, shift) }
}

/// `(OᵀO + λ)⁻¹ Oᵀε`.
fn sr_parameter_space(o: &[f64], eps: &[f64], k: usize, p: usize, shift: f64) -> Vec<f64> {
    {
        // (OᵀO + λ) δ = Oᵀε.
        let mut s = vec![0.0; p * p];
        for r in 0..k {
            let row = &o[r * p..(r + 1) * p];
            for i in 0..p {
                let oi = row[i];
                if oi == 0.0 {
                    continue;
                }
                for j in i..p {
                    s[i * p + j] += oi * row[j];
                }
            }
        }
        for i in 0..p {
            s[i * p + i] += shift;
            for j in 0..i {
                s[i * p + j] = s[j * p + i];
            }
        }
        let mut f = vec![0.0; p];
        for r in 0..k {
            for c in 0..p {
                f[c] += o[r * p + c] * eps[r];
            }
        }
        cholesky_solve(&mut s, &mut f, p);
        f
    }
}

/// `Oᵀ (OOᵀ + λ)⁻¹ ε`.
fn sr_sample_space(o: &[f64], eps: Vec<f64>, k: usize, p: usize, shift: f64) -> Vec<f64> {
    {
        // δ = Oᵀ (OOᵀ + λ)⁻¹ ε.
        let mut t = vec![0.0; k * k];
        for i in 0..k {
            for j in i..k {
                let mut acc = 0.0;
                for c in 0..p {
                    acc += o[i * p + c] * o[j * p + c];
                }
                t[i * k + j] = acc;
                t[j * k + i] = acc;
            }
            t[i * k + i] += shift;
        }
        let mut x = eps;
        cholesky_solve(&mut t, &mut x, k);
        let mut delta = vec![0.0; p];
        for r in 0..k {
            for c in 0..p {
                delta[c] += o[r * p + c] * x[r];
            }
        }
        delta
    }
}

/// Solve `A x = b` for symmetric positive-definite `A` (overwritten by its
/// Cholesky factor), `b` overwritten by `x`.
#[allow(clippy::needless_range_loop)]
fn cholesky_solve(a: &mut [f64], b: &mut [f64], n: usize) {
    for j in 0..n {
        let mut d = a[j * n + j];
        for k in 0..j {
            d -= a[j * n + k] * a[j * n + k];
        }
        let d = d.max(1e-300).sqrt();
        a[j * n + j] = d;
        for i in j + 1..n {
            let mut s = a[i * n + j];
            for k in 0..j {
                s -= a[i * n + k] * a[j * n + k];
            }
            a[i * n + j] = s / d;
        }
    }
    for i in 0..n {
        let mut s = b[i];
        for k in 0..i {
            s -= a[i * n + k] * b[k];
        }
        b[i] = s / a[i * n + i];
    }
    for i in (0..n).rev() {
        let mut s = b[i];
        for k in i + 1..n {
            s -= a[k * n + i] * b[k];
        }
        b[i] = s / a[i * n + i];
    }
}

/// A random configuration, at the given magnetisation if any.
fn start_configuration(n: usize, magnetisation: Option<i32>, rng: &mut Rng) -> Vec<i8> {
    match magnetisation {
        None => (0..n).map(|_| if rng.next() & 1 == 0 { 1 } else { -1 }).collect(),
        Some(m) => {
            let ups = ((n as i32 + m) / 2).clamp(0, n as i32) as usize;
            let mut s: Vec<i8> = (0..n).map(|i| if i < ups { 1 } else { -1 }).collect();
            // Fisher–Yates.
            for i in (1..n).rev() {
                let j = rng.below(i + 1);
                s.swap(i, j);
            }
            s
        }
    }
}

/// `moves` Metropolis moves on `s` (single flips, or swaps of antiparallel
/// spins), accepted with probability `|ψ(s′)/ψ(s)|²`. Returns the number
/// accepted.
fn metropolis(rbm: &Rbm, s: &mut [i8], rng: &mut Rng, moves: usize, exchange: bool) -> usize {
    let n = s.len();
    let mut theta = rbm.theta(s);
    let mut accepted = 0;
    for _ in 0..moves {
        let flips: Vec<usize> = if exchange {
            let i = rng.below(n);
            let j = rng.below(n);
            if s[i] == s[j] {
                continue;
            }
            vec![i, j]
        } else {
            vec![rng.below(n)]
        };
        let r = rbm.ratio(s, &theta, &flips);
        if rng.uniform() < r * r {
            for &i in &flips {
                let ds = -2.0 * f64::from(s[i]);
                for (j, t) in theta.iter_mut().enumerate() {
                    *t += rbm.w[j * n + i] * ds;
                }
                s[i] = -s[i];
            }
            accepted += 1;
        }
    }
    accepted
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantum_dmrg::{Chain, exact_ground_energy};

    #[test]
    fn exact_diagonalisation_matches_the_chain_referees() {
        let n = 10;
        let e = chain_edges(n, false);
        let a = SpinModel::heisenberg(n, &e, 1.0).exact_ground_energy(Some(0)).unwrap();
        let b = exact_ground_energy(&Chain::heisenberg(n, 1.0, 1.0, 0.0)).unwrap();
        assert!((a - b).abs() < 1e-9, "{a} vs {b}");
        let a = SpinModel::ising(n, &e, 1.0, 0.7).exact_ground_energy(None).unwrap();
        let b = exact_ground_energy(&Chain::ising(n, 1.0, 0.7)).unwrap();
        assert!((a - b).abs() < 1e-9, "{a} vs {b}");
    }

    #[test]
    fn the_operator_sum_is_the_model() {
        // A 3 × 3 lattice with every kind of term.
        let n = 9;
        let mut m = SpinModel::heisenberg(n, &square_edges(3, 3, false), 0.8);
        m.x.push((4, 0.3));
        m.z.push((2, -0.45));
        m.zz.push((0, 8, 0.2));
        let dense = m.opsum().mpo().to_dense().unwrap();
        let dim = 1usize << n;
        let mut col = vec![0.0; dim];
        let mut out = vec![0.0; dim];
        for c in 0..dim {
            col.iter_mut().for_each(|x| *x = 0.0);
            col[c] = 1.0;
            m.apply_dense(&col, &mut out);
            for r in 0..dim {
                assert!((out[r] - dense[r * dim + c]).abs() < 1e-12, "({r}, {c})");
            }
        }
    }

    #[test]
    fn ratios_and_derivatives_are_the_machines() {
        let n = 7;
        let mut rbm = Rbm::new(n, 2, 0.4, 3).with_marshall_sign(checkerboard(n, 1));
        let mut rng = Rng(5);
        for v in rbm.a.iter_mut().chain(rbm.b.iter_mut()) {
            *v = 0.3 * (rng.uniform() - 0.5);
        }
        for _ in 0..20 {
            let s = start_configuration(n, None, &mut rng);
            let theta = rbm.theta(&s);
            for flips in [vec![rng.below(n)], vec![0, 3], vec![2, 6]] {
                let mut t = s.clone();
                for &i in &flips {
                    t[i] = -t[i];
                }
                let want = rbm.sign(&t) * rbm.sign(&s) * exp(rbm.log_amplitude(&t) - rbm.log_amplitude(&s));
                let got = rbm.ratio(&s, &theta, &flips);
                assert!((got - want).abs() < 1e-12 * want.abs().max(1.0), "{got} vs {want}");
            }
            // Derivatives by central differences.
            let p = rbm.params();
            let mut d = vec![0.0; p];
            rbm.log_derivatives(&s, &theta, &mut d);
            let h = 1e-6;
            for k in [0, n, n + 3, n + rbm.m + 5, p - 1] {
                let mut e = vec![0.0; p];
                e[k] = -h;
                let mut up = rbm.clone();
                up.shift(&e, 1.0);
                e[k] = h;
                let mut dn = rbm.clone();
                dn.shift(&e, 1.0);
                let fd = (up.log_amplitude(&s) - dn.log_amplitude(&s)) / (2.0 * h);
                assert!((fd - d[k]).abs() < 1e-7, "param {k}: {fd} vs {}", d[k]);
            }
        }
    }

    #[test]
    fn both_spaces_take_the_same_step() {
        let mut rng = Rng(9);
        for (k, p) in [(12usize, 5usize), (5, 12)] {
            let o: Vec<f64> = (0..k * p).map(|_| rng.uniform() - 0.5).collect();
            let eps: Vec<f64> = (0..k).map(|_| rng.uniform() - 0.5).collect();
            let a = sr_parameter_space(&o, &eps, k, p, 1e-3);
            let b = sr_sample_space(&o, eps.clone(), k, p, 1e-3);
            for (x, y) in a.iter().zip(&b) {
                assert!((x - y).abs() < 1e-9 * x.abs().max(1.0), "{x} vs {y}");
            }
        }
    }

    #[test]
    fn exact_vmc_finds_the_ground_state() {
        let n = 8;
        let model = SpinModel::heisenberg(n, &chain_edges(n, false), 1.0);
        let exact = model.exact_ground_energy(Some(0)).unwrap();
        let rbm = Rbm::new(n, 2, 0.05, 1).with_marshall_sign(checkerboard(n, 1));
        let cfg = VmcConfig { exact: true, magnetisation: Some(0), lr: 0.1, shift: 1e-4, ..VmcConfig::default() };
        let mut vmc = Vmc::new(model, rbm, cfg);
        let mut last = vmc.step();
        for _ in 0..300 {
            let e = vmc.step();
            last = e;
        }
        let rel = (last.energy - exact) / exact.abs();
        assert!(last.energy >= exact - 1e-9, "variational: {} vs {exact}", last.energy);
        assert!(rel < 1e-3, "{} vs {exact} (relative {rel:e})", last.energy);
        assert!(last.variance < 1e-2, "{}", last.variance);
    }

    #[test]
    fn sampling_agrees_with_enumeration_and_repeats() {
        let n = 8;
        let model = SpinModel::ising(n, &chain_edges(n, false), 1.0, 1.0);
        let rbm = Rbm::new(n, 1, 0.3, 4);
        let exact = Vmc::new(model.clone(), rbm.clone(), VmcConfig { exact: true, ..VmcConfig::default() }).estimate();
        let cfg = VmcConfig { samples: 8000, chains: 8, threads: 1, ..VmcConfig::default() };
        let sampled = Vmc::new(model.clone(), rbm.clone(), cfg).estimate();
        assert!((sampled.energy - exact.energy).abs() < 5.0 * sampled.error, "{sampled:?} vs {exact:?}");
        // The same bits at any thread count, and again.
        let mut one = Vmc::new(model.clone(), rbm.clone(), VmcConfig { samples: 400, threads: 1, ..VmcConfig::default() });
        let mut many = Vmc::new(model, rbm, VmcConfig { samples: 400, threads: 3, ..VmcConfig::default() });
        for _ in 0..3 {
            assert_eq!(one.step(), many.step());
        }
        assert_eq!(one.rbm, many.rbm);
        assert_eq!(one.chains, many.chains);
    }
}