wai-quantum 0.3.37

//! 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, the frustrated
//!   J₁–J₂ model, and transverse-field Ising, on chains and square lattices.
//! - **The network.** [`Rbm`] is a restricted Boltzmann machine with `α·n`
//!   hidden units, its parameters real or complex ([`Amplitude`]).
//!   - A real network is positive. On a bipartite lattice it can carry
//!     Marshall's sign, which makes the Heisenberg ground state positive in
//!     what remains.
//!   - A complex network carries a phase, so training can find a sign
//!     structure nobody supplied. Frustrated models need that, because no
//!     sign rule is known for them.
//! - **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.
//! - **Complex networks.** A complex network with real parameters is the real
//!   one, to `10⁻¹³`. Its ratios equal ratios of its amplitudes. Its
//!   holomorphic log-derivatives agree with differences along a real step.
//!   Both spaces take the same Hermitian reconfiguration step.
//! - **A sign nobody supplied.** Exact diagonalisation gives the
//!   Majumdar–Ghosh ring (`J₂ = J₁/2`) its exact `−3n/8`. On 8-site rings
//!   with no sign rule given, Heisenberg and Majumdar–Ghosh, a real network
//!   stays more than 10 % off. A complex one comes within `2·10⁻³`, with
//!   every configuration enumerated. On 12-site rings, 400 steps reach:
//!
//!   | ring (no sign given)  | relative error |
//!   |-----------------------|----------------|
//!   | Majumdar–Ghosh        | 2·10⁻⁶         |
//!   | J₁–J₂, `J₂ = 0.3`     | 3.5·10⁻⁴       |
//!   | Heisenberg            | 2·10⁻³         |
//! - **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.
//!
//! # Measured
//!
//! On 10-site rings every model above settled about 9 % high, on a state
//! with essentially zero variance: an eigenstate, but not the ground state.
//! The cause is a symmetry the start imposes.
//!
//! - With its biases at zero, the network is exactly even under flipping
//!   every spin. Training keeps it even, because the biases' gradients vanish
//!   on an even state.
//! - A singlet is even or odd under that flip as half the ring's length is
//!   even or odd. On 10-site rings the singlet ground state is therefore out
//!   of reach, and training finds the lowest even state instead.
//! - Starting the biases away from zero frees it. The same 10-site runs then
//!   reach `3·10⁻⁴` (Heisenberg), `1·10⁻⁶` (Majumdar–Ghosh) and `1.4·10⁻⁵`
//!   (J₁–J₂ 0.3) in 400 steps.
//! - On 12-site rings, where the ground state is even, the even start is the
//!   better one (the table above). Drawn biases there reach only `2·10⁻²`,
//!   `6·10⁻⁷` and `3·10⁻³`.
//!
//! The variance cannot catch the trap, since an excited eigenstate has none.
//! Only a referee can.
//!
//! The same trap caught the referee once. Lanczos stopped a restart once its
//! lowest Ritz pair had converged, which is safe from DMRG's warm starts.
//! From a cold start that pair can be an excited state, and the exact
//! diagonalisation of the 8-site Majumdar–Ghosh ring returned `−2.595`
//! instead of `−3`. The network's energy, a true upper bound, came out below
//! it. The referees now build every Krylov vector.
//!
//! # Honest boundaries
//!
//! - **The hardest point is not reached.** At `J₂ = 0.5` on the periodic
//!   4 × 4 lattice, exact diagonalisation gives `−8.45792335`. At `α = 1`,
//!   600 steps of 1000 samples leave all three networks about 3 % above it:
//!   complex with Marshall's sign, complex without, and real with it. At
//!   this size and budget the network is the limit, not the sampling.
//! - **Symmetry is not chosen.** A network is not projected onto a momentum
//!   or spin-flip sector. Started with zero biases it stays even under a
//!   global spin flip, and that settles the sector, as above.
//! - **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::{C, Op, OpSum, Scalar, default_threads, lanczos};
use crate::repro::{atan2, exp, ln, sin_cos};

// ---------------------------------------------------------------------------
// 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 next-nearest-neighbour bonds of an open (or periodic) chain.
pub fn chain_next_nearest(n: usize, periodic: bool) -> Vec<(usize, usize)> {
    let mut e: Vec<(usize, usize)> = (0..n.saturating_sub(2)).map(|i| (i, i + 2)).collect();
    if periodic && n > 4 {
        e.push((n - 2, 0));
        e.push((n - 1, 1));
    }
    e
}

/// The diagonal (next-nearest) bonds of an `lx × ly` square lattice, open or
/// periodic in both directions.
pub fn square_diagonals(lx: usize, ly: usize, periodic: bool) -> Vec<(usize, usize)> {
    let mut e = Vec::new();
    let wrap = |v: isize, l: usize| -> Option<usize> {
        if v >= 0 && (v as usize) < l {
            Some(v as usize)
        } else if periodic && l > 2 {
            Some(v.rem_euclid(l as isize) as usize)
        } else {
            None
        }
    };
    for y in 0..ly {
        for x in 0..lx {
            let i = y * lx + x;
            for dx in [1isize, -1] {
                if let (Some(xx), Some(yy)) = (wrap(x as isize + dx, lx), wrap(y as isize + 1, ly)) {
                    let j = yy * lx + xx;
                    if i != j && !e.contains(&(i, j)) && !e.contains(&(j, i)) {
                        e.push((i, j));
                    }
                }
            }
        }
    }
    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 frustrated `J₁ Σ_nn Sᵢ·Sⱼ + J₂ Σ_nnn Sᵢ·Sⱼ`. With `J₂ > 0` no
    /// sign rule is known in general, and a real network cannot follow it.
    pub fn j1j2(n: usize, nn: &[(usize, usize)], nnn: &[(usize, usize)], j1: f64, j2: f64) -> SpinModel {
        let mut m = SpinModel::heisenberg(n, nn, j1);
        let second = SpinModel::heisenberg(n, nnn, j2);
        m.exchange.extend(second.exchange);
        m.zz.extend(second.zz);
        m
    }

    /// 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, false).0)
    }
}

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

/// The numbers a network's parameters are: real, or complex so the network
/// can carry a sign structure nobody supplied.
pub trait Amplitude: Scalar {
    fn exp(self) -> Self;
    /// `ln cosh`, stable for large real parts (shifted by `ln 2`).
    fn lncosh(self) -> Self;
    fn tanh(self) -> Self;
    /// `self / d` for a real `d`.
    fn div_re(self, d: f64) -> Self;
    /// An initial weight, uniform in `±scale` (each part, when complex).
    fn draw(uniform: &mut dyn FnMut() -> f64, scale: f64) -> Self;
}

impl Amplitude for f64 {
    fn exp(self) -> f64 {
        exp(self)
    }
    fn lncosh(self) -> f64 {
        let a = self.abs();
        a + ln(1.0 + exp(-2.0 * a)) - core::f64::consts::LN_2
    }
    fn tanh(self) -> f64 {
        let e = exp(-2.0 * self.abs());
        let t = (1.0 - e) / (1.0 + e);
        if self < 0.0 { -t } else { t }
    }
    fn div_re(self, d: f64) -> f64 {
        self / d
    }
    fn draw(uniform: &mut dyn FnMut() -> f64, scale: f64) -> f64 {
        scale * (2.0 * uniform() - 1.0)
    }
}

fn c_ln(z: C) -> C {
    C { re: ln((z.re * z.re + z.im * z.im).sqrt()), im: atan2(z.im, z.re) }
}

fn c_div(a: C, b: C) -> C {
    let d = b.re * b.re + b.im * b.im;
    C { re: (a.re * b.re + a.im * b.im) / d, im: (a.im * b.re - a.re * b.im) / d }
}

impl Amplitude for C {
    fn exp(self) -> C {
        let m = exp(self.re);
        let (s, c) = sin_cos(self.im);
        C { re: m * c, im: m * s }
    }
    fn lncosh(self) -> C {
        // cosh is even; with Re z ≥ 0, ln cosh z = z + ln(1 + e^{−2z}) − ln 2.
        if self.re < 0.0 {
            return Scalar::scale(self, -1.0).lncosh();
        }
        let e = Scalar::scale(self, -2.0).exp();
        let l = c_ln(C { re: 1.0 + e.re, im: e.im });
        C { re: self.re + l.re - core::f64::consts::LN_2, im: self.im + l.im }
    }
    fn tanh(self) -> C {
        if self.re < 0.0 {
            return Scalar::scale(Scalar::scale(self, -1.0).tanh(), -1.0);
        }
        let e = Scalar::scale(self, -2.0).exp();
        c_div(C { re: 1.0 - e.re, im: -e.im }, C { re: 1.0 + e.re, im: e.im })
    }
    fn div_re(self, d: f64) -> C {
        C { re: self.re / d, im: self.im / d }
    }
    fn draw(uniform: &mut dyn FnMut() -> f64, scale: f64) -> C {
        let re = scale * (2.0 * uniform() - 1.0);
        let im = scale * (2.0 * uniform() - 1.0);
        C { re, im }
    }
}

/// A restricted Boltzmann machine: `ψ(s) = sign(s) · exp(Σ aᵢsᵢ) ·
/// Π_j cosh(b_j + Σ W_jᵢ sᵢ)`, with `m = α·n` hidden units. Its parameters are
/// real (`Rbm<f64>`), so `ψ` is positive up to the optional sign, or complex
/// (`Rbm<C>`), so `ψ` carries a phase the training finds. 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<T = f64> {
    pub n: usize,
    pub m: usize,
    pub a: Vec<T>,
    pub b: Vec<T>,
    /// `W[j·n + i]`.
    pub w: Vec<T>,
    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<f64> {
    /// A real machine: `α·n` hidden units, weights uniform in `±scale`,
    /// biases zero.
    pub fn new(n: usize, alpha: usize, scale: f64, seed: u64) -> Rbm {
        Rbm::random(n, alpha, scale, seed)
    }
}

impl Rbm<C> {
    /// A complex machine: weights with real and imaginary parts uniform in
    /// `±scale`, biases zero.
    pub fn complex(n: usize, alpha: usize, scale: f64, seed: u64) -> Rbm<C> {
        Rbm::random(n, alpha, scale, seed)
    }
}

impl<T: Amplitude> Rbm<T> {
    fn random(n: usize, alpha: usize, scale: f64, seed: u64) -> Rbm<T> {
        let m = alpha.max(1) * n;
        let mut rng = Rng(seed);
        let mut u = || rng.uniform();
        let w = (0..m * n).map(|_| T::draw(&mut u, scale)).collect();
        Rbm { n, m, a: vec![T::ZERO; n], b: vec![T::ZERO; m], w, sublattice: None }
    }

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

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

    fn theta(&self, s: &[i8]) -> Vec<T> {
        (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 = t.add(wi.scale(f64::from(si)));
                }
                t
            })
            .collect()
    }

    /// `ln ψ(s)` up to a constant, without the sign: real part `ln |ψ|`,
    /// imaginary part the phase.
    pub fn log_psi(&self, s: &[i8]) -> T {
        let mut l = T::ZERO;
        for (a, &si) in self.a.iter().zip(s) {
            l = l.add(a.scale(f64::from(si)));
        }
        for t in self.theta(s) {
            l = l.add(t.lncosh());
        }
        l
    }

    /// `ln |ψ(s)|` up to a constant.
    pub fn log_amplitude(&self, s: &[i8]) -> f64 {
        self.log_psi(s).re()
    }

    /// The sign of `ψ(s)` from Marshall's rule (one without it).
    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: &[T], flips: &[usize]) -> T {
        let mut l = T::ZERO;
        for &i in flips {
            l = l.sub(self.a[i].scale(2.0).scale(f64::from(s[i])));
        }
        for (j, &t) in theta.iter().enumerate() {
            let mut d = T::ZERO;
            for &i in flips {
                d = d.sub(self.w[j * self.n + i].scale(2.0).scale(f64::from(s[i])));
            }
            l = l.add(t.add(d).lncosh().sub(t.lncosh()));
        }
        let mut r = l.exp();
        if let Some(sub) = &self.sublattice
            && flips.iter().filter(|&&i| sub[i]).count() % 2 == 1
        {
            r = r.scale(-1.0);
        }
        r
    }

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

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

/// `E_loc(s) = Σ_{s′} H_{ss′} ψ(s′)/ψ(s)`.
fn local_energy<T: Amplitude>(model: &SpinModel, rbm: &Rbm<T>, s: &[i8], theta: &[T]) -> T {
    let mut e = T::ZERO;
    for &(i, j, c) in &model.zz {
        e = e.add(T::ONE.scale(c * f64::from(s[i]) * f64::from(s[j])));
    }
    for &(i, c) in &model.z {
        e = e.add(T::ONE.scale(c * f64::from(s[i])));
    }
    for &(i, j, c) in &model.exchange {
        if s[i] != s[j] {
            e = e.add(rbm.ratio(s, theta, &[i, j]).scale(2.0 * c));
        }
    }
    for &(i, c) in &model.x {
        e = e.add(rbm.ratio(s, theta, &[i]).scale(c));
    }
    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 {
    /// `Re ⟨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, real
/// or complex.
#[derive(Clone, Debug, PartialEq)]
pub struct Vmc<T = f64> {
    pub model: SpinModel,
    pub rbm: Rbm<T>,
    pub cfg: VmcConfig,
    chains: Vec<(Vec<i8>, Rng)>,
}

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

impl<T: Amplitude> Vmc<T> {
    pub fn new(model: SpinModel, rbm: Rbm<T>, cfg: VmcConfig) -> Vmc<T> {
        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<T> {
        if self.cfg.exact { self.enumerate() } else { self.sample() }
    }

    fn enumerate(&self) -> Batch<T> {
        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![T::ZERO; k];
        let mut derivs = vec![T::ZERO; 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 [T], derivs: &mut [T]| {
            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 [T] = &mut energies;
                let mut drest: &mut [T] = &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<T> {
        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<T>, Vec<T>, f64) {
            let (s, rng) = chain;
            let mut energies = Vec::with_capacity(per_chain);
            let mut derivs = vec![T::ZERO; 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<T>, Vec<T>, 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 (Hermitian ones, for a complex machine). 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<T: Amplitude>(b: &Batch<T>, exact: bool) -> Estimate {
    let energy: f64 = b.weights.iter().zip(&b.energies).map(|(w, e)| w * e.re()).sum();
    let et = T::ONE.scale(energy);
    let variance: f64 = b
        .weights
        .iter()
        .zip(&b.energies)
        .map(|(&w, e)| {
            let d = e.sub(et);
            d.scale(w).mul(d.conj()).re()
        })
        .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<T: Amplitude>(b: &Batch<T>, p: usize, shift: f64) -> Vec<T> {
    let k = b.weights.len();
    let energy: f64 = b.weights.iter().zip(&b.energies).map(|(w, e)| w * e.re()).sum();
    let et = T::ONE.scale(energy);
    let mut mean = vec![T::ZERO; 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 = m.add(d.scale(w));
        }
    }
    // Rows √w (O − Ō) and √w (E − Ē).
    let mut o = vec![T::ZERO; k * p];
    let mut eps = vec![T::ZERO; k];
    for r in 0..k {
        let sw = b.weights[r].sqrt();
        for c in 0..p {
            o[r * p + c] = b.derivs[r * p + c].sub(mean[c]).scale(sw);
        }
        eps[r] = b.energies[r].sub(et).scale(sw);
    }
    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<T: Amplitude>(o: &[T], eps: &[T], k: usize, p: usize, shift: f64) -> Vec<T> {
    let mut s = vec![T::ZERO; p * p];
    for r in 0..k {
        let row = &o[r * p..(r + 1) * p];
        for i in 0..p {
            let oi = row[i].conj();
            if oi.is_zero() {
                continue;
            }
            for j in i..p {
                s[i * p + j] = s[i * p + j].add(oi.mul(row[j]));
            }
        }
    }
    for i in 0..p {
        s[i * p + i] = s[i * p + i].add(T::ONE.scale(shift));
        for j in 0..i {
            s[i * p + j] = s[j * p + i].conj();
        }
    }
    let mut f = vec![T::ZERO; p];
    for r in 0..k {
        for c in 0..p {
            f[c] = f[c].add(o[r * p + c].conj().mul(eps[r]));
        }
    }
    cholesky_solve(&mut s, &mut f, p);
    f
}

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

/// Solve `A x = b` for Hermitian positive-definite `A` (overwritten by its
/// Cholesky factor `L`, `A = L L†`), `b` overwritten by `x`.
#[allow(clippy::needless_range_loop)]
fn cholesky_solve<T: Amplitude>(a: &mut [T], b: &mut [T], n: usize) {
    for j in 0..n {
        let mut d = a[j * n + j].re();
        for k in 0..j {
            d -= a[j * n + k].norm2();
        }
        let d = d.max(1e-300).sqrt();
        a[j * n + j] = T::ONE.scale(d);
        for i in j + 1..n {
            let mut s = a[i * n + j];
            for k in 0..j {
                s = s.sub(a[i * n + k].mul(a[j * n + k].conj()));
            }
            a[i * n + j] = s.div_re(d);
        }
    }
    for i in 0..n {
        let mut s = b[i];
        for k in 0..i {
            s = s.sub(a[i * n + k].mul(b[k]));
        }
        b[i] = s.div_re(a[i * n + i].re());
    }
    for i in (0..n).rev() {
        let mut s = b[i];
        for k in i + 1..n {
            s = s.sub(a[k * n + i].conj().mul(b[k]));
        }
        b[i] = s.div_re(a[i * n + i].re());
    }
}

/// 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<T: Amplitude>(rbm: &Rbm<T>, 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.norm2() {
            for &i in &flips {
                let ds = -2.0 * f64::from(s[i]);
                for (j, t) in theta.iter_mut().enumerate() {
                    *t = t.add(rbm.w[j * n + i].scale(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);
    }

    #[test]
    fn a_complex_machine_with_real_parameters_is_the_real_one() {
        let n = 6;
        let real = Rbm::new(n, 2, 0.4, 5);
        let cplx = Rbm { n: real.n, m: real.m, a: real.a.iter().map(|&re| C { re, im: 0.0 }).collect(), b: real.b.iter().map(|&re| C { re, im: 0.0 }).collect(), w: real.w.iter().map(|&re| C { re, im: 0.0 }).collect(), sublattice: None };
        let mut rng = Rng(8);
        for _ in 0..10 {
            let s = start_configuration(n, None, &mut rng);
            let (tr, tc) = (real.theta(&s), cplx.theta(&s));
            assert!((real.log_amplitude(&s) - cplx.log_amplitude(&s)).abs() < 1e-14);
            assert!(cplx.log_psi(&s).im.abs() < 1e-14);
            for flips in [vec![1], vec![0, 4]] {
                let (r, c) = (real.ratio(&s, &tr, &flips), cplx.ratio(&s, &tc, &flips));
                assert!((r - c.re).abs() < 1e-13 * r.abs().max(1.0) && c.im.abs() < 1e-13);
            }
        }
    }

    #[test]
    fn complex_ratios_and_derivatives_are_the_machines() {
        let n = 6;
        let mut rbm = Rbm::complex(n, 2, 0.4, 6);
        let mut rng = Rng(12);
        for v in rbm.a.iter_mut().chain(rbm.b.iter_mut()) {
            *v = C { re: 0.3 * (rng.uniform() - 0.5), im: 0.3 * (rng.uniform() - 0.5) };
        }
        let p = rbm.params();
        for _ in 0..10 {
            let s = start_configuration(n, None, &mut rng);
            let theta = rbm.theta(&s);
            for flips in [vec![rng.below(n)], vec![1, 5]] {
                let mut t = s.clone();
                for &i in &flips {
                    t[i] = -t[i];
                }
                let want = rbm.log_psi(&t).sub(rbm.log_psi(&s)).exp();
                let got = rbm.ratio(&s, &theta, &flips);
                assert!(got.sub(want).norm2().sqrt() < 1e-12 * want.norm2().sqrt().max(1.0), "{got:?} vs {want:?}");
            }
            // Holomorphic: the derivative along a real step is the derivative.
            let mut d = vec![C::ZERO; p];
            rbm.log_derivatives(&s, &theta, &mut d);
            let h = 1e-6;
            for k in [0, n + 1, n + rbm.m + 3, p - 1] {
                let mut e = vec![C::ZERO; p];
                e[k] = C { re: -h, im: 0.0 };
                let mut up = rbm.clone();
                up.shift(&e, 1.0);
                e[k] = C { re: h, im: 0.0 };
                let mut dn = rbm.clone();
                dn.shift(&e, 1.0);
                let fd = up.log_psi(&s).sub(dn.log_psi(&s)).scale(0.5 / h);
                assert!(fd.sub(d[k]).norm2().sqrt() < 1e-7, "param {k}: {fd:?} vs {:?}", d[k]);
            }
        }
    }

    #[test]
    fn both_spaces_take_the_same_complex_step() {
        let mut rng = Rng(19);
        for (k, p) in [(12usize, 5usize), (5, 12)] {
            let o: Vec<C> = (0..k * p).map(|_| C { re: rng.uniform() - 0.5, im: rng.uniform() - 0.5 }).collect();
            let eps: Vec<C> = (0..k).map(|_| C { re: rng.uniform() - 0.5, im: 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.sub(*y).norm2().sqrt() < 1e-9 * x.norm2().sqrt().max(1.0), "{x:?} vs {y:?}");
            }
        }
    }

    #[test]
    fn exact_diagonalisation_solves_majumdar_ghosh() {
        // At J₂ = J₁/2 a ring's ground states are the two dimer coverings,
        // at exactly −3n/8.
        for n in [8usize, 12] {
            let m = SpinModel::j1j2(n, &chain_edges(n, true), &chain_next_nearest(n, true), 1.0, 0.5);
            let e = m.exact_ground_energy(Some(0)).unwrap();
            assert!((e + 3.0 * n as f64 / 8.0).abs() < 1e-9, "n={n}: {e}");
        }
    }

    #[test]
    fn a_complex_machine_learns_a_sign_a_real_one_cannot() {
        // The Heisenberg ring and the Majumdar–Ghosh ring, with no sign rule
        // given. A real network is positive and stays far off; a complex one
        // finds the signs.
        let n = 8;
        let models = [
            SpinModel::heisenberg(n, &chain_edges(n, true), 1.0),
            SpinModel::j1j2(n, &chain_edges(n, true), &chain_next_nearest(n, true), 1.0, 0.5),
        ];
        for model in models {
            let exact = model.exact_ground_energy(Some(0)).unwrap();
            let cfg = VmcConfig { exact: true, magnetisation: Some(0), lr: 0.1, shift: 1e-3, ..VmcConfig::default() };
            let mut real = Vmc::new(model.clone(), Rbm::new(n, 2, 0.05, 1), cfg);
            let mut cplx = Vmc::new(model, Rbm::complex(n, 2, 0.05, 3), cfg);
            let (mut er, mut ec) = (real.step(), cplx.step());
            for _ in 0..250 {
                er = real.step();
                ec = cplx.step();
            }
            assert!((er.energy - exact) / exact.abs() > 0.1, "real: {} vs {exact}", er.energy);
            assert!(ec.energy >= exact - 1e-9, "variational: {} vs {exact}", ec.energy);
            assert!((ec.energy - exact) / exact.abs() < 2e-3, "complex: {} vs {exact}", ec.energy);
        }
    }
}