open-standards/wai/wai-quantum/src/quantum_bptn.rs

//! Tensor-network states on the hardware graph, evolved by simple update and
//! read through belief-propagation environments (`wai.quantum.bptn`).
//!
//! The state is a tensor network laid on the device's coupling graph: one
//! tensor per qubit, with a physical index and one virtual index per edge. It
//! is kept in Vidal gauge, with a vector of bond weights `λ_e` on every edge. A
//! single-qubit gate acts on one tensor. A two-qubit gate on an edge acts on
//! the two tensors and the bond between them:
//!
//! 1. Each tensor, with the weights of its other bonds absorbed, is QR-reduced.
//! 2. The gate is applied to the two small factors and the bond weight between
//!    them.
//! 3. An SVD splits the result again, keeping at most `χ` singular values.
//!    These become the new bond weights.
//!
//! An expectation value takes each tensor's environment to be the product of
//! its bond weights. That is the rank-one environment belief propagation
//! converges to, and on a tree it is exact.
//!
//! What to expect from it:
//!
//! - **On a tree** (a chain, or any loop-free coupling graph), with `χ` large
//!   enough to truncate nothing, the evolution and the expectation values are
//!   exact up to rounding. The tests hold it to sparse Pauli dynamics there.
//! - **On a graph with loops**, such as heavy-hex with its 12-qubit rings, the
//!   environments neglect correlations around the loops. Truncation discards
//!   weight as well, and [`Tns::discarded`] reports how much.
//!
//! Its errors are unrelated to those of sparse Pauli dynamics, which is what
//! makes it worth having beside that method: where the two agree, neither
//! method's approximation is likely to be what they agree on.
//!
//! Arithmetic: `+ − × ÷ √` in a fixed order (Householder QR, one-sided Jacobi
//! SVD), with `sin` and `cos` computed in-crate. Results are bit-identical on
//! every platform.

#![allow(clippy::needless_range_loop)]

use crate::linalg::{C, jacobi_svd_strict, qr};
use crate::quantum_spd::{Op, Pauli, RotCircuit};
use crate::repro::sin_cos;
use std::fmt;

/// Why a circuit could not be run on a graph.
#[derive(Clone, Debug, PartialEq)]
pub enum TnsError {
    /// A qubit outside the graph.
    QubitOutOfRange(u32),
    /// A two-qubit gate on qubits the graph does not connect.
    NotAnEdge(u32, u32),
    /// A rotation about a string of more than two qubits.
    TooWide(usize),
    /// A circuit `quantum_spd` would refuse.
    Invalid(String),
}

impl fmt::Display for TnsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TnsError::QubitOutOfRange(q) => write!(f, "qubit {q} is not on the graph"),
            TnsError::NotAnEdge(a, b) => write!(f, "no edge between {a} and {b}"),
            TnsError::TooWide(w) => write!(f, "a rotation on {w} qubits; at most 2 here"),
            TnsError::Invalid(e) => write!(f, "{e}"),
        }
    }
}

impl std::error::Error for TnsError {}

/// The state and how it was truncated.
pub struct Tns {
    n: usize,
    edges: Vec<(u32, u32)>,
    /// Per vertex: incident edge ids, ascending. A tensor's virtual indices
    /// follow this order.
    nbrs: Vec<Vec<usize>>,
    /// Per vertex: physical index outermost, then the virtual indices in
    /// `nbrs` order, row-major.
    t: Vec<Vec<C>>,
    /// Per edge: the bond weights, descending, with unit 2-norm.
    lam: Vec<Vec<f64>>,
    chi: usize,
    discarded: f64,
}

/// A singular value is dropped below this fraction of the largest. Dividing
/// a tensor by a smaller weight would amplify rounding past use.
const CUTOFF: f64 = 1e-10;

impl Tns {
    /// `|0…0⟩` on the graph with these edges, keeping at most `chi` singular
    /// values per bond.
    pub fn new(n: u32, edges: &[(u32, u32)], chi: usize) -> Result<Tns, TnsError> {
        let n = n as usize;
        let mut nbrs = vec![Vec::new(); n];
        for (e, &(a, b)) in edges.iter().enumerate() {
            for q in [a, b] {
                if q as usize >= n {
                    return Err(TnsError::QubitOutOfRange(q));
                }
            }
            if a == b {
                return Err(TnsError::NotAnEdge(a, b));
            }
            nbrs[a as usize].push(e);
            nbrs[b as usize].push(e);
        }
        // Every bond starts at dimension 1, so a tensor holds two numbers:
        // amplitude 1 for |0⟩, 0 for |1⟩.
        let t = (0..n).map(|_| vec![C::ONE, C::ZERO]).collect();
        Ok(Tns { n, edges: edges.to_vec(), nbrs, t, lam: vec![vec![1.0]; edges.len()], chi: chi.max(1), discarded: 0.0 })
    }

    /// The summed fraction of squared singular-value weight the truncations
    /// have discarded: 0 means no truncation has happened yet.
    pub fn discarded(&self) -> f64 {
        self.discarded
    }

    /// The largest bond dimension in the network.
    pub fn max_bond(&self) -> usize {
        self.lam.iter().map(Vec::len).max().unwrap_or(1)
    }

    fn bond_dims(&self, v: usize) -> Vec<usize> {
        self.nbrs[v].iter().map(|&e| self.lam[e].len()).collect()
    }

    fn edge_of(&self, a: u32, b: u32) -> Option<usize> {
        self.nbrs.get(a as usize)?.iter().copied().find(|&e| {
            let (x, y) = self.edges[e];
            (x, y) == (a, b) || (x, y) == (b, a)
        })
    }

    /// Apply every op of `circuit` in order.
    ///
    /// Consecutive two-qubit gates on vertex-disjoint edges, such as a layer of
    /// the device's couplings, touch disjoint tensors and bonds and read no
    /// bond another of them writes. They are computed in parallel and
    /// committed in order, which gives the same bits as applying them one by
    /// one (and is what happens on targets without threads).
    pub fn run(&mut self, circuit: &RotCircuit) -> Result<(), TnsError> {
        self.run_inner(circuit, None)
    }

    /// [`Self::run`], re-gauging the network into the belief-propagation gauge
    /// (at most `sweeps` sweeps to tolerance `tol`) after every layer of
    /// two-qubit gates, so each truncation is made against the environment
    /// belief propagation sees rather than the bond weights alone.
    pub fn run_gauged(&mut self, circuit: &RotCircuit, sweeps: usize, tol: f64) -> Result<(), TnsError> {
        self.run_inner(circuit, Some((sweeps, tol)))
    }

    fn run_inner(&mut self, circuit: &RotCircuit, gauge: Option<(usize, f64)>) -> Result<(), TnsError> {
        circuit.validate().map_err(|e| TnsError::Invalid(e.to_string()))?;
        if circuit.n as usize != self.n {
            return Err(TnsError::Invalid(format!("a {}-qubit circuit on a {}-qubit graph", circuit.n, self.n)));
        }
        let mut batch: Vec<(u32, u32, [[C; 4]; 4])> = Vec::new();
        let mut busy = vec![false; self.n];
        for op in &circuit.ops {
            match gate_of(op)? {
                Gate::One(q, g) => {
                    if busy[q as usize] {
                        self.flush(&mut batch, &mut busy, gauge)?;
                    }
                    self.gate1(q, g)?;
                }
                Gate::Two(a, b, g) => {
                    for q in [a, b] {
                        if q as usize >= self.n {
                            return Err(TnsError::QubitOutOfRange(q));
                        }
                    }
                    if busy[a as usize] || busy[b as usize] {
                        self.flush(&mut batch, &mut busy, gauge)?;
                    }
                    busy[a as usize] = true;
                    busy[b as usize] = true;
                    batch.push((a, b, g));
                }
            }
        }
        self.flush(&mut batch, &mut busy, gauge)
    }

    fn flush(
        &mut self,
        batch: &mut Vec<(u32, u32, [[C; 4]; 4])>,
        busy: &mut [bool],
        gauge: Option<(usize, f64)>,
    ) -> Result<(), TnsError> {
        if batch.is_empty() {
            return Ok(());
        }
        let threads = if cfg!(target_arch = "wasm32") {
            1
        } else {
            std::thread::available_parallelism().map_or(1, |n| n.get()).min(batch.len())
        };
        let updates: Vec<Result<Update, TnsError>> = if threads <= 1 {
            batch.iter().map(|&(a, b, g)| self.compute2(a, b, g)).collect()
        } else {
            let this = &*self;
            let chunk = batch.len().div_ceil(threads);
            std::thread::scope(|scope| {
                let handles: Vec<_> = batch
                    .chunks(chunk)
                    .map(|part| scope.spawn(move || part.iter().map(|&(a, b, g)| this.compute2(a, b, g)).collect::<Vec<_>>()))
                    .collect();
                handles.into_iter().flat_map(|h| h.join().expect("a gate's computation panicked")).collect()
            })
        };
        for up in updates {
            self.commit(up?);
        }
        for &(a, b, _) in batch.iter() {
            busy[a as usize] = false;
            busy[b as usize] = false;
        }
        batch.clear();
        if let Some((sweeps, tol)) = gauge {
            let env = self.belief_propagation(sweeps, tol);
            self.regauge(&env);
        }
        Ok(())
    }

    /// Apply one op.
    pub fn apply(&mut self, op: &Op) -> Result<(), TnsError> {
        match gate_of(op)? {
            Gate::One(q, g) => self.gate1(q, g),
            Gate::Two(a, b, g) => {
                let up = self.compute2(a, b, g)?;
                self.commit(up);
                Ok(())
            }
        }
    }

    fn gate1(&mut self, q: u32, g: [[C; 2]; 2]) -> Result<(), TnsError> {
        let v = q as usize;
        if v >= self.n {
            return Err(TnsError::QubitOutOfRange(q));
        }
        let half = self.t[v].len() / 2;
        let t = &mut self.t[v];
        for i in 0..half {
            let (a, b) = (t[i], t[half + i]);
            t[i] = g[0][0].mul(a).add(g[0][1].mul(b));
            t[half + i] = g[1][0].mul(a).add(g[1][1].mul(b));
        }
        Ok(())
    }

    /// What a two-qubit gate `g` (row and column index `2·s_a + s_b`) on edge
    /// `(a, b)` does, computed without changing the state: [`Self::commit`]
    /// applies it.
    fn compute2(&self, a: u32, b: u32, g: [[C; 4]; 4]) -> Result<Update, TnsError> {
        for q in [a, b] {
            if q as usize >= self.n {
                return Err(TnsError::QubitOutOfRange(q));
            }
        }
        let e = self.edge_of(a, b).ok_or(TnsError::NotAnEdge(a, b))?;
        let (u, v) = (a as usize, b as usize);
        let d = self.lam[e].len();

        // 1. Each side as a matrix: rows over its other bonds (their weights
        //    absorbed), columns over (physical, this bond); then QR.
        let (mu, ru, rest_u) = self.side_matrix(u, e);
        let (mv, rv, rest_v) = self.side_matrix(v, e);
        let (qu, rru, ku) = qr(&mu, ru, 2 * d);
        let (qv, rrv, kv) = qr(&mv, rv, 2 * d);

        // 2. θ[(x,s),(t,y)] = Σ_i Ru[x,(s,i)] λ[i] Rv[y,(t,i)], then the gate.
        let lam = &self.lam[e];
        let rows = 2 * ku;
        let cols = 2 * kv;
        let mut theta = vec![C::ZERO; rows * cols];
        for x in 0..ku {
            for s in 0..2 {
                for y in 0..kv {
                    for t in 0..2 {
                        let mut acc = C::ZERO;
                        for i in 0..d {
                            acc = acc.add(rru[x * 2 * d + s * d + i].scale(lam[i]).mul(rrv[y * 2 * d + t * d + i]));
                        }
                        theta[(x * 2 + s) * cols + t * kv + y] = acc;
                    }
                }
            }
        }
        let mut gated = vec![C::ZERO; rows * cols];
        for x in 0..ku {
            for y in 0..kv {
                for s2 in 0..2 {
                    for t2 in 0..2 {
                        let mut acc = C::ZERO;
                        for s in 0..2 {
                            for t in 0..2 {
                                acc = acc.add(g[2 * s2 + t2][2 * s + t].mul(theta[(x * 2 + s) * cols + t * kv + y]));
                            }
                        }
                        gated[(x * 2 + s2) * cols + t2 * kv + y] = acc;
                    }
                }
            }
        }

        // 3. SVD, truncated.
        let (left, sv, right) = svd_rows(&gated, rows, cols);
        let total: f64 = sv.iter().map(|x| x * x).sum();
        let top = sv.first().copied().unwrap_or(0.0);
        let keep = sv.iter().take(self.chi).take_while(|&&x| x > CUTOFF * top).count().max(1);
        let kept: f64 = sv[..keep].iter().map(|x| x * x).sum();
        let discarded = if total > 0.0 { (total - kept) / total } else { 0.0 };
        let norm = kept.sqrt();
        let new_lam: Vec<f64> = sv[..keep].iter().map(|x| x / norm).collect();

        // 4. New tensors: Q times the singular vectors, the other bonds'
        //    weights divided back out.
        let mut new_u = vec![C::ZERO; rest_u * 2 * keep];
        for o in 0..rest_u {
            for s in 0..2 {
                for k in 0..keep {
                    let mut acc = C::ZERO;
                    for x in 0..ku {
                        acc = acc.add(qu[o * ku + x].mul(left[k][x * 2 + s]));
                    }
                    new_u[(o * 2 + s) * keep + k] = acc;
                }
            }
        }
        let mut new_v = vec![C::ZERO; rest_v * 2 * keep];
        for o in 0..rest_v {
            for t in 0..2 {
                for k in 0..keep {
                    let mut acc = C::ZERO;
                    for y in 0..kv {
                        acc = acc.add(qv[o * kv + y].mul(right[k][t * kv + y].conj()));
                    }
                    new_v[(o * 2 + t) * keep + k] = acc;
                }
            }
        }
        Ok(Update {
            e,
            u,
            v,
            tu: self.rebuilt(u, e, &new_u, keep),
            tv: self.rebuilt(v, e, &new_v, keep),
            lam: new_lam,
            discarded,
        })
    }

    fn commit(&mut self, up: Update) {
        self.discarded += up.discarded;
        self.lam[up.e] = up.lam;
        self.t[up.u] = up.tu;
        self.t[up.v] = up.tv;
    }


    /// Vertex `v`'s tensor as a matrix: rows over its bonds other than `e`
    /// (row-major in `nbrs` order, each scaled by its bond weight), columns
    /// over (physical, bond `e`). Returns the matrix, its row count, and that
    /// count again for the caller's reshaping.
    fn side_matrix(&self, v: usize, e: usize) -> (Vec<C>, usize, usize) {
        let dims = self.bond_dims(v);
        let pos = self.nbrs[v].iter().position(|&x| x == e).expect("the edge is incident");
        let d = dims[pos];
        let bonds: usize = dims.iter().product();
        let rest = bonds / d;
        let mut m = vec![C::ZERO; rest * 2 * d];
        let mut idx = vec![0usize; dims.len()];
        for flat in 0..bonds {
            // unflatten
            let mut r = flat;
            for j in (0..dims.len()).rev() {
                idx[j] = r % dims[j];
                r /= dims[j];
            }
            let mut row = 0usize;
            let mut w = 1.0;
            for j in 0..dims.len() {
                if j != pos {
                    row = row * dims[j] + idx[j];
                    w *= self.lam[self.nbrs[v][j]][idx[j]];
                }
            }
            for s in 0..2 {
                m[row * 2 * d + s * d + idx[pos]] = self.t[v][s * bonds + flat].scale(w);
            }
        }
        (m, rest, rest)
    }

    /// Vertex `v`'s new tensor from `m` (rows over the other bonds, as in
    /// `side_matrix`, then physical, then the new bond `e` of size `keep`),
    /// with the other bonds' weights divided back out.
    fn rebuilt(&self, v: usize, e: usize, m: &[C], keep: usize) -> Vec<C> {
        let pos = self.nbrs[v].iter().position(|&x| x == e).expect("the edge is incident");
        let mut dims = self.bond_dims(v);
        dims[pos] = keep;
        let bonds: usize = dims.iter().product();
        let mut t = vec![C::ZERO; 2 * bonds];
        let mut idx = vec![0usize; dims.len()];
        for flat in 0..bonds {
            let mut r = flat;
            for j in (0..dims.len()).rev() {
                idx[j] = r % dims[j];
                r /= dims[j];
            }
            let mut row = 0usize;
            let mut w = 1.0;
            for j in 0..dims.len() {
                if j != pos {
                    row = row * dims[j] + idx[j];
                    w *= self.lam[self.nbrs[v][j]][idx[j]];
                }
            }
            let inv = 1.0 / w;
            for s in 0..2 {
                t[s * bonds + flat] = m[(row * 2 + s) * keep + idx[pos]].scale(inv);
            }
        }
        t
    }

    /// `⟨P⟩` for a Pauli on one qubit, from the tensor and the squared weights
    /// of its bonds: the rank-one environment belief propagation converges to.
    pub fn expect1(&self, q: u32, p: Pauli) -> Result<f64, TnsError> {
        let v = q as usize;
        if v >= self.n {
            return Err(TnsError::QubitOutOfRange(q));
        }
        let dims = self.bond_dims(v);
        let bonds: usize = dims.iter().product();
        let mut rho = [[C::ZERO; 2]; 2];
        let mut idx = vec![0usize; dims.len()];
        for flat in 0..bonds {
            let mut r = flat;
            for j in (0..dims.len()).rev() {
                idx[j] = r % dims[j];
                r /= dims[j];
            }
            let w: f64 = (0..dims.len()).map(|j| self.lam[self.nbrs[v][j]][idx[j]]).map(|l| l * l).product();
            for s in 0..2 {
                for s2 in 0..2 {
                    let x = self.t[v][s * bonds + flat].mul(self.t[v][s2 * bonds + flat].conj()).scale(w);
                    rho[s][s2] = rho[s][s2].add(x);
                }
            }
        }
        let pm = pauli(p);
        let mut num = C::ZERO;
        for s in 0..2 {
            for s2 in 0..2 {
                num = num.add(rho[s][s2].mul(pm[s2][s]));
            }
        }
        let den = rho[0][0].re + rho[1][1].re;
        Ok(num.re / den)
    }
}

/// Belief-propagation messages on the norm network `⟨ψ|ψ⟩`. For each edge
/// `e = (a, b)` there is a message from `a` to `b` and one from `b` to `a`:
/// `χ_e × χ_e` matrices, the row on the ket's bond index and the column on the
/// bra's, normalised to unit trace.
pub struct BpEnv {
    msgs: Vec<[Vec<C>; 2]>,
    /// Sweeps run.
    pub sweeps: usize,
    /// The largest change in any message entry during the last sweep.
    pub residual: f64,
}

impl Tns {
    /// Vertex `v`'s tensor with the square root of each of its bond weights
    /// absorbed, so that the norm network is the tensors and their conjugates
    /// alone.
    fn absorbed(&self, v: usize) -> Vec<C> {
        let dims = self.bond_dims(v);
        let bonds: usize = dims.iter().product();
        let mut out = self.t[v].clone();
        let mut idx = vec![0usize; dims.len()];
        for flat in 0..bonds {
            let mut r = flat;
            for j in (0..dims.len()).rev() {
                idx[j] = r % dims[j];
                r /= dims[j];
            }
            let w: f64 = (0..dims.len()).map(|j| self.lam[self.nbrs[v][j]][idx[j]].sqrt()).product();
            for s in 0..2 {
                out[s * bonds + flat] = out[s * bonds + flat].scale(w);
            }
        }
        out
    }

    /// The message arriving at `v` across its bond at position `j`.
    fn incoming<'a>(&self, msgs: &'a [[Vec<C>; 2]], v: usize, j: usize) -> &'a [C] {
        let e = self.nbrs[v][j];
        if self.edges[e].0 as usize == v { &msgs[e][1] } else { &msgs[e][0] }
    }

    /// `a` (vertex `v`'s absorbed tensor) with the incoming message contracted
    /// into every bond but `skip`: each such index now runs over the bra side.
    fn with_messages(&self, a: &[C], v: usize, msgs: &[[Vec<C>; 2]], skip: Option<usize>) -> Vec<C> {
        let dims = self.bond_dims(v);
        let bonds: usize = dims.iter().product();
        let mut cur = a.to_vec();
        for j in 0..dims.len() {
            if skip == Some(j) {
                continue;
            }
            let m = self.incoming(msgs, v, j);
            let d = dims[j];
            let stride: usize = dims[j + 1..].iter().product();
            let mut next = vec![C::ZERO; cur.len()];
            for s in 0..2 {
                for flat in 0..bonds {
                    let x = cur[s * bonds + flat];
                    if x.re == 0.0 && x.im == 0.0 {
                        continue;
                    }
                    let i = (flat / stride) % d;
                    let base = flat - i * stride;
                    for ip in 0..d {
                        let k = s * bonds + base + ip * stride;
                        next[k] = next[k].add(x.mul(m[i * d + ip]));
                    }
                }
            }
            cur = next;
        }
        cur
    }

    /// The message from `u` across its bond at position `p`, given `msgs`.
    fn message(&self, abs: &[Vec<C>], msgs: &[[Vec<C>; 2]], u: usize, p: usize) -> Vec<C> {
        let a = &abs[u];
        let c = self.with_messages(a, u, msgs, Some(p));
        let dims = self.bond_dims(u);
        let bonds: usize = dims.iter().product();
        let d = dims[p];
        let stride: usize = dims[p + 1..].iter().product();
        let mut m = vec![C::ZERO; d * d];
        for s in 0..2 {
            for base in 0..bonds {
                if (base / stride) % d != 0 {
                    continue;
                }
                for i in 0..d {
                    let ci = c[s * bonds + base + i * stride];
                    for ip in 0..d {
                        m[i * d + ip] = m[i * d + ip].add(ci.mul(a[s * bonds + base + ip * stride].conj()));
                    }
                }
            }
        }
        let tr: f64 = (0..d).map(|i| m[i * d + i].re).sum();
        if tr > 0.0 {
            for x in m.iter_mut() {
                *x = x.scale(1.0 / tr);
            }
        }
        m
    }

    /// Run belief propagation on the norm network to a fixed point: at most
    /// `max_sweeps` sweeps, stopping once no message entry changes by more
    /// than `tol`. Each sweep computes every message from the previous
    /// sweep's (in parallel where there are threads), so the result is the
    /// same everywhere.
    pub fn belief_propagation(&self, max_sweeps: usize, tol: f64) -> BpEnv {
        let abs: Vec<Vec<C>> = (0..self.n).map(|v| self.absorbed(v)).collect();
        let mut msgs: Vec<[Vec<C>; 2]> = self
            .lam
            .iter()
            .map(|l| {
                // In Vidal form the consistent message is diag(λ): start there.
                let d = l.len();
                let sum: f64 = l.iter().sum();
                let diag: Vec<C> = (0..d * d).map(|k| if k % (d + 1) == 0 { C::new(l[k / (d + 1)] / sum, 0.0) } else { C::ZERO }).collect();
                [diag.clone(), diag]
            })
            .collect();
        // Directed edges, in order: (edge, direction, sender, sender's bond position).
        let directed: Vec<(usize, usize, usize, usize)> = (0..self.edges.len())
            .flat_map(|e| {
                let (a, b) = (self.edges[e].0 as usize, self.edges[e].1 as usize);
                let pa = self.nbrs[a].iter().position(|&x| x == e).expect("incident");
                let pb = self.nbrs[b].iter().position(|&x| x == e).expect("incident");
                [(e, 0, a, pa), (e, 1, b, pb)]
            })
            .collect();
        let mut sweeps = 0;
        let mut residual = f64::INFINITY;
        while sweeps < max_sweeps && residual > tol {
            let threads = if cfg!(target_arch = "wasm32") {
                1
            } else {
                std::thread::available_parallelism().map_or(1, |n| n.get()).min(directed.len())
            };
            let fresh: Vec<Vec<C>> = if threads <= 1 {
                directed.iter().map(|&(_, _, u, p)| self.message(&abs, &msgs, u, p)).collect()
            } else {
                let (this, abs_ref, msgs_ref) = (&*self, &abs, &msgs);
                let chunk = directed.len().div_ceil(threads);
                std::thread::scope(|scope| {
                    let handles: Vec<_> = directed
                        .chunks(chunk)
                        .map(|part| scope.spawn(move || part.iter().map(|&(_, _, u, p)| this.message(abs_ref, msgs_ref, u, p)).collect::<Vec<_>>()))
                        .collect();
                    handles.into_iter().flat_map(|h| h.join().expect("a message's computation panicked")).collect()
                })
            };
            residual = 0.0;
            for (&(e, dir, _, _), m) in directed.iter().zip(fresh) {
                for (old, new) in msgs[e][dir].iter().zip(&m) {
                    residual = residual.max(old.sub(*new).norm2().sqrt());
                }
                msgs[e][dir] = m;
            }
            sweeps += 1;
        }
        BpEnv { msgs, sweeps, residual }
    }

    /// Bring the network into the gauge belief propagation's messages define:
    /// on each edge, factor the two messages as `M = X X†`, take the SVD
    /// `Xₐᵀ X_b = U Λ V†`, and put `(Xₐ⁺)ᵀ U` on one side and `(V† X_b⁺)ᵀ` on the
    /// other. That inserts the identity on the bond (on the support the
    /// messages see), so the state does not change, and `Λ` becomes the bond's
    /// weights: afterwards [`Self::expect1`]'s environment is the BP one.
    pub fn regauge(&mut self, env: &BpEnv) {
        let abs: Vec<Vec<C>> = (0..self.n).map(|v| self.absorbed(v)).collect();
        let mut transforms: Vec<[Vec<C>; 2]> = Vec::with_capacity(self.edges.len());
        let mut new_lam: Vec<Vec<f64>> = Vec::with_capacity(self.edges.len());
        for e in 0..self.edges.len() {
            let d = self.lam[e].len();
            let (xa, xap, ra) = psd_sqrt(&env.msgs[e][0], d);
            let (xb, xbp, rb) = psd_sqrt(&env.msgs[e][1], d);
            let mut bm = vec![C::ZERO; ra * rb];
            for p in 0..ra {
                for q in 0..rb {
                    let mut acc = C::ZERO;
                    for i in 0..d {
                        acc = acc.add(xa[i * ra + p].mul(xb[i * rb + q]));
                    }
                    bm[p * rb + q] = acc;
                }
            }
            let (u, sv, v) = svd_rows(&bm, ra, rb);
            let top = sv.first().copied().unwrap_or(0.0);
            let k = sv.iter().take_while(|&&x| x > CUTOFF * top).count().max(1);
            let norm = sv[..k].iter().map(|x| x * x).sum::<f64>().sqrt();
            new_lam.push(sv[..k].iter().map(|x| x / norm).collect());
            let mut ta = vec![C::ZERO; d * k];
            let mut tb = vec![C::ZERO; d * k];
            for i in 0..d {
                for n in 0..k {
                    let mut a_acc = C::ZERO;
                    for p in 0..ra {
                        a_acc = a_acc.add(xap[p * d + i].mul(u[n][p]));
                    }
                    ta[i * k + n] = a_acc;
                    let mut b_acc = C::ZERO;
                    for q in 0..rb {
                        b_acc = b_acc.add(v[n][q].conj().mul(xbp[q * d + i]));
                    }
                    tb[i * k + n] = b_acc;
                }
            }
            transforms.push([ta, tb]);
        }
        for v in 0..self.n {
            let mut dims = self.bond_dims(v);
            let mut cur = abs[v].clone();
            for j in 0..dims.len() {
                let e = self.nbrs[v][j];
                let side = if self.edges[e].0 as usize == v { 0 } else { 1 };
                let k = new_lam[e].len();
                cur = transform_bond(&cur, &dims, j, &transforms[e][side], k);
                dims[j] = k;
            }
            self.t[v] = cur;
        }
        self.lam = new_lam;
    }

    /// `⟨P⟩` for a Pauli on one qubit, with the environment `env` belief
    /// propagation converged to.
    pub fn expect1_bp(&self, env: &BpEnv, q: u32, p: Pauli) -> Result<f64, TnsError> {
        let v = q as usize;
        if v >= self.n {
            return Err(TnsError::QubitOutOfRange(q));
        }
        let a = self.absorbed(v);
        let c = self.with_messages(&a, v, &env.msgs, None);
        let bonds = a.len() / 2;
        let mut rho = [[C::ZERO; 2]; 2];
        for s in 0..2 {
            for s2 in 0..2 {
                for k in 0..bonds {
                    rho[s][s2] = rho[s][s2].add(c[s * bonds + k].mul(a[s2 * bonds + k].conj()));
                }
            }
        }
        let pm = pauli(p);
        let mut num = C::ZERO;
        for s in 0..2 {
            for s2 in 0..2 {
                num = num.add(rho[s][s2].mul(pm[s2][s]));
            }
        }
        Ok(num.re / (rho[0][0].re + rho[1][1].re))
    }
}

/// A gate as matrices: one qubit, or two (row and column index `2·s_a + s_b`).
enum Gate {
    One(u32, [[C; 2]; 2]),
    Two(u32, u32, [[C; 4]; 4]),
}

/// What a two-qubit gate does to the state, ready to commit.
struct Update {
    e: usize,
    u: usize,
    v: usize,
    tu: Vec<C>,
    tv: Vec<C>,
    lam: Vec<f64>,
    discarded: f64,
}

fn gate_of(op: &Op) -> Result<Gate, TnsError> {
    Ok(match op {
        Op::H(q) | Op::S(q) | Op::Sdg(q) | Op::SX(q) | Op::SXdg(q) | Op::X(q) | Op::Y(q) | Op::Z(q) => Gate::One(*q, fixed1(op)),
        Op::CX(a, b) => Gate::Two(*a, *b, cx()),
        Op::CZ(a, b) => Gate::Two(*a, *b, cz()),
        Op::Swap(a, b) => Gate::Two(*a, *b, swap()),
        Op::Rot { axis, theta } => {
            let (s, c) = sin_cos(0.5 * theta);
            rotation(axis, c, s)?
        }
        Op::QuarterRot { axis, k } => {
            // exp(−i(kπ/2)P/2): cos and sin of kπ/4, exactly.
            let h = std::f64::consts::FRAC_1_SQRT_2;
            let (c, s) = match k {
                1 => (h, h),
                2 => (0.0, 1.0),
                _ => (-h, h),
            };
            rotation(axis, c, s)?
        }
    })
}

/// `exp(−iθP/2) = cos(θ/2)·I − i·sin(θ/2)·P`, given `cos(θ/2), sin(θ/2)`.
fn rotation(axis: &[(u32, Pauli)], c: f64, s: f64) -> Result<Gate, TnsError> {
    let minus_is = C::new(0.0, -s);
    match axis {
        [(q, p)] => {
            let pm = pauli(*p);
            let mut g = [[C::ZERO; 2]; 2];
            for r in 0..2 {
                for col in 0..2 {
                    let id = if r == col { C::new(c, 0.0) } else { C::ZERO };
                    g[r][col] = id.add(minus_is.mul(pm[r][col]));
                }
            }
            Ok(Gate::One(*q, g))
        }
        [(a, pa), (b, pb)] => {
            let (ma, mb) = (pauli(*pa), pauli(*pb));
            let mut g = [[C::ZERO; 4]; 4];
            for r in 0..4 {
                for col in 0..4 {
                    let kron = ma[r >> 1][col >> 1].mul(mb[r & 1][col & 1]);
                    let id = if r == col { C::new(c, 0.0) } else { C::ZERO };
                    g[r][col] = id.add(minus_is.mul(kron));
                }
            }
            Ok(Gate::Two(*a, *b, g))
        }
        _ => Err(TnsError::TooWide(axis.len())),
    }
}

/// A Hermitian positive semi-definite `d × d` matrix (row-major) as `X X†`:
/// `X` (`d × r`, row-major), its pseudo-inverse `X⁺` (`r × d`), and `r`, the
/// number of eigenvalues above the cutoff.
fn psd_sqrt(m: &[C], d: usize) -> (Vec<C>, Vec<C>, usize) {
    let (u, s, _) = svd_rows(m, d, d);
    let top = s.first().copied().unwrap_or(0.0);
    let r = s.iter().take_while(|&&x| x > CUTOFF * top).count().max(1);
    let mut x = vec![C::ZERO; d * r];
    let mut xp = vec![C::ZERO; r * d];
    for k in 0..r {
        let root = s[k].sqrt();
        for i in 0..d {
            x[i * r + k] = u[k][i].scale(root);
            xp[k * d + i] = u[k][i].conj().scale(1.0 / root);
        }
    }
    (x, xp, r)
}

/// Contract bond `j` of a tensor (physical index outermost, then the bonds
/// row-major with dimensions `dims`) with the `dims[j] × k` matrix `m`.
fn transform_bond(t: &[C], dims: &[usize], j: usize, m: &[C], k: usize) -> Vec<C> {
    let d = dims[j];
    let bonds: usize = dims.iter().product();
    let stride: usize = dims[j + 1..].iter().product();
    let new_bonds = bonds / d * k;
    let mut out = vec![C::ZERO; 2 * new_bonds];
    for s in 0..2 {
        for flat in 0..bonds {
            let x = t[s * bonds + flat];
            if x.re == 0.0 && x.im == 0.0 {
                continue;
            }
            let pre = flat / (d * stride);
            let i = (flat / stride) % d;
            let post = flat % stride;
            for n in 0..k {
                let at = s * new_bonds + pre * k * stride + n * stride + post;
                out[at] = out[at].add(x.mul(m[i * k + n]));
            }
        }
    }
    out
}

/// SVD of a row-major `rows × cols` matrix: left singular vectors (each of
/// length `rows`), singular values (descending), right singular vectors (each
/// of length `cols`), with `A = Σ_k s_k · u_k · v_k†`. The one-sided Jacobi
/// sweeps run over the shorter dimension, to convergence: the singular vectors
/// must be orthonormal to rounding, or every later expectation inherits the error.
fn svd_rows(a: &[C], rows: usize, cols: usize) -> (Vec<Vec<C>>, Vec<f64>, Vec<Vec<C>>) {
    if cols <= rows {
        let columns: Vec<Vec<C>> = (0..cols).map(|c| (0..rows).map(|r| a[r * cols + c]).collect()).collect();
        jacobi_svd_strict(columns, rows, cols)
    } else {
        // A† = V S U†: decompose the conjugate transpose, whose columns are
        // the conjugated rows of A, and swap the roles.
        let columns: Vec<Vec<C>> = (0..rows).map(|r| (0..cols).map(|c| a[r * cols + c].conj()).collect()).collect();
        let (v, s, u) = jacobi_svd_strict(columns, cols, rows);
        (u, s, v)
    }
}

fn pauli(p: Pauli) -> [[C; 2]; 2] {
    let (z, o) = (C::ZERO, C::ONE);
    match p {
        Pauli::X => [[z, o], [o, z]],
        Pauli::Y => [[z, C::new(0.0, -1.0)], [C::new(0.0, 1.0), z]],
        Pauli::Z => [[o, z], [z, C::new(-1.0, 0.0)]],
    }
}

fn fixed1(op: &Op) -> [[C; 2]; 2] {
    let (z, o) = (C::ZERO, C::ONE);
    let (i, ni) = (C::new(0.0, 1.0), C::new(0.0, -1.0));
    let h = std::f64::consts::FRAC_1_SQRT_2;
    let (hp, hm) = (C::new(h, 0.0), C::new(-h, 0.0));
    let (pp, pm) = (C::new(0.5, 0.5), C::new(0.5, -0.5));
    match op {
        Op::H(_) => [[hp, hp], [hp, hm]],
        Op::S(_) => [[o, z], [z, i]],
        Op::Sdg(_) => [[o, z], [z, ni]],
        Op::SX(_) => [[pp, pm], [pm, pp]],
        Op::SXdg(_) => [[pm, pp], [pp, pm]],
        Op::X(_) => pauli(Pauli::X),
        Op::Y(_) => pauli(Pauli::Y),
        _ => pauli(Pauli::Z),
    }
}

fn cx() -> [[C; 4]; 4] {
    let (z, o) = (C::ZERO, C::ONE);
    [[o, z, z, z], [z, o, z, z], [z, z, z, o], [z, z, o, z]]
}

fn cz() -> [[C; 4]; 4] {
    let (z, o) = (C::ZERO, C::ONE);
    [[o, z, z, z], [z, o, z, z], [z, z, o, z], [z, z, z, C::new(-1.0, 0.0)]]
}

fn swap() -> [[C; 4]; 4] {
    let (z, o) = (C::ZERO, C::ONE);
    [[o, z, z, z], [z, z, o, z], [z, o, z, z], [z, z, z, o]]
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantum_spd::{SpdConfig, expect_z, expectation, heavy_hex_127, kicked_ising};

    /// Exact sparse Pauli dynamics: an independent oracle.
    fn exact(c: &RotCircuit, q: u32, p: Pauli) -> f64 {
        expectation(c, &[(1.0, vec![(q, p)])], &[], &SpdConfig { threshold: 0.0, max_weight: None }).unwrap().value
    }

    struct Rng(u64);
    impl Rng {
        fn next(&mut self) -> u64 {
            self.0 ^= self.0 << 13;
            self.0 ^= self.0 >> 7;
            self.0 ^= self.0 << 17;
            self.0
        }
        fn below(&mut self, n: u64) -> u64 {
            self.next() % n
        }
        fn angle(&mut self) -> f64 {
            (self.next() >> 11) as f64 / (1u64 << 53) as f64 * 6.0 - 3.0
        }
    }

    /// A random circuit whose two-qubit gates all lie on `edges`.
    fn random_on(rng: &mut Rng, n: u32, edges: &[(u32, u32)], len: usize) -> RotCircuit {
        let p = |r: &mut Rng| [Pauli::X, Pauli::Y, Pauli::Z][r.below(3) as usize];
        let mut c = RotCircuit::new(n);
        for _ in 0..len {
            let q = rng.below(u64::from(n)) as u32;
            let (a, b) = edges[rng.below(edges.len() as u64) as usize];
            let (a, b) = if rng.below(2) == 0 { (a, b) } else { (b, a) };
            let op = match rng.below(12) {
                0 => Op::H(q),
                1 => Op::S(q),
                2 => Op::SX(q),
                3 => Op::SXdg(q),
                4 => Op::Y(q),
                5 => Op::CX(a, b),
                6 => Op::CZ(a, b),
                7 => Op::Swap(a, b),
                8 => Op::QuarterRot { axis: vec![(a, p(rng)), (b, p(rng))], k: 1 + rng.below(3) as u8 },
                9 => Op::Rot { axis: vec![(a, p(rng)), (b, p(rng))], theta: rng.angle() },
                _ => Op::Rot { axis: vec![(q, p(rng))], theta: rng.angle() },
            };
            c.push(op);
        }
        c
    }

    #[test]
    fn on_a_tree_with_no_truncation_it_is_exact() {
        // A chain and a branching tree. With χ = 2^⌊n/2⌋ nothing is ever
        // truncated, so the Vidal-gauge evolution and its environments are
        // exact, and must agree with exact sparse Pauli dynamics to rounding.
        let chain: Vec<(u32, u32)> = (0..7).map(|q| (q, q + 1)).collect();
        let tree = vec![(0, 1), (1, 2), (1, 3), (3, 4), (3, 5), (0, 6), (6, 7)];
        let mut rng = Rng(0x5eed);
        for edges in [chain, tree] {
            for trial in 0..25 {
                let c = random_on(&mut rng, 8, &edges, 40);
                let mut tns = Tns::new(8, &edges, 16).unwrap();
                tns.run(&c).unwrap();
                assert!(tns.discarded() < 1e-20, "nothing truncated: {}", tns.discarded());
                for q in 0..8 {
                    for p in [Pauli::X, Pauli::Y, Pauli::Z] {
                        let (got, want) = (tns.expect1(q, p).unwrap(), exact(&c, q, p));
                        assert!((got - want).abs() < 1e-12, "trial {trial} q {q} {p:?}: {got} vs {want}");
                    }
                }
            }
        }
    }

    #[test]
    fn belief_propagation_is_exact_on_a_tree() {
        let tree = vec![(0, 1), (1, 2), (1, 3), (3, 4), (3, 5), (0, 6), (6, 7)];
        let mut rng = Rng(0xbe1);
        for trial in 0..15 {
            let c = random_on(&mut rng, 8, &tree, 40);
            let mut tns = Tns::new(8, &tree, 16).unwrap();
            tns.run(&c).unwrap();
            let env = tns.belief_propagation(200, 1e-14);
            assert!(env.residual <= 1e-14, "trial {trial}: BP converged ({} sweeps, {:e})", env.sweeps, env.residual);
            for q in 0..8 {
                for p in [Pauli::X, Pauli::Y, Pauli::Z] {
                    let (got, want) = (tns.expect1_bp(&env, q, p).unwrap(), exact(&c, q, p));
                    assert!((got - want).abs() < 1e-11, "trial {trial} q {q} {p:?}: {got} vs {want}");
                }
            }
        }
    }

    #[test]
    fn regauging_changes_no_expectation_and_makes_the_bond_weights_the_environment() {
        let edges = heavy_hex_127();
        let mut tns = Tns::new(127, &edges, 8).unwrap();
        tns.run(&kicked_ising(127, &edges, 0.6, 4)).unwrap();
        let env = tns.belief_propagation(200, 1e-13);
        let before: Vec<f64> = (0..127).map(|q| tns.expect1_bp(&env, q, Pauli::Z).unwrap()).collect();
        tns.regauge(&env);
        let env2 = tns.belief_propagation(200, 1e-13);
        for q in 0..127 {
            let (after_bp, after_simple) = (tns.expect1_bp(&env2, q, Pauli::Z).unwrap(), tns.expect1(q, Pauli::Z).unwrap());
            assert!((after_bp - before[q as usize]).abs() < 1e-9, "q {q}: the state changed: {after_bp} vs {}", before[q as usize]);
            assert!((after_simple - after_bp).abs() < 1e-9, "q {q}: bond weights {after_simple} vs BP {after_bp}");
        }
    }

    #[test]
    fn gauged_evolution_is_exact_on_a_tree() {
        let tree = vec![(0, 1), (1, 2), (1, 3), (3, 4), (3, 5), (0, 6), (6, 7)];
        let mut rng = Rng(0x9a);
        for trial in 0..10 {
            let c = random_on(&mut rng, 8, &tree, 40);
            let mut tns = Tns::new(8, &tree, 16).unwrap();
            tns.run_gauged(&c, 100, 1e-14).unwrap();
            for q in 0..8 {
                let (got, want) = (tns.expect1(q, Pauli::Z).unwrap(), exact(&c, q, Pauli::Z));
                assert!((got - want).abs() < 1e-10, "trial {trial} q {q}: {got} vs {want}");
            }
        }
    }

    #[test]
    fn belief_propagation_is_bit_identical_however_it_is_parallelised() {
        let edges = heavy_hex_127();
        let mut tns = Tns::new(127, &edges, 4).unwrap();
        tns.run(&kicked_ising(127, &edges, 0.5, 3)).unwrap();
        let a = tns.belief_propagation(30, 1e-12);
        let b = tns.belief_propagation(30, 1e-12);
        assert_eq!((a.sweeps, a.residual.to_bits()), (b.sweeps, b.residual.to_bits()));
        assert_eq!(tns.expect1_bp(&a, 62, Pauli::Z).unwrap().to_bits(), tns.expect1_bp(&b, 62, Pauli::Z).unwrap().to_bits());
    }

    #[test]
    fn truncation_is_reported_and_bounded_by_chi() {
        let chain: Vec<(u32, u32)> = (0..9).map(|q| (q, q + 1)).collect();
        let mut rng = Rng(77);
        let c = random_on(&mut rng, 10, &chain, 200);
        let mut tns = Tns::new(10, &chain, 4).unwrap();
        tns.run(&c).unwrap();
        assert!(tns.max_bond() <= 4);
        assert!(tns.discarded() > 0.0, "a deep random circuit at χ = 4 must truncate");
    }

    #[test]
    fn a_gate_off_the_graph_is_refused() {
        let mut tns = Tns::new(3, &[(0, 1), (1, 2)], 4).unwrap();
        assert_eq!(tns.apply(&Op::CX(0, 2)), Err(TnsError::NotAnEdge(0, 2)));
        assert_eq!(
            tns.apply(&Op::Rot { axis: vec![(0, Pauli::Z), (1, Pauli::Z), (2, Pauli::Z)], theta: 0.1 }),
            Err(TnsError::TooWide(3))
        );
    }

    #[test]
    fn results_are_bit_identical_from_run_to_run() {
        let c = kicked_ising(127, &heavy_hex_127(), 0.6, 3);
        let run = || {
            let mut t = Tns::new(127, &heavy_hex_127(), 8).unwrap();
            t.run(&c).unwrap();
            (t.expect1(62, Pauli::Z).unwrap().to_bits(), t.discarded().to_bits())
        };
        assert_eq!(run(), run());
    }

    #[test]
    fn a_layer_computed_in_parallel_gives_the_bits_of_one_gate_at_a_time() {
        let edges = heavy_hex_127();
        let c = kicked_ising(127, &edges, 0.7, 3);
        let mut batched = Tns::new(127, &edges, 8).unwrap();
        batched.run(&c).unwrap();
        let mut serial = Tns::new(127, &edges, 8).unwrap();
        for op in &c.ops {
            serial.apply(op).unwrap();
        }
        assert_eq!(batched.discarded().to_bits(), serial.discarded().to_bits());
        for q in 0..127 {
            assert_eq!(batched.expect1(q, Pauli::Z).unwrap().to_bits(), serial.expect1(q, Pauli::Z).unwrap().to_bits());
        }
    }

    #[test]
    fn the_heavy_hex_magnetization_tracks_sparse_pauli_dynamics() {
        // On the 127-qubit lattice, with its loops, belief propagation is an
        // approximation. Five steps at small angle: the two independent
        // methods agree closely.
        let edges = heavy_hex_127();
        let c = kicked_ising(127, &edges, std::f64::consts::PI / 16.0, 5);
        let mut tns = Tns::new(127, &edges, 8).unwrap();
        tns.run(&c).unwrap();
        let spd = expect_z(&c, 62, &SpdConfig { threshold: 1e-7, max_weight: None }).unwrap().value;
        let bp = tns.expect1(62, Pauli::Z).unwrap();
        assert!((bp - spd).abs() < 2e-3, "BP {bp} vs SPD {spd}");
    }
}