wai-quantum 0.3.37

//! Time evolution of spin chains by TDVP — `wai.quantum.tdvp`.
//!
//! A claim that a quantum device simulates the dynamics of matter beyond
//! classical reach is a claim about time evolution. In one dimension the
//! time-dependent variational principle carries a matrix-product state
//! forward with its error held by the bond it keeps; until the entanglement
//! outgrows that bond, the dynamics are followed to rounding. This module is
//! that method, deterministic and held to exact answers.
//!
//! - **Two-site TDVP.** [`Evolution`] sweeps a matrix-product state left and
//!   right once a step, a symmetric second-order scheme.
//!   - Each bond's two-site tensor goes forward half a step under its
//!     effective Hamiltonian; the new centre then goes back half a step under
//!     its one-site one.
//!   - The exponentials are Lanczos with full re-orthogonalisation
//!     ([`expm_krylov`]). A step that does not converge in `krylov` vectors
//!     is halved.
//!   - Bonds are split by a singular value decomposition (one-sided Jacobi),
//!     so amplitudes are resolved to rounding, not to its square root as a
//!     density-matrix split would resolve them.
//! - **Complete bases.** TDVP's projection error comes from bases too small
//!   to express `H|ψ⟩`, and a product state's bases are as small as they
//!   come. [`Evolution::new`] completes every bond's basis with orthonormal
//!   null vectors up to `max_bond`: the state is unchanged and the tangent
//!   space grows. With `cutoff = 0` the bases stay complete, and a chain that
//!   fits within `max_bond` then evolves with every projector the identity:
//!   exact to rounding at any step size.
//! - **Measurements.** Site and bond expectations, the entanglement entropy
//!   at every bond, the energy, and the overlap with any other state, which
//!   gives the Loschmidt echo.
//! - **Exact referees.**
//!   - [`exact_evolve`] evolves all `2ⁿ` amplitudes by Krylov steps (at most
//!     20 sites).
//!   - [`ising_quench`] solves the transverse-field Ising chain started with
//!     every `σˣ = +1` as free fermions, at any length: `σˣ` on every site,
//!     `σᶻσᶻ` on every bond, and the Loschmidt echo as the overlap of two
//!     Gaussian states.
//!
//! The engine is the DMRG module's, on complex numbers. Every sum runs in a
//! fixed order and the effective Hamiltonians split across threads by rows,
//! so a run is the same bits on every machine and at any thread count.
//!
//! # Checked
//!
//! - **The referees agree.** The free-fermion solution matches exact
//!   evolution to `10⁻¹⁰` on 8 sites for three quenches, echo included. The
//!   Krylov exponential matches the dense one to `10⁻¹¹`, including a long
//!   step it has to halve.
//! - **Exact when the basis is complete.** On 8 sites with complete bases,
//!   four steps of 0.5 match exact evolution to `10⁻¹¹` in the state. Two
//!   cases are checked: Heisenberg from the Néel state, and Ising from a
//!   tilted product state.
//! - **Truncated, against free fermions.** 16 sites at bond 8 follow the
//!   critical Ising quench to `10⁻⁶` in `σˣ`, `σᶻσᶻ` and the echo.
//! - **At scale** (`tdvp_quench` example; the largest error in `σˣ` over
//!   every site, against the free-fermion answer):
//!
//!   | sites | bond | t   | error in σˣ | discarded weight | entropy, mid |
//!   |-------|------|-----|-------------|------------------|--------------|
//!   | 100   | 64   | 1   | 5.6·10⁻¹¹   | 1·10⁻¹³          | 0.74         |
//!   | 100   | 64   | 2   | 2.5·10⁻¹⁰   | 3·10⁻¹³          | 1.39         |
//!   | 100   | 64   | 3   | 9.0·10⁻¹⁰   | 4·10⁻¹⁰          | 2.02         |
//!   | 100   | 64   | 3.5 | 6.7·10⁻⁷    | 8·10⁻⁷           | 2.34         |
//!   | 100   | 64   | 4   | 1.6·10⁻⁴    | 1.5·10⁻⁴         | 2.66         |
//!   | 100   | 32   | 2.5 | 1.8·10⁻⁸    | 9·10⁻⁹           | 1.70         |
//!   | 100   | 32   | 3   | 1.5·10⁻⁵    | 8·10⁻⁶           | 2.02         |
//!   | 100   | 32   | 3.5 | 1.9·10⁻³    | 1.7·10⁻³         | 2.34         |
//!
//!   The critical quench (`J = h = 1`) at step 0.1. Until the entanglement
//!   fills the bond, the error stays below `10⁻⁹` and the Loschmidt rate
//!   agrees with the exact one to seven digits. Once it does, the error
//!   follows the discarded weight up, within a factor of three, which is why
//!   that weight is reported. The bond buys time: at `t = 3` the error is
//!   `1.5·10⁻⁵` at bond 32 and `9·10⁻¹⁰` at bond 64.
//!
//! - **The tests bite.** Lengthening the backward half-step by 0.1 % fails
//!   both accuracy tests.
//! - **Conserved and repeatable.** Energy holds to `10⁻⁹` over ten truncated
//!   steps. Runs repeat bit for bit at any thread count, and the same bits
//!   come out natively and under `wasm32-wasip2`.
//!
//! # Measured, and not kept
//!
//! A fourth-order composition of the symmetric step (Yoshida's three-stage
//! one) was built and measured on the 40-site critical quench at bond 32. At
//! time steps 0.05, 0.1 and 0.2 it was less accurate than the second-order
//! step: `1.1·10⁻⁹` against `2.3·10⁻¹⁰` at `t = 2` for step 0.05. It was also
//! three to ten times slower. Its backward middle stage entangles the state
//! the next split must truncate, and the time-step error it removes is not
//! where TDVP's error lies.
//!
//! # Honest boundaries
//!
//! - **The same Hamiltonians as DMRG**: one dimension, nearest neighbours,
//!   real terms.
//! - **Entanglement sets the reach.** After a quench it grows linearly in
//!   time, so the bond needed grows exponentially. Past that point a result
//!   carries truncation error. The discarded weight is reported; it is not
//!   converted into a bound on observables.
//! - **The cutoff trades accuracy for speed.** A cutoff drops basis vectors
//!   the state does not yet use, and with them some of the tangent space.
//!   On the 40-site critical quench at bond 32 and step 0.05, the error in
//!   `σˣ` at `t = 2` is `4·10⁻⁶` at cutoff `10⁻¹²`, `3·10⁻⁸` at `10⁻¹⁶` and
//!   `2·10⁻¹⁰` at `10⁻²⁰`, the default. At zero it is `4·10⁻¹⁰`: by then the
//!   bond is full and truncation sets the error whatever the cutoff, at three
//!   times the cost.

use crate::linalg::jacobi_svd_strict;
use crate::quantum_dmrg::{C, Chain, Engine, Mpo, Mps, Op, Scalar, apply_dense, complete_bases, default_threads, entropy, hermitian_eigen, right_canonicalise, symmetric_eigen, transpose, two_site};
use crate::repro::sin_cos;

// ---------------------------------------------------------------------------
// Dense helpers
// ---------------------------------------------------------------------------

fn dotc(a: &[C], b: &[C]) -> C {
    let mut s = C::ZERO;
    for (x, y) in a.iter().zip(b) {
        s = s.add(x.conj().mul(*y));
    }
    s
}

fn norm(v: &[C]) -> f64 {
    v.iter().map(|x| x.norm2()).sum::<f64>().sqrt()
}

// ---------------------------------------------------------------------------
// The Krylov exponential
// ---------------------------------------------------------------------------

/// `exp(−i τ H) v` for a Hermitian `H` given by `apply`: Lanczos with full
/// re-orthogonalisation, then the exponential of the small tridiagonal
/// matrix. It stops once the next Krylov vector would carry less than `tol`
/// of the result; if `krylov` vectors are not enough, the step is halved and
/// taken twice.
pub fn expm_krylov(apply: &dyn Fn(&[C], &mut [C]), v: &[C], tau: f64, krylov: usize, tol: f64) -> Vec<C> {
    expm_halving(apply, v, tau, krylov, tol, 0)
}

fn expm_halving(apply: &dyn Fn(&[C], &mut [C]), v: &[C], tau: f64, krylov: usize, tol: f64, depth: u32) -> Vec<C> {
    match expm_once(apply, v, tau, krylov, tol) {
        Some(out) => out,
        None => {
            // Each halving shrinks the error estimate by at least half once
            // two vectors are allowed, so this depth means the input is not
            // finite or the operator not Hermitian.
            assert!(depth < 60, "the Krylov exponential did not converge: is the state finite and the operator Hermitian?");
            let half = expm_halving(apply, v, tau / 2.0, krylov, tol, depth + 1);
            expm_halving(apply, &half, tau / 2.0, krylov, tol, depth + 1)
        }
    }
}

fn expm_once(apply: &dyn Fn(&[C], &mut [C]), v: &[C], tau: f64, krylov: usize, tol: f64) -> Option<Vec<C>> {
    let dim = v.len();
    let beta0 = norm(v);
    if beta0 == 0.0 || tau == 0.0 {
        return Some(v.to_vec());
    }
    // Two vectors at least: with one, the error estimate does not fall as
    // the step does.
    let k = krylov.max(2).min(dim);
    let mut basis: Vec<Vec<C>> = vec![v.iter().map(|x| x.scale(1.0 / beta0)).collect()];
    let (mut alpha, mut beta): (Vec<f64>, Vec<f64>) = (Vec::new(), Vec::new());
    let mut w = vec![C::ZERO; dim];
    loop {
        let j = basis.len() - 1;
        apply(&basis[j], &mut w);
        alpha.push(dotc(&basis[j], &w).re);
        // Full re-orthogonalisation, twice.
        for _ in 0..2 {
            for b in &basis {
                let c = dotc(b, &w);
                for (x, y) in w.iter_mut().zip(b) {
                    *x = x.sub(c.mul(*y));
                }
            }
        }
        let bn = norm(&w);
        let m = alpha.len();
        let coeffs = exp_tridiagonal(&alpha, &beta, tau);
        let err = bn * coeffs[m - 1].norm2().sqrt();
        if err <= tol || m == dim {
            let mut out = vec![C::ZERO; dim];
            for (c, b) in coeffs.iter().zip(&basis) {
                let c = c.scale(beta0);
                for (o, x) in out.iter_mut().zip(b) {
                    *o = o.add(c.mul(*x));
                }
            }
            return Some(out);
        }
        if m == k {
            return None;
        }
        beta.push(bn);
        basis.push(w.iter().map(|x| x.scale(1.0 / bn)).collect());
    }
}

/// `exp(−i τ T) e₁` for the real symmetric tridiagonal `T` with diagonal
/// `alpha` and off-diagonal `beta`.
fn exp_tridiagonal(alpha: &[f64], beta: &[f64], tau: f64) -> Vec<C> {
    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);
    let phases: Vec<C> = vals
        .iter()
        .map(|&l| {
            let (s, c) = sin_cos(tau * l);
            C { re: c, im: -s }
        })
        .collect();
    (0..m)
        .map(|r| {
            let mut s = C::ZERO;
            for (l, p) in phases.iter().enumerate() {
                s = s.add(p.scale(vecs[r * m + l] * vecs[l]));
            }
            s
        })
        .collect()
}

// ---------------------------------------------------------------------------
// TDVP
// ---------------------------------------------------------------------------

/// Settings for [`Evolution`].
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TdvpConfig {
    /// The time step.
    pub dt: f64,
    pub max_bond: usize,
    /// Discard Schmidt weight below this at each split. Zero keeps every
    /// basis vector up to `max_bond`, null ones included, so the tangent
    /// space is as large as the bond allows.
    pub cutoff: f64,
    /// Lanczos vectors per exponential before its step is halved.
    pub krylov: usize,
    /// The error allowed in each exponential, relative to the state.
    pub tol: f64,
    /// Threads for the effective Hamiltonians; the result does not depend
    /// on it.
    pub threads: usize,
}

impl Default for TdvpConfig {
    fn default() -> TdvpConfig {
        TdvpConfig { dt: 0.05, max_bond: 64, cutoff: 1e-20, krylov: 30, tol: 1e-12, threads: default_threads() }
    }
}

/// What one step did.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct StepReport {
    /// Density-matrix weight discarded, summed over every split of the step.
    pub discarded: f64,
    /// The largest bond after the step.
    pub max_bond: usize,
}

/// A matrix-product state evolving under a chain's Hamiltonian by the
/// two-site time-dependent variational principle. Between steps the state is
/// right-canonical with its centre on site 0.
#[derive(Clone, Debug, PartialEq)]
pub struct Evolution {
    mpo: Mpo,
    mps: Mps<C>,
    /// `rights[i]` is the environment of sites `i..n`.
    rights: Vec<Vec<C>>,
    cfg: TdvpConfig,
    steps: u64,
}

impl Evolution {
    /// Start from `state` (normalised here) under `chain`.
    pub fn new(chain: &Chain, state: &Mps<C>, cfg: TdvpConfig) -> Evolution {
        let n = chain.n;
        assert!(n >= 2, "TDVP needs at least two sites");
        assert_eq!(state.n(), n, "the state and the chain differ in length");
        let mpo = chain.mpo();
        let mut mps = state.clone();
        right_canonicalise(&mut mps);
        let nrm = norm(&mps.sites[0]);
        assert!(nrm > 0.0, "the state is zero");
        mps.sites[0].iter_mut().for_each(|x| *x = x.scale(1.0 / nrm));
        complete_bases(&mut mps, cfg.max_bond);
        let eng = Engine::new(&mpo);
        let mut rights: Vec<Vec<C>> = vec![Vec::new(); n + 1];
        rights[n] = eng.right_edge();
        for i in (1..n).rev() {
            rights[i] = eng.grow_right(i, &rights[i + 1], &mps.sites[i], mps.dims[i], mps.dims[i + 1]);
        }
        Evolution { mpo, mps, rights, cfg, steps: 0 }
    }

    /// Advance by one time step.
    pub fn step(&mut self) -> StepReport {
        let discarded = self.sweep_pair(self.cfg.dt);
        self.steps += 1;
        StepReport { discarded, max_bond: self.max_bond() }
    }

    /// One symmetric step of length `tau`: left to right and back, each
    /// bond forward by `tau / 2` and each centre in between back by as much.
    /// Returns the weight discarded.
    fn sweep_pair(&mut self, tau: f64) -> f64 {
        let n = self.mps.n();
        let cfg = self.cfg;
        let half = tau / 2.0;
        let eng = Engine::new(&self.mpo);
        let mps = &mut self.mps;
        let rights = &mut self.rights;
        let mut lefts: Vec<Vec<C>> = vec![Vec::new(); n + 1];
        lefts[0] = eng.left_edge();
        let mut discarded = 0.0;
        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 theta = expm_krylov(&|x: &[C], y: &mut [C]| eng.apply(i, l, r, x, y, dl, dr, cfg.threads), &theta, half, cfg.krylov, cfg.tol);
            let (u, mut rest, disc, weights) = svd_split(&theta, dl * 2, 2 * dr, cfg.max_bond, cfg.cutoff);
            if disc > 0.0 {
                renormalise(&mut rest, norm(&theta));
            }
            discarded += disc;
            let k = weights.len();
            mps.sites[i] = u;
            mps.dims[i + 1] = k;
            lefts[i + 1] = eng.grow_left(i, &lefts[i], &mps.sites[i], dl, k);
            mps.sites[i + 1] = if i + 2 < n {
                let (l, r) = (&lefts[i + 1], &rights[i + 2]);
                expm_krylov(&|x: &[C], y: &mut [C]| eng.apply1(i + 1, l, r, x, y, k, dr, cfg.threads), &rest, -half, cfg.krylov, cfg.tol)
            } else {
                rest
            };
        }
        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 theta = expm_krylov(&|x: &[C], y: &mut [C]| eng.apply(i, l, r, x, y, dl, dr, cfg.threads), &theta, half, cfg.krylov, cfg.tol);
            let (u, rest, disc, weights) = svd_split(&transpose(&theta, dl * 2, 2 * dr), 2 * dr, dl * 2, cfg.max_bond, cfg.cutoff);
            discarded += disc;
            let k = weights.len();
            mps.sites[i + 1] = transpose(&u, 2 * dr, k);
            mps.dims[i + 1] = k;
            rights[i + 1] = eng.grow_right(i + 1, &rights[i + 2], &mps.sites[i + 1], k, dr);
            let mut centre = transpose(&rest, k, dl * 2);
            if disc > 0.0 {
                renormalise(&mut centre, norm(&theta));
            }
            mps.sites[i] = if i > 0 {
                let (l, r) = (&lefts[i], &rights[i + 1]);
                expm_krylov(&|x: &[C], y: &mut [C]| eng.apply1(i, l, r, x, y, dl, k, cfg.threads), &centre, -half, cfg.krylov, cfg.tol)
            } else {
                centre
            };
        }
        discarded
    }

    /// The time evolved so far.
    pub fn time(&self) -> f64 {
        self.steps as f64 * self.cfg.dt
    }

    pub fn state(&self) -> &Mps<C> {
        &self.mps
    }

    pub fn max_bond(&self) -> usize {
        self.mps.dims.iter().copied().max().unwrap_or(1)
    }

    /// `‖ψ‖`: one, less whatever truncation has not put back.
    pub fn norm(&self) -> f64 {
        norm(&self.mps.sites[0])
    }

    /// `⟨ψ|H|ψ⟩`, from the centre and its environments.
    pub fn energy(&self) -> f64 {
        let eng = Engine::new(&self.mpo);
        let c = &self.mps.sites[0];
        let (dl, dr) = (self.mps.dims[0], self.mps.dims[1]);
        let mut hc = vec![C::ZERO; c.len()];
        eng.apply1(0, &eng.left_edge(), &self.rights[1], c, &mut hc, dl, dr, self.cfg.threads);
        dotc(c, &hc).re / dotc(c, c).re
    }

    /// `⟨O_i⟩` on every site. For a real symmetric `op` this is its
    /// expectation; otherwise its real part.
    pub fn expect(&self, op: Op) -> Vec<f64> {
        let mut out = Vec::with_capacity(self.mps.n());
        let mut env = vec![C::ONE];
        for i in 0..self.mps.n() {
            let (dl, dr) = (self.mps.dims[i], self.mps.dims[i + 1]);
            let m = &self.mps.sites[i];
            let t = env_times(&env, m, dl, 2 * dr);
            let mut v = C::ZERO;
            for ap in 0..dl {
                for sp in 0..2 {
                    for s in 0..2 {
                        let o = op.0[sp * 2 + s];
                        if o == 0.0 {
                            continue;
                        }
                        for b in 0..dr {
                            v = v.add(m[(ap * 2 + sp) * dr + b].conj().mul(t[(ap * 2 + s) * dr + b]).scale(o));
                        }
                    }
                }
            }
            out.push(v.re);
            env = env_close(m, &t, dl, dr);
        }
        out
    }

    /// `⟨A_i B_{i+1}⟩` on every bond, real parts as for [`Evolution::expect`].
    pub fn bond_expect(&self, a: Op, b: Op) -> Vec<f64> {
        let n = self.mps.n();
        let mut out = Vec::with_capacity(n - 1);
        let mut env = vec![C::ONE];
        for i in 0..n - 1 {
            let (dl, dm, dr) = (self.mps.dims[i], self.mps.dims[i + 1], self.mps.dims[i + 2]);
            let theta = two_site(&self.mps.sites[i], &self.mps.sites[i + 1], dl, dm, dr);
            let t = env_times(&env, &theta, dl, 4 * dr);
            let mut v = C::ZERO;
            for ap in 0..dl {
                for s1p in 0..2 {
                    for s2p in 0..2 {
                        for s1 in 0..2 {
                            let oa = a.0[s1p * 2 + s1];
                            if oa == 0.0 {
                                continue;
                            }
                            for s2 in 0..2 {
                                let ob = b.0[s2p * 2 + s2];
                                if ob == 0.0 {
                                    continue;
                                }
                                for c in 0..dr {
                                    let bra = theta[((ap * 2 + s1p) * 2 + s2p) * dr + c].conj();
                                    v = v.add(bra.mul(t[((ap * 2 + s1) * 2 + s2) * dr + c]).scale(oa * ob));
                                }
                            }
                        }
                    }
                }
            }
            out.push(v.re);
            let m = &self.mps.sites[i];
            let t1 = env_times(&env, m, dl, 2 * dm);
            env = env_close(m, &t1, dl, dm);
        }
        out
    }

    /// Von Neumann entanglement entropy at each of the `n − 1` bonds: the
    /// state's Schmidt weights there are the eigenvalues of the Gram matrix
    /// of its left part, since its right part is orthonormal.
    pub fn entropies(&self) -> Vec<f64> {
        let n = self.mps.n();
        let mut out = Vec::with_capacity(n - 1);
        let mut env = vec![C::ONE];
        for i in 0..n - 1 {
            let (dl, dr) = (self.mps.dims[i], self.mps.dims[i + 1]);
            let m = &self.mps.sites[i];
            let t = env_times(&env, m, dl, 2 * dr);
            env = env_close(m, &t, dl, dr);
            let (vals, _) = hermitian_eigen(&env, dr);
            let total: f64 = vals.iter().map(|v| v.max(0.0)).sum();
            let weights: Vec<f64> = vals.iter().map(|v| v.max(0.0) / total).collect();
            out.push(entropy(&weights));
        }
        out
    }

    /// `⟨φ|ψ⟩` with another state of the same length.
    pub fn overlap(&self, phi: &Mps<C>) -> C {
        overlap(phi, &self.mps)
    }
}

/// Split `m` (rows × cols) by its singular value decomposition, keeping at
/// most `max_bond` (and never more than its rank allows) and dropping weight
/// below `cutoff`. Returns the kept left vectors (rows × k), `u† m`
/// (k × cols), the discarded weight and the kept normalised weights.
///
/// One-sided Jacobi runs on `m†`, so the left vectors are its accumulated
/// rotations: orthonormal to rounding however small the singular value, and
/// a complete basis to draw null vectors from when the cutoff is zero.
/// Amplitudes are resolved to rounding, not to its square root as a
/// density-matrix split would.
fn svd_split(m: &[C], rows: usize, cols: usize, max_bond: usize, cutoff: f64) -> (Vec<C>, Vec<C>, f64, Vec<f64>) {
    // Columns of m† are the conjugated rows of m.
    let adj: Vec<Vec<C>> = (0..rows).map(|r| m[r * cols..(r + 1) * cols].iter().map(|x| x.conj()).collect()).collect();
    let (_, s, v) = jacobi_svd_strict(adj, cols, rows);
    let total: f64 = s.iter().map(|x| x * x).sum();
    let cap = max_bond.min(rows).min(cols).max(1);
    let mut k = 0;
    let mut kept = 0.0;
    for &sv in &s {
        let w = if total > 0.0 { sv * sv / total } else { 0.0 };
        if k >= cap || (k > 0 && w < cutoff) {
            break;
        }
        kept += w;
        k += 1;
    }
    let k = k.max(1);
    let mut u = vec![C::ZERO; rows * k];
    for (c, col) in v.iter().take(k).enumerate() {
        for r in 0..rows {
            u[r * k + c] = col[r];
        }
    }
    let mut rest = vec![C::ZERO; k * cols];
    for c in 0..k {
        for r in 0..rows {
            let uv = u[r * k + c];
            if uv.is_zero() {
                continue;
            }
            let uv = uv.conj();
            for j in 0..cols {
                rest[c * cols + j] = rest[c * cols + j].add(uv.mul(m[r * cols + j]));
            }
        }
    }
    let weights = s.iter().take(k).map(|x| if total > 0.0 { x * x / total } else { 0.0 }).collect();
    (u, rest, (1.0 - kept).max(0.0), weights)
}

fn renormalise(v: &mut [C], target: f64) {
    let nv = norm(v);
    if nv > 0.0 {
        v.iter_mut().for_each(|x| *x = x.scale(target / nv));
    }
}

/// `t[a', x] = Σ_a env[a', a] m[a, x]` for `m` with `dl` rows of `cols`.
fn env_times(env: &[C], m: &[C], dl: usize, cols: usize) -> Vec<C> {
    let mut t = vec![C::ZERO; dl * cols];
    for ap in 0..dl {
        let dst = &mut t[ap * cols..(ap + 1) * cols];
        for a in 0..dl {
            let e = env[ap * dl + a];
            if e.is_zero() {
                continue;
            }
            for (d, &x) in dst.iter_mut().zip(&m[a * cols..(a + 1) * cols]) {
                *d = d.add(e.mul(x));
            }
        }
    }
    t
}

/// `env'[b', b] = Σ_{a', s} m*[a', s, b'] t[a', s, b]`.
fn env_close(m: &[C], t: &[C], dl: usize, dr: usize) -> Vec<C> {
    let mut out = vec![C::ZERO; dr * dr];
    for ap in 0..dl {
        for s in 0..2 {
            for bp in 0..dr {
                let bra = m[(ap * 2 + s) * dr + bp];
                if bra.is_zero() {
                    continue;
                }
                let bra = bra.conj();
                let dst = &mut out[bp * dr..(bp + 1) * dr];
                for (d, &x) in dst.iter_mut().zip(&t[(ap * 2 + s) * dr..(ap * 2 + s + 1) * dr]) {
                    *d = d.add(bra.mul(x));
                }
            }
        }
    }
    out
}

/// `⟨φ|ψ⟩` for two matrix-product states of the same length.
pub fn overlap(phi: &Mps<C>, psi: &Mps<C>) -> C {
    assert_eq!(phi.n(), psi.n(), "the states differ in length");
    // env[a', a]: a' indexes φ's bond, a indexes ψ's.
    let mut env = vec![C::ONE];
    for i in 0..psi.n() {
        let (pl, pr) = (phi.dims[i], phi.dims[i + 1]);
        let (dl, dr) = (psi.dims[i], psi.dims[i + 1]);
        // t[a', s, b] = Σ_a env[a', a] ψ[a, s, b]
        let mut t = vec![C::ZERO; pl * 2 * dr];
        for ap in 0..pl {
            for a in 0..dl {
                let e = env[ap * dl + a];
                if e.is_zero() {
                    continue;
                }
                for (d, &x) in t[ap * 2 * dr..(ap + 1) * 2 * dr].iter_mut().zip(&psi.sites[i][a * 2 * dr..(a + 1) * 2 * dr]) {
                    *d = d.add(e.mul(x));
                }
            }
        }
        let mut next = vec![C::ZERO; pr * dr];
        for ap in 0..pl {
            for s in 0..2 {
                for bp in 0..pr {
                    let bra = phi.sites[i][(ap * 2 + s) * pr + bp];
                    if bra.is_zero() {
                        continue;
                    }
                    let bra = bra.conj();
                    for (d, &x) in next[bp * dr..(bp + 1) * dr].iter_mut().zip(&t[(ap * 2 + s) * dr..(ap * 2 + s + 1) * dr]) {
                        *d = d.add(bra.mul(x));
                    }
                }
            }
        }
        env = next;
    }
    env[0]
}

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

/// `exp(−iHt) ψ` for `chain` on all `2ⁿ` amplitudes (at most 20 sites), in
/// `steps` Krylov steps.
pub fn exact_evolve(chain: &Chain, psi: &[C], t: f64, steps: usize) -> Vec<C> {
    let n = chain.n;
    assert!((1..=20).contains(&n), "exact evolution takes 1 to 20 sites");
    assert_eq!(psi.len(), 1 << n);
    let apply = |x: &[C], y: &mut [C]| {
        let re: Vec<f64> = x.iter().map(|c| c.re).collect();
        let im: Vec<f64> = x.iter().map(|c| c.im).collect();
        let (mut hre, mut him) = (vec![0.0; x.len()], vec![0.0; x.len()]);
        apply_dense(chain, &re, &mut hre);
        apply_dense(chain, &im, &mut him);
        for ((o, r), i) in y.iter_mut().zip(hre).zip(him) {
            *o = C { re: r, im: i };
        }
    };
    let steps = steps.max(1);
    let mut v = psi.to_vec();
    for _ in 0..steps {
        v = expm_krylov(&apply, &v, t / steps as f64, 40, 1e-14);
    }
    v
}

/// `⟨ψ|O_site|ψ⟩` on a dense state of `n` sites (real part).
pub fn dense_expect(psi: &[C], n: usize, site: usize, op: Op) -> f64 {
    let shift = n - 1 - site;
    let mut v = C::ZERO;
    for (x, &amp) in psi.iter().enumerate() {
        let s = (x >> shift) & 1;
        for sp in 0..2 {
            let o = op.0[sp * 2 + s];
            if o != 0.0 {
                let y = (x & !(1 << shift)) | (sp << shift);
                v = v.add(psi[y].conj().mul(amp).scale(o));
            }
        }
    }
    v.re
}

/// `⟨ψ|A_site B_{site+1}|ψ⟩` on a dense state of `n` sites (real part).
pub fn dense_bond_expect(psi: &[C], n: usize, site: usize, a: Op, b: Op) -> f64 {
    let (sa, sb) = (n - 1 - site, n - 2 - site);
    let mut v = C::ZERO;
    for (x, &amp) in psi.iter().enumerate() {
        let (s1, s2) = ((x >> sa) & 1, (x >> sb) & 1);
        for s1p in 0..2 {
            let oa = a.0[s1p * 2 + s1];
            if oa == 0.0 {
                continue;
            }
            for s2p in 0..2 {
                let ob = b.0[s2p * 2 + s2];
                if ob == 0.0 {
                    continue;
                }
                let y = (x & !((1 << sa) | (1 << sb))) | (s1p << sa) | (s2p << sb);
                v = v.add(psi[y].conj().mul(amp).scale(oa * ob));
            }
        }
    }
    v.re
}

/// The transverse-field Ising chain after a quench, solved exactly.
#[derive(Clone, Debug, PartialEq)]
pub struct IsingQuench {
    /// `⟨σˣ_i(t)⟩` on every site.
    pub x: Vec<f64>,
    /// `⟨σᶻ_i σᶻ_{i+1}(t)⟩` on every bond.
    pub zz: Vec<f64>,
    /// The Loschmidt echo `|⟨ψ(0)|ψ(t)⟩|²`.
    pub echo: f64,
}

/// The open chain `H = −J Σ σᶻσᶻ − h Σ σˣ`, started with every `σˣ = +1`
/// and evolved for time `t`, as free fermions.
///
/// With Majorana operators `γ_{2i} = (Π_{k<i} σˣ_k) σᶻ_i` and
/// `γ_{2i+1} = (Π_{k<i} σˣ_k) σʸ_i`, the Hamiltonian is quadratic,
/// `H = (i/4) Σ h_ab γ_a γ_b` with `h_{2i,2i+1} = −2h` and
/// `h_{2i+1,2i+2} = −2J`, and the operators evolve as `γ(t) = e^{ht} γ`. The
/// correlation matrix `Γ_ab = ⟨i γ_a γ_b⟩` then evolves as `R Γ Rᵀ`, which
/// gives `σˣ_i = i γ_{2i} γ_{2i+1}` and `σᶻ_i σᶻ_{i+1} = i γ_{2i+1} γ_{2i+2}`
/// directly. The echo is the overlap of two Gaussian states,
/// `|⟨ψ₁|ψ₂⟩|² = |det((Γ₁ + Γ₂)/2)|^{1/2}`.
pub fn ising_quench(n: usize, j: f64, h: f64, t: f64) -> IsingQuench {
    assert!(n >= 1, "the chain needs a site");
    let m = 2 * n;
    let mut a = vec![0.0; m * m];
    for i in 0..n {
        a[(2 * i) * m + 2 * i + 1] = -2.0 * h * t;
        a[(2 * i + 1) * m + 2 * i] = 2.0 * h * t;
        if i + 1 < n {
            a[(2 * i + 1) * m + 2 * i + 2] = -2.0 * j * t;
            a[(2 * i + 2) * m + 2 * i + 1] = 2.0 * j * t;
        }
    }
    let r = expm_real(&a, m);
    // Γ(t)_ab = Σ_k (R_{a,2k} R_{b,2k+1} − R_{a,2k+1} R_{b,2k}).
    let gamma = |p: usize, q: usize| -> f64 {
        let mut s = 0.0;
        for k in 0..n {
            s += r[p * m + 2 * k] * r[q * m + 2 * k + 1] - r[p * m + 2 * k + 1] * r[q * m + 2 * k];
        }
        s
    };
    let x = (0..n).map(|i| gamma(2 * i, 2 * i + 1)).collect();
    let zz = (0..n.saturating_sub(1)).map(|i| gamma(2 * i + 1, 2 * i + 2)).collect();
    // (Γ(0) + Γ(t)) / 2.
    let mut s = vec![0.0; m * m];
    for p in 0..m {
        for q in 0..m {
            if p != q {
                s[p * m + q] = 0.5 * gamma(p, q);
            }
        }
    }
    for i in 0..n {
        s[(2 * i) * m + 2 * i + 1] += 0.5;
        s[(2 * i + 1) * m + 2 * i] -= 0.5;
    }
    let echo = determinant(s, m).abs().sqrt();
    IsingQuench { x, zz, echo }
}

/// `e^A` for a real `n × n` matrix: Taylor's series to degree 18 after
/// halving `A` until its 1-norm is at most ½, then squaring back.
fn expm_real(a: &[f64], n: usize) -> Vec<f64> {
    let norm1 = (0..n).map(|c| (0..n).map(|r| a[r * n + c].abs()).sum::<f64>()).fold(0.0, f64::max);
    let mut halvings = 0;
    let mut scale = 1.0;
    while norm1 * scale > 0.5 {
        scale *= 0.5;
        halvings += 1;
    }
    let b: Vec<f64> = a.iter().map(|x| x * scale).collect();
    let identity = |m: &mut [f64]| {
        for i in 0..n {
            m[i * n + i] += 1.0;
        }
    };
    // Horner: I + B (I + B/2 (I + … (I + B/18))).
    let mut e = vec![0.0; n * n];
    identity(&mut e);
    for k in (1..=18).rev() {
        let mut next = matmul(&b, &e, n);
        next.iter_mut().for_each(|x| *x /= k as f64);
        identity(&mut next);
        e = next;
    }
    for _ in 0..halvings {
        e = matmul(&e, &e, n);
    }
    e
}

fn matmul(a: &[f64], b: &[f64], n: usize) -> Vec<f64> {
    let mut c = vec![0.0; n * n];
    for i in 0..n {
        for k in 0..n {
            let av = a[i * n + k];
            if av == 0.0 {
                continue;
            }
            for (d, &x) in c[i * n..(i + 1) * n].iter_mut().zip(&b[k * n..(k + 1) * n]) {
                *d += av * x;
            }
        }
    }
    c
}

/// The determinant by Gaussian elimination with partial pivoting.
fn determinant(mut a: Vec<f64>, n: usize) -> f64 {
    let mut det = 1.0;
    for c in 0..n {
        let mut p = c;
        for r in c + 1..n {
            if a[r * n + c].abs() > a[p * n + c].abs() {
                p = r;
            }
        }
        if a[p * n + c] == 0.0 {
            return 0.0;
        }
        if p != c {
            for k in 0..n {
                a.swap(c * n + k, p * n + k);
            }
            det = -det;
        }
        let piv = a[c * n + c];
        det *= piv;
        for r in c + 1..n {
            let f = a[r * n + c] / piv;
            if f == 0.0 {
                continue;
            }
            for k in c..n {
                a[r * n + k] -= f * a[c * n + k];
            }
        }
    }
    det
}

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

    fn rng(seed: u64) -> impl FnMut() -> f64 {
        let mut s = seed;
        move || {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((s >> 11) as f64 / 9_007_199_254_740_992.0) - 0.5
        }
    }

    #[test]
    fn the_krylov_exponential_matches_the_dense_one() {
        let mut next = rng(5);
        let n = 24;
        let mut h = vec![C::ZERO; n * n];
        for i in 0..n {
            h[i * n + i] = C { re: 3.0 * next(), im: 0.0 };
            for j in i + 1..n {
                let v = C { re: next(), im: next() };
                h[i * n + j] = v;
                h[j * n + i] = v.conj();
            }
        }
        let v: Vec<C> = (0..n).map(|_| C { re: next(), im: next() }).collect();
        let (vals, vecs) = hermitian_eigen(&h, n);
        let apply = |x: &[C], y: &mut [C]| {
            for r in 0..n {
                let mut s = C::ZERO;
                for c in 0..n {
                    s = s.add(h[r * n + c].mul(x[c]));
                }
                y[r] = s;
            }
        };
        // Short steps converge in a few vectors; a long one with few vectors
        // has to halve itself.
        for (tau, krylov) in [(0.1, 30), (-0.7, 30), (6.0, 8)] {
            let got = expm_krylov(&apply, &v, tau, krylov, 1e-13);
            // Dense: X e^{−iτΛ} X† v.
            for r in 0..n {
                let mut want = C::ZERO;
                for k in 0..n {
                    let mut proj = C::ZERO;
                    for c in 0..n {
                        proj = proj.add(vecs[c * n + k].conj().mul(v[c]));
                    }
                    let (s, co) = sin_cos(tau * vals[k]);
                    want = want.add(vecs[r * n + k].mul(C { re: co, im: -s }).mul(proj));
                }
                assert!(got[r].sub(want).norm2().sqrt() < 1e-11, "tau={tau} row {r}");
            }
        }
    }

    #[test]
    fn free_fermions_solve_the_quench() {
        let n = 8;
        let plus = Mps::product(&vec![[C::ONE, C::ONE]; n]).to_dense().unwrap();
        for (j, h, t) in [(1.0, 1.0, 0.7), (1.0, 0.4, 2.0), (0.6, 1.3, 1.1)] {
            let chain = Chain::ising(n, j, h);
            let psi = exact_evolve(&chain, &plus, t, 8);
            let ff = ising_quench(n, j, h, t);
            for i in 0..n {
                let x = dense_expect(&psi, n, i, Op::x());
                assert!((x - ff.x[i]).abs() < 1e-10, "J={j} h={h} t={t} x[{i}]: {x} vs {}", ff.x[i]);
            }
            for i in 0..n - 1 {
                let zz = dense_bond_expect(&psi, n, i, Op::z(), Op::z());
                assert!((zz - ff.zz[i]).abs() < 1e-10, "zz[{i}]: {zz} vs {}", ff.zz[i]);
            }
            let echo = dotc(&plus, &psi).norm2() / dotc(&plus, &plus).re.powi(2);
            assert!((echo - ff.echo).abs() < 1e-10, "echo: {echo} vs {}", ff.echo);
        }
    }

    #[test]
    fn tdvp_is_exact_when_the_basis_is_complete() {
        // Completed bases make every projector the identity: no projection
        // error and no splitting error, at any step, from product states.
        let n = 8;
        let neel: Vec<bool> = (0..n).map(|i| i % 2 == 1).collect();
        let mut next = rng(9);
        let tilted: Vec<[C; 2]> = (0..n).map(|_| [C { re: 1.0, im: 0.0 }, C { re: next(), im: next() }]).collect();
        let cases = [(Chain::heisenberg(n, 1.0, 0.6, 0.2), Mps::basis(&neel)), (Chain::ising(n, 1.0, 0.7), Mps::product(&tilted))];
        for (chain, start) in cases {
            let cfg = TdvpConfig { dt: 0.5, max_bond: 16, cutoff: 0.0, ..TdvpConfig::default() };
            let mut ev = Evolution::new(&chain, &start, cfg);
            for _ in 0..4 {
                ev.step();
            }
            let exact = exact_evolve(&chain, &start.to_dense().unwrap(), ev.time(), 8);
            let got = ev.state().to_dense().unwrap();
            let err = got.iter().zip(&exact).map(|(a, b)| a.sub(*b).norm2()).sum::<f64>().sqrt();
            assert!(err < 1e-11, "{chain:?}: state error {err:e}");
            let z = ev.expect(Op::z());
            for (i, zi) in z.iter().enumerate() {
                assert!((zi - dense_expect(&exact, n, i, Op::z())).abs() < 1e-11);
            }
        }
    }

    #[test]
    fn truncated_tdvp_follows_the_quench() {
        // Sixteen sites at bond 8: the state no longer fits, so this is
        // truncation and projection error, held to the free-fermion answer.
        let n = 16;
        let chain = Chain::ising(n, 1.0, 1.0);
        let plus = Mps::product(&vec![[C::ONE, C::ONE]; n]);
        let cfg = TdvpConfig { dt: 0.1, max_bond: 8, ..TdvpConfig::default() };
        let mut ev = Evolution::new(&chain, &plus, cfg);
        for _ in 0..10 {
            ev.step();
        }
        let ff = ising_quench(n, 1.0, 1.0, ev.time());
        let x = ev.expect(Op::x());
        let zz = ev.bond_expect(Op::z(), Op::z());
        for (i, (got, want)) in x.iter().zip(&ff.x).enumerate() {
            assert!((got - want).abs() < 1e-6, "x[{i}]: {got} vs {want}");
        }
        for (i, (got, want)) in zz.iter().zip(&ff.zz).enumerate() {
            assert!((got - want).abs() < 1e-6, "zz[{i}]: {got} vs {want}");
        }
        let echo = ev.overlap(&plus).norm2();
        assert!((echo - ff.echo).abs() < 1e-6, "echo: {echo} vs {}", ff.echo);
        // Entanglement peaks mid-chain.
        let s = ev.entropies();
        assert!(s[n / 2 - 1] > s[0] && s[n / 2 - 1] > s[n - 2]);
    }

    #[test]
    fn energy_is_conserved_and_runs_repeat() {
        let n = 12;
        let chain = Chain::heisenberg(n, 1.0, 0.8, 0.1);
        let start = Mps::basis(&(0..n).map(|i| i % 2 == 0).collect::<Vec<_>>());
        let cfg = TdvpConfig { dt: 0.1, max_bond: 12, ..TdvpConfig::default() };
        let mut ev = Evolution::new(&chain, &start, cfg);
        let e0 = ev.energy();
        let want = 0.8 * 0.25 * -(n as f64 - 1.0);
        assert!((e0 - want).abs() < 1e-12, "{e0} vs {want}");
        for _ in 0..10 {
            ev.step();
        }
        assert!((ev.energy() - e0).abs() < 1e-9, "{} vs {e0}", ev.energy());
        assert!((ev.norm() - 1.0).abs() < 1e-12);
        // Same inputs, same bits, at any thread count.
        let mut one = Evolution::new(&chain, &start, TdvpConfig { threads: 1, ..cfg });
        let mut many = Evolution::new(&chain, &start, TdvpConfig { threads: 5, ..cfg });
        for _ in 0..10 {
            one.step();
            many.step();
        }
        assert_eq!(one.state(), many.state());
        assert_eq!(one.state(), ev.state());
    }

    #[test]
    fn threaded_complex_effective_hamiltonians_are_bit_identical() {
        let mpo = Chain::heisenberg(8, 1.0, 0.7, 0.2).mpo();
        let eng = Engine::new(&mpo);
        let (dl, dr, w) = (56usize, 48usize, mpo.dims[2]);
        let mut next = rng(3);
        let mut c = || C { re: next(), im: next() };
        let l: Vec<C> = (0..dl * w * dl).map(|_| c()).collect();
        let r: Vec<C> = (0..dr * w * dr).map(|_| c()).collect();
        let theta: Vec<C> = (0..dl * 4 * dr).map(|_| c()).collect();
        let site: Vec<C> = (0..dl * 2 * dr).map(|_| c()).collect();
        let r1: Vec<C> = r.clone();
        let mut one = vec![C::ZERO; theta.len()];
        let mut one1 = vec![C::ZERO; site.len()];
        eng.apply(2, &l, &r, &theta, &mut one, dl, dr, 1);
        eng.apply1(2, &l, &r1, &site, &mut one1, dl, dr, 1);
        for t in [2, 3, 7] {
            let mut many = vec![C::ZERO; theta.len()];
            let mut many1 = vec![C::ZERO; site.len()];
            eng.apply(2, &l, &r, &theta, &mut many, dl, dr, t);
            eng.apply1(2, &l, &r1, &site, &mut many1, dl, dr, t);
            assert_eq!(one, many, "threads={t}");
            assert_eq!(one1, many1, "threads={t}");
        }
    }
}