wai-quantum 0.3.34

//! Ground states of spin chains by DMRG — `wai.quantum.dmrg`.
//!
//! Most claims of a quantum advantage in simulating matter are claims that
//! classical methods cannot reach the answer. For one-dimensional systems the
//! density-matrix renormalization group usually can, with the error
//! controlled by the entanglement it keeps. This module is that method,
//! deterministic and held to exact answers.
//!
//! - **Hamiltonians.** [`Chain`] is a nearest-neighbour spin-½ chain: real
//!   two-site terms on every bond and a one-site term on every site. Presets
//!   are the XXZ Heisenberg chain in a field and the transverse-field Ising
//!   chain. [`Chain::mpo`] writes it as a matrix-product operator, the finite
//!   state machine of its terms.
//! - **Two-site DMRG.** [`dmrg`] sweeps a matrix-product state left and right.
//!   - At each bond it finds the lowest eigenvector of the effective
//!     Hamiltonian by Lanczos, with full re-orthogonalisation.
//!   - It splits that vector by the eigenvectors of its reduced density
//!     matrix. It keeps at most `max_bond` of them and drops those below a
//!     weight cutoff.
//!   - The discarded weight, and the entanglement entropy at every bond, are
//!     reported.
//! - **Exact referees.**
//!   - [`exact_ground_energy`] runs Lanczos on the full `2ⁿ` space for short
//!     chains.
//!   - [`ising_exact_energy`] is the transverse-field Ising chain solved as
//!     free fermions: the ground energy of an open chain of any length, from
//!     the eigenvalues of an `n × n` matrix.
//!
//! The eigensolver is Householder tridiagonalisation followed by the implicit
//! QL iteration, with hypotenuses taken in-crate. The effective Hamiltonian
//! is applied in blocks of rows across threads, each row in the same order.
//! Every sum runs in a fixed order, so a run is the same bits on every
//! machine and at any thread count.
//!
//! # Checked
//!
//! - **The eigensolver.** It reconstructs random symmetric matrices up to
//!   40 × 40 to `10⁻¹²` with orthonormal vectors.
//! - **The free-fermion solution.** It equals exact diagonalisation of the
//!   Ising chain at three couplings, for 4, 7 and 10 sites.
//! - **DMRG against exact diagonalisation.** On Heisenberg chains (isotropic,
//!   and anisotropic in a field) and Ising chains of 10–12 sites, energies
//!   agree to `10⁻⁸`.
//! - **At scale** (`dmrg_chain` example; energy error per site against the
//!   exact answer):
//!
//!   | chain                  | sites | bond | error per site | discarded weight |
//!   |------------------------|-------|------|----------------|------------------|
//!   | Ising, critical        | 100   | ≤ 32 | 1.1·10⁻¹¹      | 2.5·10⁻¹²        |
//!   | Ising, critical        | 200   | ≤ 64 | 1.4·10⁻¹¹      | 3.0·10⁻¹²        |
//!   | Heisenberg             | 20    | 64   | 4.4·10⁻¹²      | 4.4·10⁻¹²        |
//!
//!   The 100-site Heisenberg chain has no exact solution at that length. Its
//!   energy falls from above as the bond grows: −44.127739250 at bond 64
//!   (discarded 4.9·10⁻⁹), then −44.127739889 at bond 128 (2.5·10⁻¹¹).
//! - **Variational and repeatable.** The energy never goes below the exact
//!   value, never rises from one sweep to the next, and repeats bit for bit at
//!   any thread count.
//!
//! # Honest boundaries
//!
//! - **One dimension, nearest neighbours, real Hamiltonians.** Longer-range
//!   couplings, complex terms and two-dimensional lattices are not handled,
//!   and nor is time evolution.
//! - **No symmetry sectors.** Conserved quantities are not used to block the
//!   matrices, so a run costs more than a symmetric code would.
//! - **Variational.** The energy is an upper bound, converging from above as
//!   `max_bond` grows. Critical chains converge slowest.

use crate::repro::ln;

// ---------------------------------------------------------------------------
// The symmetric eigensolver
// ---------------------------------------------------------------------------

fn hypot(a: f64, b: f64) -> f64 {
    let (a, b) = (a.abs(), b.abs());
    let m = a.max(b);
    if m == 0.0 {
        return 0.0;
    }
    let (x, y) = (a / m, b / m);
    m * (x * x + y * y).sqrt()
}

/// Eigenvalues (ascending) and eigenvectors of a real symmetric `n × n`
/// matrix given row-major. Eigenvector `i` is column `i` of the returned
/// row-major matrix.
// The loops keep the index form of the published tred2/tql2 algorithms, so
// the code can be checked against them line by line.
#[allow(clippy::needless_range_loop, clippy::manual_memcpy)]
pub fn symmetric_eigen(a: &[f64], n: usize) -> (Vec<f64>, Vec<f64>) {
    if n == 0 {
        return (Vec::new(), Vec::new());
    }
    let mut v: Vec<Vec<f64>> = (0..n).map(|i| a[i * n..(i + 1) * n].to_vec()).collect();
    let mut d = vec![0.0; n];
    let mut e = vec![0.0; n];
    // Householder reduction to tridiagonal form.
    for j in 0..n {
        d[j] = v[n - 1][j];
    }
    for i in (1..n).rev() {
        let mut scale = 0.0;
        let mut h = 0.0;
        for k in 0..i {
            scale += d[k].abs();
        }
        if scale == 0.0 {
            e[i] = d[i - 1];
            for j in 0..i {
                d[j] = v[i - 1][j];
                v[i][j] = 0.0;
                v[j][i] = 0.0;
            }
        } else {
            for k in 0..i {
                d[k] /= scale;
                h += d[k] * d[k];
            }
            let mut f = d[i - 1];
            let mut g = h.sqrt();
            if f > 0.0 {
                g = -g;
            }
            e[i] = scale * g;
            h -= f * g;
            d[i - 1] = f - g;
            for item in e.iter_mut().take(i) {
                *item = 0.0;
            }
            for j in 0..i {
                f = d[j];
                v[j][i] = f;
                g = e[j] + v[j][j] * f;
                for k in j + 1..i {
                    g += v[k][j] * d[k];
                    e[k] += v[k][j] * f;
                }
                e[j] = g;
            }
            f = 0.0;
            for j in 0..i {
                e[j] /= h;
                f += e[j] * d[j];
            }
            let hh = f / (h + h);
            for j in 0..i {
                e[j] -= hh * d[j];
            }
            for j in 0..i {
                f = d[j];
                g = e[j];
                for k in j..i {
                    v[k][j] -= f * e[k] + g * d[k];
                }
                d[j] = v[i - 1][j];
                v[i][j] = 0.0;
            }
        }
        d[i] = h;
    }
    for i in 0..n - 1 {
        v[n - 1][i] = v[i][i];
        v[i][i] = 1.0;
        let h = d[i + 1];
        if h != 0.0 {
            for k in 0..=i {
                d[k] = v[k][i + 1] / h;
            }
            for j in 0..=i {
                let mut g = 0.0;
                for k in 0..=i {
                    g += v[k][i + 1] * v[k][j];
                }
                for k in 0..=i {
                    v[k][j] -= g * d[k];
                }
            }
        }
        for row in v.iter_mut().take(i + 1) {
            row[i + 1] = 0.0;
        }
    }
    for j in 0..n {
        d[j] = v[n - 1][j];
        v[n - 1][j] = 0.0;
    }
    v[n - 1][n - 1] = 1.0;
    e[0] = 0.0;
    // Implicit QL on the tridiagonal matrix.
    for i in 1..n {
        e[i - 1] = e[i];
    }
    e[n - 1] = 0.0;
    let mut f = 0.0;
    let mut tst1: f64 = 0.0;
    let eps = f64::EPSILON;
    for l in 0..n {
        tst1 = tst1.max(d[l].abs() + e[l].abs());
        let mut m = l;
        while m < n {
            if e[m].abs() <= eps * tst1 {
                break;
            }
            m += 1;
        }
        if m > l {
            loop {
                let mut g = d[l];
                let mut p = (d[l + 1] - g) / (2.0 * e[l]);
                let mut r = hypot(p, 1.0);
                if p < 0.0 {
                    r = -r;
                }
                d[l] = e[l] / (p + r);
                d[l + 1] = e[l] * (p + r);
                let dl1 = d[l + 1];
                let mut h = g - d[l];
                for item in d.iter_mut().skip(l + 2) {
                    *item -= h;
                }
                f += h;
                p = d[m];
                let (mut c, mut c2, mut c3) = (1.0, 1.0, 1.0);
                let el1 = e[l + 1];
                let (mut s, mut s2) = (0.0, 0.0);
                for i in (l..m).rev() {
                    c3 = c2;
                    c2 = c;
                    s2 = s;
                    g = c * e[i];
                    h = c * p;
                    r = hypot(p, e[i]);
                    e[i + 1] = s * r;
                    s = e[i] / r;
                    c = p / r;
                    p = c * d[i] - s * g;
                    d[i + 1] = h + s * (c * g + s * d[i]);
                    for row in v.iter_mut() {
                        h = row[i + 1];
                        row[i + 1] = s * row[i] + c * h;
                        row[i] = c * row[i] - s * h;
                    }
                }
                p = -s * s2 * c3 * el1 * e[l] / dl1;
                e[l] = s * p;
                d[l] = c * p;
                if e[l].abs() <= eps * tst1 {
                    break;
                }
            }
        }
        d[l] += f;
        e[l] = 0.0;
    }
    // Sort ascending, carrying the vectors.
    for i in 0..n - 1 {
        let mut k = i;
        let mut p = d[i];
        for (j, &dj) in d.iter().enumerate().skip(i + 1) {
            if dj < p {
                k = j;
                p = dj;
            }
        }
        if k != i {
            d.swap(k, i);
            for row in v.iter_mut() {
                row.swap(i, k);
            }
        }
    }
    (d, v.into_iter().flatten().collect())
}

// ---------------------------------------------------------------------------
// Hamiltonians
// ---------------------------------------------------------------------------

/// A real one-site operator, row-major in the basis `|↑⟩ = 0`, `|↓⟩ = 1`.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Op(pub [f64; 4]);

impl Op {
    pub fn id() -> Op {
        Op([1.0, 0.0, 0.0, 1.0])
    }
    pub fn x() -> Op {
        Op([0.0, 1.0, 1.0, 0.0])
    }
    pub fn z() -> Op {
        Op([1.0, 0.0, 0.0, -1.0])
    }
    /// `S⁺ = |↑⟩⟨↓|`.
    pub fn plus() -> Op {
        Op([0.0, 1.0, 0.0, 0.0])
    }
    /// `S⁻ = |↓⟩⟨↑|`.
    pub fn minus() -> Op {
        Op([0.0, 0.0, 1.0, 0.0])
    }
    /// `S^z = σ^z / 2`.
    pub fn sz() -> Op {
        Op([0.5, 0.0, 0.0, -0.5])
    }
    fn scaled(self, c: f64) -> Op {
        Op(self.0.map(|v| v * c))
    }
    fn add(self, o: Op) -> Op {
        Op([self.0[0] + o.0[0], self.0[1] + o.0[1], self.0[2] + o.0[2], self.0[3] + o.0[3]])
    }
}

/// A uniform nearest-neighbour chain: `H = Σ_i Σ_t c_t A_t(i) B_t(i+1) +
/// Σ_i Σ_f c_f O_f(i)`, open boundaries.
#[derive(Clone, Debug, PartialEq)]
pub struct Chain {
    pub n: usize,
    pub bonds: Vec<(f64, Op, Op)>,
    pub sites: Vec<(f64, Op)>,
}

impl Chain {
    /// The XXZ Heisenberg chain in a field:
    /// `H = Σ [J (S^x S^x + S^y S^y) + J_z S^z S^z] − h Σ S^z`.
    pub fn heisenberg(n: usize, j: f64, jz: f64, h: f64) -> Chain {
        let mut sites = Vec::new();
        if h != 0.0 {
            sites.push((-h, Op::sz()));
        }
        Chain { n, bonds: vec![(j / 2.0, Op::plus(), Op::minus()), (j / 2.0, Op::minus(), Op::plus()), (jz, Op::sz(), Op::sz())], sites }
    }

    /// The transverse-field Ising chain: `H = −J Σ σ^z σ^z − h Σ σ^x`.
    pub fn ising(n: usize, j: f64, h: f64) -> Chain {
        Chain { n, bonds: vec![(-j, Op::z(), Op::z())], sites: vec![(-h, Op::x())] }
    }

    /// The chain as a matrix-product operator. The finite-state machine is:
    /// state 0 is "nothing placed yet", state `t + 1` is "A_t placed, B_t
    /// owed", and the last state is "done".
    pub fn mpo(&self) -> Mpo {
        let m = self.bonds.len();
        let w = m + 2;
        let onsite = self.sites.iter().fold(Op([0.0; 4]), |acc, &(c, o)| acc.add(o.scaled(c)));
        let mut site = vec![Op([0.0; 4]); w * w];
        site[0] = Op::id();
        site[(w - 1) * w + (w - 1)] = Op::id();
        site[w - 1] = onsite;
        for (t, &(c, a, b)) in self.bonds.iter().enumerate() {
            site[t + 1] = a.scaled(c);
            site[(t + 1) * w + (w - 1)] = b;
        }
        Mpo { n: self.n, w, site }
    }
}

/// A uniform matrix-product operator: `site[a·w + b]` is the one-site
/// operator from state `a` to state `b`; the chain starts in state 0 and
/// ends in state `w − 1`.
#[derive(Clone, Debug, PartialEq)]
pub struct Mpo {
    pub n: usize,
    pub w: usize,
    site: Vec<Op>,
}

// ---------------------------------------------------------------------------
// Exact referees
// ---------------------------------------------------------------------------

/// Lanczos for the lowest eigenpair of a symmetric operator given by
/// `apply`, from `start`, with full re-orthogonalisation.
fn lanczos(apply: &dyn Fn(&[f64], &mut [f64]), start: &[f64], krylov: usize, restarts: usize, tol: f64) -> (f64, Vec<f64>) {
    let dim = start.len();
    let norm = |v: &[f64]| v.iter().map(|x| x * x).sum::<f64>().sqrt();
    let mut v = start.to_vec();
    let nv = norm(&v);
    if nv == 0.0 {
        v = (0..dim).map(|i| 1.0 + (i % 7) as f64 * 0.1).collect();
    }
    let nv = norm(&v);
    v.iter_mut().for_each(|x| *x /= nv);
    let mut energy = 0.0;
    let mut w = vec![0.0; dim];
    for _ in 0..restarts.max(1) {
        let k = krylov.min(dim).max(1);
        let mut basis: Vec<Vec<f64>> = vec![v.clone()];
        let (mut alpha, mut beta) = (Vec::new(), Vec::new());
        for j in 0..k {
            apply(&basis[j], &mut w);
            let a: f64 = w.iter().zip(&basis[j]).map(|(x, y)| x * y).sum();
            alpha.push(a);
            // Full re-orthogonalisation, twice.
            for _ in 0..2 {
                for b in &basis {
                    let c: f64 = w.iter().zip(b).map(|(x, y)| x * y).sum();
                    w.iter_mut().zip(b).for_each(|(x, y)| *x -= c * y);
                }
            }
            let bnorm = norm(&w);
            if j + 1 == k || bnorm < 1e-13 {
                break;
            }
            beta.push(bnorm);
            basis.push(w.iter().map(|x| x / bnorm).collect());
        }
        let m = alpha.len();
        let mut t = vec![0.0; m * m];
        for i in 0..m {
            t[i * m + i] = alpha[i];
            if i + 1 < m {
                t[i * m + i + 1] = beta[i];
                t[(i + 1) * m + i] = beta[i];
            }
        }
        let (vals, vecs) = symmetric_eigen(&t, m);
        energy = vals[0];
        let mut next = vec![0.0; dim];
        for (i, b) in basis.iter().enumerate().take(m) {
            let c = vecs[i * m];
            next.iter_mut().zip(b).for_each(|(x, y)| *x += c * y);
        }
        let nn = norm(&next);
        next.iter_mut().for_each(|x| *x /= nn);
        v = next;
        apply(&v, &mut w);
        let resid: f64 = w.iter().zip(&v).map(|(x, y)| (x - energy * y) * (x - energy * y)).sum::<f64>().sqrt();
        if resid < tol {
            break;
        }
    }
    (energy, v)
}

/// The exact ground energy of `chain` by Lanczos on all `2ⁿ` states (at most
/// 22 sites).
pub fn exact_ground_energy(chain: &Chain) -> Option<f64> {
    let n = chain.n;
    if n == 0 || n > 22 {
        return None;
    }
    let dim = 1usize << n;
    let bit = |x: usize, i: usize| (x >> (n - 1 - i)) & 1;
    let apply = |inp: &[f64], out: &mut [f64]| {
        out.iter_mut().for_each(|o| *o = 0.0);
        for (x, &amp) in inp.iter().enumerate() {
            if amp == 0.0 {
                continue;
            }
            for i in 0..n {
                let s = bit(x, i);
                for &(c, o) in &chain.sites {
                    for sp in 0..2 {
                        let m = o.0[sp * 2 + s];
                        if m != 0.0 {
                            let y = (x & !(1 << (n - 1 - i))) | (sp << (n - 1 - i));
                            out[y] += c * m * amp;
                        }
                    }
                }
                if i + 1 < n {
                    let s2 = bit(x, i + 1);
                    for &(c, a, b) in &chain.bonds {
                        for ap in 0..2 {
                            let ma = a.0[ap * 2 + s];
                            if ma == 0.0 {
                                continue;
                            }
                            for bp in 0..2 {
                                let mb = b.0[bp * 2 + s2];
                                if mb == 0.0 {
                                    continue;
                                }
                                let mask = (1 << (n - 1 - i)) | (1 << (n - 2 - i));
                                let y = (x & !mask) | (ap << (n - 1 - i)) | (bp << (n - 2 - i));
                                out[y] += c * ma * mb * amp;
                            }
                        }
                    }
                }
            }
        }
    };
    // u64, not usize: the start vector must not depend on the pointer width.
    let start: Vec<f64> = (0..dim as u64).map(|x| 1.0 + (x.wrapping_mul(2_654_435_761) % 1000) as f64 / 1000.0).collect();
    Some(lanczos(&apply, &start, 60, 20, 1e-10).0)
}

/// The exact ground energy of the open transverse-field Ising chain
/// `−J Σ σ^z σ^z − h Σ σ^x` as free fermions. After a Jordan–Wigner
/// transformation, the chain is quadratic in fermions with
/// `A_ii = 2h`, `A_{i,i+1} = A_{i+1,i} = −J` and `B_{i,i+1} = −B_{i+1,i} = −J`.
/// The quasiparticle energies `Λ_k` are the square roots of the eigenvalues
/// of `(A − B)(A + B)`, and `E₀ = −½ Σ_k Λ_k`.
pub fn ising_exact_energy(n: usize, j: f64, h: f64) -> f64 {
    let mut a = vec![0.0; n * n];
    let mut b = vec![0.0; n * n];
    for i in 0..n {
        a[i * n + i] = 2.0 * h;
        if i + 1 < n {
            a[i * n + i + 1] = -j;
            a[(i + 1) * n + i] = -j;
            b[i * n + i + 1] = -j;
            b[(i + 1) * n + i] = j;
        }
    }
    // M = (A − B)(A + B), symmetric positive semidefinite.
    let mut m = vec![0.0; n * n];
    for r in 0..n {
        for c in 0..n {
            let mut s = 0.0;
            for k in 0..n {
                s += (a[r * n + k] - b[r * n + k]) * (a[k * n + c] + b[k * n + c]);
            }
            m[r * n + c] = s;
        }
    }
    // Symmetrise against rounding.
    for r in 0..n {
        for c in r + 1..n {
            let s = 0.5 * (m[r * n + c] + m[c * n + r]);
            m[r * n + c] = s;
            m[c * n + r] = s;
        }
    }
    let (vals, _) = symmetric_eigen(&m, n);
    -0.5 * vals.iter().map(|v| v.max(0.0).sqrt()).sum::<f64>()
}

// ---------------------------------------------------------------------------
// Matrix-product states and DMRG
// ---------------------------------------------------------------------------

/// A real matrix-product state: site `i` is `left × 2 × right`, stored
/// `[(a·2 + s)·right + b]`.
#[derive(Clone, Debug, PartialEq)]
pub struct Mps {
    pub dims: Vec<usize>,
    pub sites: Vec<Vec<f64>>,
}

impl Mps {
    /// A deterministic pseudo-random state with bond dimension `bond` (capped
    /// by what the chain allows).
    pub fn random(n: usize, bond: usize, seed: u64) -> Mps {
        let mut s = seed;
        let mut next = move || {
            s = s.wrapping_add(0x9e37_79b9_7f4a_7c15);
            let mut z = s;
            z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
            ((z ^ (z >> 31)) >> 11) as f64 / 9_007_199_254_740_992.0 - 0.5
        };
        let mut dims = vec![1usize; n + 1];
        for (i, d) in dims.iter_mut().enumerate().take(n).skip(1) {
            let cap = 1usize.checked_shl(i.min(n - i).min(30) as u32).unwrap_or(usize::MAX);
            *d = bond.min(cap);
        }
        let sites = (0..n).map(|i| (0..dims[i] * 2 * dims[i + 1]).map(|_| next()).collect()).collect();
        Mps { dims, sites }
    }

    pub fn n(&self) -> usize {
        self.sites.len()
    }
}

/// Settings for [`dmrg`].
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DmrgConfig {
    pub max_bond: usize,
    /// Discard density-matrix weight below this.
    pub cutoff: f64,
    /// Full sweeps (left to right and back).
    pub sweeps: usize,
    /// Lanczos vectors per restart, and restarts.
    pub krylov: usize,
    pub restarts: usize,
    /// Stop Lanczos once the residual `‖Hv − Ev‖` falls below this. The
    /// energy error goes as its square.
    pub tol: f64,
    pub seed: u64,
    /// Threads for the effective Hamiltonian; the result does not depend on
    /// it.
    pub threads: usize,
}

impl Default for DmrgConfig {
    fn default() -> DmrgConfig {
        let threads = if cfg!(target_arch = "wasm32") { 1 } else { std::thread::available_parallelism().map_or(1, |n| n.get()) };
        DmrgConfig { max_bond: 64, cutoff: 1e-12, sweeps: 6, krylov: 24, restarts: 4, tol: 1e-8, seed: 1, threads }
    }
}

/// The outcome of a DMRG run.
#[derive(Clone, Debug, PartialEq)]
pub struct DmrgResult {
    pub energy: f64,
    /// The energy after each sweep.
    pub sweep_energies: Vec<f64>,
    /// The largest discarded weight in the last sweep.
    pub discarded: f64,
    /// Von Neumann entanglement entropy at each of the `n − 1` bonds (last
    /// sweep).
    pub entropies: Vec<f64>,
    pub mps: Mps,
}

/// Contract environments and the two-site effective Hamiltonian.
struct Engine<'a> {
    mpo: &'a Mpo,
    w: usize,
}

impl Engine<'_> {
    fn op(&self, a: usize, b: usize) -> &Op {
        &self.mpo.site[a * self.w + b]
    }

    fn nonzero(&self, a: usize, b: usize) -> bool {
        self.op(a, b).0.iter().any(|&v| v != 0.0)
    }

    /// `L'[b', w2, b] = Σ A[a',s',b'] L[a',w1,a] W[w1,w2,s',s] A[a,s,b]`.
    fn grow_left(&self, l: &[f64], a: &[f64], dl: usize, dr: usize) -> Vec<f64> {
        let w = self.w;
        // T[a', w1, s, b] = Σ_a L[a', w1, a] A[a, s, b]
        let mut t = vec![0.0; dl * w * 2 * dr];
        for ap in 0..dl {
            for w1 in 0..w {
                let lrow = &l[(ap * w + w1) * dl..(ap * w + w1 + 1) * dl];
                for (aa, &lv) in lrow.iter().enumerate() {
                    if lv == 0.0 {
                        continue;
                    }
                    for s in 0..2 {
                        let src = &a[(aa * 2 + s) * dr..(aa * 2 + s + 1) * dr];
                        let dst = &mut t[((ap * w + w1) * 2 + s) * dr..((ap * w + w1) * 2 + s + 1) * dr];
                        for (d, &x) in dst.iter_mut().zip(src) {
                            *d += lv * x;
                        }
                    }
                }
            }
        }
        // U[a', w2, s', b] = Σ_{w1, s} W[w1, w2, s', s] T[a', w1, s, b]
        let mut u = vec![0.0; dl * w * 2 * dr];
        for ap in 0..dl {
            for w1 in 0..w {
                for w2 in 0..w {
                    if !self.nonzero(w1, w2) {
                        continue;
                    }
                    let o = self.op(w1, w2).0;
                    for sp in 0..2 {
                        for s in 0..2 {
                            let m = o[sp * 2 + s];
                            if m == 0.0 {
                                continue;
                            }
                            let src = ((ap * w + w1) * 2 + s) * dr;
                            let dst = ((ap * w + w2) * 2 + sp) * dr;
                            for b in 0..dr {
                                u[dst + b] += m * t[src + b];
                            }
                        }
                    }
                }
            }
        }
        // L'[b', w2, b] = Σ_{a', s'} A[a', s', b'] U[a', w2, s', b]
        let mut out = vec![0.0; dr * w * dr];
        for ap in 0..dl {
            for sp in 0..2 {
                let arow = &a[(ap * 2 + sp) * dr..(ap * 2 + sp + 1) * dr];
                for (bp, &av) in arow.iter().enumerate() {
                    if av == 0.0 {
                        continue;
                    }
                    for w2 in 0..w {
                        let src = &u[((ap * w + w2) * 2 + sp) * dr..((ap * w + w2) * 2 + sp + 1) * dr];
                        let dst = &mut out[(bp * w + w2) * dr..(bp * w + w2 + 1) * dr];
                        for (d, &x) in dst.iter_mut().zip(src) {
                            *d += av * x;
                        }
                    }
                }
            }
        }
        out
    }

    /// `R'[a', w1, a] = Σ B[a',s',b'] W[w1,w2,s',s] R[b',w2,b] B[a,s,b]`.
    fn grow_right(&self, r: &[f64], bt: &[f64], dl: usize, dr: usize) -> Vec<f64> {
        let w = self.w;
        // T[a, s, b', w2] = Σ_b B[a, s, b] R[b', w2, b]
        let mut t = vec![0.0; dl * 2 * dr * w];
        for a in 0..dl {
            for s in 0..2 {
                let brow = &bt[(a * 2 + s) * dr..(a * 2 + s + 1) * dr];
                for bp in 0..dr {
                    for w2 in 0..w {
                        let rrow = &r[(bp * w + w2) * dr..(bp * w + w2 + 1) * dr];
                        let mut acc = 0.0;
                        for (x, y) in brow.iter().zip(rrow) {
                            acc += x * y;
                        }
                        t[((a * 2 + s) * dr + bp) * w + w2] = acc;
                    }
                }
            }
        }
        // U[a, s', b', w1] = Σ_{w2, s} W[w1, w2, s', s] T[a, s, b', w2]
        let mut u = vec![0.0; dl * 2 * dr * w];
        for w1 in 0..w {
            for w2 in 0..w {
                if !self.nonzero(w1, w2) {
                    continue;
                }
                let o = self.op(w1, w2).0;
                for sp in 0..2 {
                    for s in 0..2 {
                        let m = o[sp * 2 + s];
                        if m == 0.0 {
                            continue;
                        }
                        for a in 0..dl {
                            for bp in 0..dr {
                                u[((a * 2 + sp) * dr + bp) * w + w1] += m * t[((a * 2 + s) * dr + bp) * w + w2];
                            }
                        }
                    }
                }
            }
        }
        // R'[a', w1, a] = Σ_{s', b'} B[a', s', b'] U[a, s', b', w1]
        let mut out = vec![0.0; dl * w * dl];
        for ap in 0..dl {
            for sp in 0..2 {
                for bp in 0..dr {
                    let bv = bt[(ap * 2 + sp) * dr + bp];
                    if bv == 0.0 {
                        continue;
                    }
                    for a in 0..dl {
                        for w1 in 0..w {
                            out[(ap * w + w1) * dl + a] += bv * u[((a * 2 + sp) * dr + bp) * w + w1];
                        }
                    }
                }
            }
        }
        out
    }

    /// Rows `rows` (bra indices `a'`) of `H_eff θ`, written to `out` (which
    /// holds just those rows).
    #[allow(clippy::too_many_arguments)]
    fn apply_rows(&self, l: &[f64], r: &[f64], theta: &[f64], out: &mut [f64], dl: usize, dr: usize, rows: core::ops::Range<usize>) {
        let w = self.w;
        let first = rows.start;
        let nrows = rows.len();
        // T1[a', w0, s1, s2, b] = Σ_a L[a', w0, a] θ[a, s1, s2, b]
        let blk = 4 * dr;
        let mut t1 = vec![0.0; nrows * w * blk];
        for ap in 0..nrows {
            for w0 in 0..w {
                let lrow = &l[((first + ap) * w + w0) * dl..((first + ap) * w + w0 + 1) * dl];
                let dst = &mut t1[(ap * w + w0) * blk..(ap * w + w0 + 1) * blk];
                for (a, &lv) in lrow.iter().enumerate() {
                    if lv == 0.0 {
                        continue;
                    }
                    for (d, &x) in dst.iter_mut().zip(&theta[a * blk..(a + 1) * blk]) {
                        *d += lv * x;
                    }
                }
            }
        }
        // T2[a', w1, s1', s2, b] = Σ_{w0, s1} W[w0, w1, s1', s1] T1[a', w0, s1, s2, b]
        let half = 2 * dr;
        let mut t2 = vec![0.0; nrows * w * blk];
        for ap in 0..nrows {
            for w0 in 0..w {
                for w1 in 0..w {
                    if !self.nonzero(w0, w1) {
                        continue;
                    }
                    let o = self.op(w0, w1).0;
                    for s1p in 0..2 {
                        for s1 in 0..2 {
                            let m = o[s1p * 2 + s1];
                            if m == 0.0 {
                                continue;
                            }
                            let src = (ap * w + w0) * blk + s1 * half;
                            let dst = (ap * w + w1) * blk + s1p * half;
                            for k in 0..half {
                                t2[dst + k] += m * t1[src + k];
                            }
                        }
                    }
                }
            }
        }
        // T3[a', s1', w2, s2', b] = Σ_{w1, s2} W[w1, w2, s2', s2] T2[a', w1, s1', s2, b]
        let mut t3 = vec![0.0; nrows * 2 * w * 2 * dr];
        for ap in 0..nrows {
            for w1 in 0..w {
                for w2 in 0..w {
                    if !self.nonzero(w1, w2) {
                        continue;
                    }
                    let o = self.op(w1, w2).0;
                    for s2p in 0..2 {
                        for s2 in 0..2 {
                            let m = o[s2p * 2 + s2];
                            if m == 0.0 {
                                continue;
                            }
                            for s1p in 0..2 {
                                let src = (ap * w + w1) * blk + s1p * half + s2 * dr;
                                let dst = (((ap * 2 + s1p) * w + w2) * 2 + s2p) * dr;
                                for b in 0..dr {
                                    t3[dst + b] += m * t2[src + b];
                                }
                            }
                        }
                    }
                }
            }
        }
        // out[a', s1', s2', b'] = Σ_{w2, b} T3[a', s1', w2, s2', b] R[b', w2, b]
        for ap in 0..nrows {
            for s1p in 0..2 {
                for s2p in 0..2 {
                    for bp in 0..dr {
                        let mut acc = 0.0;
                        for w2 in 0..w {
                            let trow = &t3[(((ap * 2 + s1p) * w + w2) * 2 + s2p) * dr..(((ap * 2 + s1p) * w + w2) * 2 + s2p + 1) * dr];
                            let rrow = &r[(bp * w + w2) * dr..(bp * w + w2 + 1) * dr];
                            for (x, y) in trow.iter().zip(rrow) {
                                acc += x * y;
                            }
                        }
                        out[((ap * 2 + s1p) * 2 + s2p) * dr + bp] = acc;
                    }
                }
            }
        }
    }

    /// `out = H_eff θ` for `θ[a, s1, s2, b]` between `l` (dl) and `r` (dr).
    /// Blocks of bra rows run on separate threads; each row is computed in
    /// the same order whatever the split, so the result is the same bits.
    #[allow(clippy::too_many_arguments)]
    fn apply(&self, l: &[f64], r: &[f64], theta: &[f64], out: &mut [f64], dl: usize, dr: usize, threads: usize) {
        let row = 4 * dr;
        // At least eight rows a thread; small bonds stay on one thread, where
        // spawning would cost more than it saves.
        let threads = if cfg!(target_arch = "wasm32") { 1 } else { threads.max(1).min(dl / 8) };
        if threads <= 1 || dl * dr < 2048 {
            self.apply_rows(l, r, theta, out, dl, dr, 0..dl);
            return;
        }
        let per = dl.div_ceil(threads);
        std::thread::scope(|scope| {
            for (c, chunk) in out.chunks_mut(per * row).enumerate() {
                let start = c * per;
                let end = (start + per).min(dl);
                scope.spawn(move || self.apply_rows(l, r, theta, chunk, dl, dr, start..end));
            }
        });
    }
}

/// Split `m` (rows × cols) by the eigenvectors of `m mᵀ`, keeping at most
/// `max_bond` and dropping weight below `cutoff`. Returns the kept vectors
/// (rows × k), `uᵀ m` (k × cols), the discarded weight and the kept
/// normalised weights.
fn split(m: &[f64], rows: usize, cols: usize, max_bond: usize, cutoff: f64) -> (Vec<f64>, Vec<f64>, f64, Vec<f64>) {
    let mut rho = vec![0.0; rows * rows];
    for i in 0..rows {
        for j in i..rows {
            let mut s = 0.0;
            for c in 0..cols {
                s += m[i * cols + c] * m[j * cols + c];
            }
            rho[i * rows + j] = s;
            rho[j * rows + i] = s;
        }
    }
    let (vals, vecs) = symmetric_eigen(&rho, rows);
    let total: f64 = vals.iter().map(|v| v.max(0.0)).sum();
    // Largest first.
    let order: Vec<usize> = (0..rows).rev().collect();
    let mut k = 0;
    let mut kept_weight = 0.0;
    for &idx in &order {
        let w = vals[idx].max(0.0) / total;
        if k >= max_bond || (k > 0 && w < cutoff) {
            break;
        }
        kept_weight += w;
        k += 1;
    }
    let k = k.max(1);
    let mut u = vec![0.0; rows * k];
    let mut weights = Vec::with_capacity(k);
    for (c, &idx) in order.iter().take(k).enumerate() {
        weights.push(vals[idx].max(0.0) / total);
        for r in 0..rows {
            u[r * k + c] = vecs[r * rows + idx];
        }
    }
    let mut rest = vec![0.0; k * cols];
    for c in 0..k {
        for r in 0..rows {
            let uv = u[r * k + c];
            if uv == 0.0 {
                continue;
            }
            for j in 0..cols {
                rest[c * cols + j] += uv * m[r * cols + j];
            }
        }
    }
    (u, rest, (1.0 - kept_weight).max(0.0), weights)
}

fn transpose(m: &[f64], rows: usize, cols: usize) -> Vec<f64> {
    let mut t = vec![0.0; rows * cols];
    for r in 0..rows {
        for c in 0..cols {
            t[c * rows + r] = m[r * cols + c];
        }
    }
    t
}

fn entropy(weights: &[f64]) -> f64 {
    -weights.iter().filter(|&&p| p > 0.0).map(|&p| p * ln(p)).sum::<f64>()
}

/// Two-site DMRG for the ground state of `chain`.
pub fn dmrg(chain: &Chain, cfg: &DmrgConfig) -> DmrgResult {
    let n = chain.n;
    assert!(n >= 2, "DMRG needs at least two sites");
    let mpo = chain.mpo();
    let eng = Engine { mpo: &mpo, w: mpo.w };
    let w = mpo.w;
    let mut mps = Mps::random(n, cfg.max_bond.min(8), cfg.seed);
    // Bring the state into right-canonical form from the right end, and
    // build the right environments.
    let left_edge = {
        let mut l = vec![0.0; w];
        l[0] = 1.0;
        l
    };
    let right_edge = {
        let mut r = vec![0.0; w];
        r[w - 1] = 1.0;
        r
    };
    for i in (1..n).rev() {
        let (dl, dr) = (mps.dims[i], mps.dims[i + 1]);
        // Site as dl × (2·dr); orthonormalise its rows.
        let m = mps.sites[i].clone();
        let (u, rest, _, _) = split(&transpose(&m, dl, 2 * dr), 2 * dr, dl, dl, 0.0);
        // rest is k × dl = (uᵀ mᵀ); new site = uᵀ (k × 2dr), absorb mᵀ... keep it simple:
        let k = rest.len() / dl;
        let new_site = transpose(&u, 2 * dr, k);
        let carry = transpose(&rest, k, dl); // dl × k
        // Previous site absorbs `carry` on its right index.
        let (pl, pr) = (mps.dims[i - 1], mps.dims[i]);
        let prev = &mps.sites[i - 1];
        let mut np = vec![0.0; pl * 2 * k];
        for a in 0..pl * 2 {
            for c in 0..k {
                let mut s = 0.0;
                for b in 0..pr {
                    s += prev[a * pr + b] * carry[b * k + c];
                }
                np[a * k + c] = s;
            }
        }
        mps.sites[i - 1] = np;
        mps.sites[i] = new_site;
        mps.dims[i] = k;
    }
    let mut rights: Vec<Vec<f64>> = vec![Vec::new(); n + 1];
    rights[n] = right_edge.clone();
    for i in (1..n).rev() {
        rights[i] = eng.grow_right(&rights[i + 1], &mps.sites[i], mps.dims[i], mps.dims[i + 1]);
    }
    let mut lefts: Vec<Vec<f64>> = vec![Vec::new(); n + 1];
    lefts[0] = left_edge;
    let mut energy = 0.0;
    let mut sweep_energies = Vec::new();
    let mut entropies = vec![0.0; n - 1];
    let mut discarded = 0.0;
    for _sweep in 0..cfg.sweeps {
        discarded = 0.0;
        // Left to right: optimise (i, i+1), keep the left part.
        for i in 0..n - 1 {
            let (dl, dm, dr) = (mps.dims[i], mps.dims[i + 1], mps.dims[i + 2]);
            let theta = two_site(&mps.sites[i], &mps.sites[i + 1], dl, dm, dr);
            let (l, r) = (&lefts[i], &rights[i + 2]);
            let apply = |x: &[f64], y: &mut [f64]| eng.apply(l, r, x, y, dl, dr, cfg.threads);
            let (e, v) = lanczos(&apply, &theta, cfg.krylov, cfg.restarts, cfg.tol);
            energy = e;
            let (u, rest, disc, weights) = split(&v, dl * 2, 2 * dr, cfg.max_bond, cfg.cutoff);
            let k = weights.len();
            discarded = f64::max(discarded, disc);
            entropies[i] = entropy(&weights);
            mps.sites[i] = u;
            mps.sites[i + 1] = rest;
            mps.dims[i + 1] = k;
            lefts[i + 1] = eng.grow_left(&lefts[i], &mps.sites[i], dl, k);
        }
        // Right to left: optimise (i, i+1), keep the right part.
        for i in (0..n - 1).rev() {
            let (dl, dm, dr) = (mps.dims[i], mps.dims[i + 1], mps.dims[i + 2]);
            let theta = two_site(&mps.sites[i], &mps.sites[i + 1], dl, dm, dr);
            let (l, r) = (&lefts[i], &rights[i + 2]);
            let apply = |x: &[f64], y: &mut [f64]| eng.apply(l, r, x, y, dl, dr, cfg.threads);
            let (e, v) = lanczos(&apply, &theta, cfg.krylov, cfg.restarts, cfg.tol);
            energy = e;
            // Split on the right: rows of the right factor are eigenvectors
            // of θᵀθ.
            let vt = transpose(&v, dl * 2, 2 * dr);
            let (u, rest, disc, weights) = split(&vt, 2 * dr, dl * 2, cfg.max_bond, cfg.cutoff);
            let k = weights.len();
            discarded = f64::max(discarded, disc);
            entropies[i] = entropy(&weights);
            mps.sites[i + 1] = transpose(&u, 2 * dr, k);
            mps.sites[i] = transpose(&rest, k, dl * 2);
            mps.dims[i + 1] = k;
            rights[i + 1] = eng.grow_right(&rights[i + 2], &mps.sites[i + 1], k, dr);
        }
        sweep_energies.push(energy);
    }
    DmrgResult { energy, sweep_energies, discarded, entropies, mps }
}

/// `θ[a, s1, s2, b] = Σ_c A[a, s1, c] B[c, s2, b]`.
fn two_site(a: &[f64], b: &[f64], dl: usize, dm: usize, dr: usize) -> Vec<f64> {
    let mut t = vec![0.0; dl * 4 * dr];
    for x in 0..dl * 2 {
        for c in 0..dm {
            let av = a[x * dm + c];
            if av == 0.0 {
                continue;
            }
            for s2 in 0..2 {
                let src = &b[(c * 2 + s2) * dr..(c * 2 + s2 + 1) * dr];
                let dst = &mut t[(x * 2 + s2) * dr..(x * 2 + s2 + 1) * dr];
                for (d, &v) in dst.iter_mut().zip(src) {
                    *d += av * v;
                }
            }
        }
    }
    t
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_eigensolver_reconstructs_its_matrix() {
        let mut s = 7u64;
        for n in [1usize, 2, 5, 17, 40] {
            let mut a = vec![0.0; n * n];
            for i in 0..n {
                for j in i..n {
                    s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
                    let v = ((s >> 11) as f64 / 9_007_199_254_740_992.0) - 0.5;
                    a[i * n + j] = v;
                    a[j * n + i] = v;
                }
            }
            let (vals, vecs) = symmetric_eigen(&a, n);
            for w in vals.windows(2) {
                assert!(w[0] <= w[1]);
            }
            for r in 0..n {
                for c in 0..n {
                    let rec: f64 = (0..n).map(|k| vecs[r * n + k] * vals[k] * vecs[c * n + k]).sum();
                    assert!((rec - a[r * n + c]).abs() < 1e-12, "n={n}");
                    let orth: f64 = (0..n).map(|k| vecs[k * n + r] * vecs[k * n + c]).sum();
                    assert!((orth - if r == c { 1.0 } else { 0.0 }).abs() < 1e-12, "n={n}");
                }
            }
        }
    }

    #[test]
    fn free_fermions_solve_the_ising_chain() {
        for n in [4usize, 7, 10] {
            for (j, h) in [(1.0, 0.3), (1.0, 1.0), (0.7, 1.8)] {
                let exact = exact_ground_energy(&Chain::ising(n, j, h)).unwrap();
                let ff = ising_exact_energy(n, j, h);
                assert!((exact - ff).abs() < 1e-9, "n={n} J={j} h={h}: {exact} vs {ff}");
            }
        }
    }

    #[test]
    fn dmrg_matches_exact_diagonalisation() {
        let cfg = DmrgConfig { max_bond: 32, sweeps: 4, ..DmrgConfig::default() };
        for chain in [Chain::heisenberg(12, 1.0, 1.0, 0.0), Chain::heisenberg(10, 1.0, 0.5, 0.3), Chain::ising(12, 1.0, 1.0), Chain::ising(11, 1.0, 0.6)] {
            let exact = exact_ground_energy(&chain).unwrap();
            let r = dmrg(&chain, &cfg);
            assert!((r.energy - exact).abs() < 1e-8, "{chain:?}: {} vs {exact}", r.energy);
        }
    }

    #[test]
    fn dmrg_reaches_the_critical_ising_chain_at_forty_sites() {
        let n = 40;
        let exact = ising_exact_energy(n, 1.0, 1.0);
        let r = dmrg(&Chain::ising(n, 1.0, 1.0), &DmrgConfig { max_bond: 24, sweeps: 3, ..DmrgConfig::default() });
        assert!(r.energy >= exact - 1e-9, "variational: {} vs {exact}", r.energy);
        assert!((r.energy - exact).abs() / (n as f64) < 1e-8, "{} vs {exact}", r.energy);
        // Entanglement peaks mid-chain on an open critical chain.
        let mid = r.entropies[n / 2 - 1];
        assert!(mid > r.entropies[2] && mid > r.entropies[n - 4]);
    }

    #[test]
    fn threaded_effective_hamiltonian_is_bit_identical() {
        // Environments big enough to split across threads.
        let mpo = Chain::heisenberg(8, 1.0, 0.7, 0.2).mpo();
        let eng = Engine { mpo: &mpo, w: mpo.w };
        let (dl, dr, w) = (56usize, 48usize, mpo.w);
        let mut s = 3u64;
        let mut next = move || {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((s >> 11) as f64 / 9_007_199_254_740_992.0) - 0.5
        };
        let l: Vec<f64> = (0..dl * w * dl).map(|_| next()).collect();
        let r: Vec<f64> = (0..dr * w * dr).map(|_| next()).collect();
        let theta: Vec<f64> = (0..dl * 4 * dr).map(|_| next()).collect();
        let mut one = vec![0.0; theta.len()];
        eng.apply(&l, &r, &theta, &mut one, dl, dr, 1);
        for t in [2, 3, 7] {
            let mut many = vec![0.0; theta.len()];
            eng.apply(&l, &r, &theta, &mut many, dl, dr, t);
            assert!(one.iter().zip(&many).all(|(a, b)| a.to_bits() == b.to_bits()), "threads={t}");
        }
    }

    #[test]
    fn energy_falls_sweep_by_sweep_and_runs_repeat() {
        let chain = Chain::heisenberg(16, 1.0, 1.0, 0.0);
        let cfg = DmrgConfig { max_bond: 16, sweeps: 3, ..DmrgConfig::default() };
        let a = dmrg(&chain, &cfg);
        for w in a.sweep_energies.windows(2) {
            assert!(w[1] <= w[0] + 1e-10, "{:?}", a.sweep_energies);
        }
        assert_eq!(a, dmrg(&chain, &cfg));
        // Threads split the effective Hamiltonian by rows; no bit moves.
        let one = dmrg(&chain, &DmrgConfig { threads: 1, ..cfg });
        let many = dmrg(&chain, &DmrgConfig { threads: 5, ..cfg });
        assert_eq!(one, many);
    }
}