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

//! Dense complex linear algebra shared by the tensor-network backends: a
//! complex number, a one-sided Jacobi SVD and a Householder QR. Only
//! `+ − × ÷ √`, in a fixed order, so every result is the same on every machine.

#![allow(clippy::should_implement_trait, clippy::needless_range_loop)]

// ===========================================================================
// Complex f64 (deterministic: only + - * / sqrt, no FMA/transcendentals)
// ===========================================================================

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct C {
    pub re: f64,
    pub im: f64,
}
// The tensor-network layer uses only the type and its constants; the
// arithmetic helpers serve the decomposition backends.
#[cfg_attr(not(any(feature = "quantum_mps", feature = "quantum_bptn")), allow(dead_code))]
impl C {
    pub(crate) const ZERO: C = C { re: 0.0, im: 0.0 };
    pub(crate) const ONE: C = C { re: 1.0, im: 0.0 };
    #[inline]
    pub(crate) fn new(re: f64, im: f64) -> C {
        C { re, im }
    }
    #[inline]
    pub(crate) fn add(self, o: C) -> C {
        C { re: self.re + o.re, im: self.im + o.im }
    }
    #[inline]
    pub(crate) fn sub(self, o: C) -> C {
        C { re: self.re - o.re, im: self.im - o.im }
    }
    #[inline]
    pub(crate) fn mul(self, o: C) -> C {
        C { re: self.re * o.re - self.im * o.im, im: self.re * o.im + self.im * o.re }
    }
    #[inline]
    pub(crate) fn conj(self) -> C {
        C { re: self.re, im: -self.im }
    }
    #[inline]
    pub(crate) fn scale(self, s: f64) -> C {
        C { re: self.re * s, im: self.im * s }
    }
    #[inline]
    pub(crate) fn norm2(self) -> f64 {
        self.re * self.re + self.im * self.im
    }
}

// ===========================================================================
// Complex one-sided Jacobi SVD  (M = U · diag(S) · V†)
// ===========================================================================

/// SVD of an `rows × cols` matrix given by its **columns** (`m[c]` is column `c`, a
/// length-`rows` vector). Returns `(u_cols, s, v_cols)` where each `u_cols[k]` is a
/// length-`rows` left singular vector, `s[k]` its singular value, `v_cols[k]` a
/// length-`cols` right singular vector — sorted by `s` descending. Deterministic:
/// fixed sweep order, fixed sweep cap, only `+ − × ÷ √`.
#[cfg(feature = "quantum_mps")]
pub(crate) fn jacobi_svd(mut m: Vec<Vec<C>>, rows: usize, cols: usize) -> (Vec<Vec<C>>, Vec<f64>, Vec<Vec<C>>) {
    // V starts as identity (cols × cols), stored by columns.
    let mut v: Vec<Vec<C>> = (0..cols)
        .map(|j| (0..cols).map(|i| if i == j { C::ONE } else { C::ZERO }).collect())
        .collect();
    let eps = 1e-14;
    for _sweep in 0..60 {
        let mut off = 0.0f64;
        for p in 0..cols {
            for q in (p + 1)..cols {
                // gram entries of columns p,q
                let mut alpha = 0.0; // <c_p,c_p>
                let mut beta = 0.0; // <c_q,c_q>
                let mut gamma = C::ZERO; // <c_p,c_q> = sum conj(c_p)·c_q
                for k in 0..rows {
                    alpha += m[p][k].norm2();
                    beta += m[q][k].norm2();
                    gamma = gamma.add(m[p][k].conj().mul(m[q][k]));
                }
                let g2 = gamma.norm2();
                off += g2;
                if g2 <= eps * alpha * beta || g2 == 0.0 {
                    continue;
                }
                let gabs = g2.sqrt();
                // phase-rotate column q so <c_p,c_q> becomes real +gabs
                let ph = C::new(gamma.re / gabs, gamma.im / gabs); // unit
                let cph = ph.conj();
                for k in 0..rows {
                    m[q][k] = m[q][k].mul(cph);
                }
                for k in 0..cols {
                    v[q][k] = v[q][k].mul(cph);
                }
                // real Jacobi rotation with a=alpha,b=beta,g=gabs
                let tau = (beta - alpha) / (2.0 * gabs);
                let t = if tau >= 0.0 {
                    1.0 / (tau + (1.0 + tau * tau).sqrt())
                } else {
                    -1.0 / (-tau + (1.0 + tau * tau).sqrt())
                };
                let cs = 1.0 / (1.0 + t * t).sqrt();
                let sn = t * cs;
                for k in 0..rows {
                    let x = m[p][k];
                    let y = m[q][k];
                    m[p][k] = x.scale(cs).sub(y.scale(sn));
                    m[q][k] = x.scale(sn).add(y.scale(cs));
                }
                for k in 0..cols {
                    let x = v[p][k];
                    let y = v[q][k];
                    v[p][k] = x.scale(cs).sub(y.scale(sn));
                    v[q][k] = x.scale(sn).add(y.scale(cs));
                }
            }
        }
        if off <= eps {
            break;
        }
    }
    // singular values = column norms; left vectors = normalized columns
    let mut s = vec![0.0; cols];
    for j in 0..cols {
        s[j] = (0..rows).map(|k| m[j][k].norm2()).sum::<f64>().sqrt();
    }
    let mut u: Vec<Vec<C>> = m;
    for j in 0..cols {
        if s[j] > 1e-300 {
            let inv = 1.0 / s[j];
            for k in 0..rows {
                u[j][k] = u[j][k].scale(inv);
            }
        }
    }
    // sort by singular value descending (deterministic)
    let mut order: Vec<usize> = (0..cols).collect();
    order.sort_by(|&a, &b| s[b].partial_cmp(&s[a]).unwrap_or(std::cmp::Ordering::Equal));
    let s2: Vec<f64> = order.iter().map(|&i| s[i]).collect();
    let u2: Vec<Vec<C>> = order.iter().map(|&i| u[i].clone()).collect();
    let v2: Vec<Vec<C>> = order.iter().map(|&i| v[i].clone()).collect();
    (u2, s2, v2)
}

/// [`jacobi_svd`] run to convergence: a pair of columns is rotated until
/// their correlation is at rounding level, and the sweeps stop when one
/// rotates nothing. The original keeps its looser stopping rule (columns
/// orthogonal to about 1e-7 relative) because the MPS backend's results are
/// pinned to it.
#[cfg(any(feature = "quantum_bptn", feature = "quantum_tdvp"))]
pub(crate) fn jacobi_svd_strict(mut m: Vec<Vec<C>>, rows: usize, cols: usize) -> (Vec<Vec<C>>, Vec<f64>, Vec<Vec<C>>) {
    // V starts as identity (cols × cols), stored by columns.
    let mut v: Vec<Vec<C>> = (0..cols)
        .map(|j| (0..cols).map(|i| if i == j { C::ONE } else { C::ZERO }).collect())
        .collect();
    // |⟨c_p,c_q⟩|² ≤ eps·‖c_p‖²‖c_q‖²: columns orthogonal to rounding.
    let eps = 1e-30;
    // A column this small against the whole matrix is numerically null: it
    // holds no singular value anyone keeps, and rotating it only feeds
    // rounding noise into V.
    let null = 1e-30 * m.iter().flatten().map(|x| x.norm2()).sum::<f64>();
    for _sweep in 0..100 {
        let mut rotated = false;
        for p in 0..cols {
            for q in (p + 1)..cols {
                // gram entries of columns p,q
                let mut alpha = 0.0; // <c_p,c_p>
                let mut beta = 0.0; // <c_q,c_q>
                let mut gamma = C::ZERO; // <c_p,c_q> = sum conj(c_p)·c_q
                for k in 0..rows {
                    alpha += m[p][k].norm2();
                    beta += m[q][k].norm2();
                    gamma = gamma.add(m[p][k].conj().mul(m[q][k]));
                }
                if alpha <= null || beta <= null {
                    continue;
                }
                let g2 = gamma.norm2();
                if g2 <= eps * alpha * beta || g2 == 0.0 {
                    continue;
                }
                rotated = true;
                // |γ| scaled by its larger component, so a tiny γ keeps its
                // phase at unit modulus without underflowing. Not `hypot`:
                // libm's last bit differs between platforms; `/` and `√` do not.
                let big = gamma.re.abs().max(gamma.im.abs());
                let (r, i) = (gamma.re / big, gamma.im / big);
                let gabs = big * (r * r + i * i).sqrt();
                // phase-rotate column q so <c_p,c_q> becomes real +gabs
                let ph = C::new(gamma.re / gabs, gamma.im / gabs); // unit
                let cph = ph.conj();
                for k in 0..rows {
                    m[q][k] = m[q][k].mul(cph);
                }
                for k in 0..cols {
                    v[q][k] = v[q][k].mul(cph);
                }
                // real Jacobi rotation with a=alpha,b=beta,g=gabs
                let tau = (beta - alpha) / (2.0 * gabs);
                let t = if tau >= 0.0 {
                    1.0 / (tau + (1.0 + tau * tau).sqrt())
                } else {
                    -1.0 / (-tau + (1.0 + tau * tau).sqrt())
                };
                let cs = 1.0 / (1.0 + t * t).sqrt();
                let sn = t * cs;
                for k in 0..rows {
                    let x = m[p][k];
                    let y = m[q][k];
                    m[p][k] = x.scale(cs).sub(y.scale(sn));
                    m[q][k] = x.scale(sn).add(y.scale(cs));
                }
                for k in 0..cols {
                    let x = v[p][k];
                    let y = v[q][k];
                    v[p][k] = x.scale(cs).sub(y.scale(sn));
                    v[q][k] = x.scale(sn).add(y.scale(cs));
                }
            }
        }
        if !rotated {
            break;
        }
    }
    // singular values = column norms; left vectors = normalized columns
    let mut s = vec![0.0; cols];
    for j in 0..cols {
        s[j] = (0..rows).map(|k| m[j][k].norm2()).sum::<f64>().sqrt();
    }
    let mut u: Vec<Vec<C>> = m;
    for j in 0..cols {
        if s[j] > 1e-300 {
            let inv = 1.0 / s[j];
            for k in 0..rows {
                u[j][k] = u[j][k].scale(inv);
            }
        }
    }
    // sort by singular value descending (deterministic)
    let mut order: Vec<usize> = (0..cols).collect();
    order.sort_by(|&a, &b| s[b].partial_cmp(&s[a]).unwrap_or(std::cmp::Ordering::Equal));
    let s2: Vec<f64> = order.iter().map(|&i| s[i]).collect();
    let u2: Vec<Vec<C>> = order.iter().map(|&i| u[i].clone()).collect();
    let v2: Vec<Vec<C>> = order.iter().map(|&i| v[i].clone()).collect();
    (u2, s2, v2)
}


// ===========================================================================
// Householder QR  (A = Q · R)
// ===========================================================================

/// Thin QR of the row-major `m × n` matrix `a`: `Q` (`m × k`, orthonormal
/// columns) and `R` (`k × n`, upper triangular), both row-major, with
/// `k = min(m, n)`. Householder reflections in a fixed order; deterministic.
#[cfg(feature = "quantum_bptn")]
pub(crate) fn qr(a: &[C], m: usize, n: usize) -> (Vec<C>, Vec<C>, usize) {
    let k = m.min(n);
    let mut r = a.to_vec();
    let mut vs: Vec<Vec<C>> = Vec::with_capacity(k);
    for j in 0..k {
        let norm = (j..m).map(|i| r[i * n + j].norm2()).sum::<f64>().sqrt();
        let mut v: Vec<C> = (j..m).map(|i| r[i * n + j]).collect();
        if norm == 0.0 {
            vs.push(vec![C::ZERO; m - j]);
            continue;
        }
        // alpha = −(x₀/|x₀|)·‖x‖, so the reflection maps x to alpha·e₁
        let x0 = v[0];
        let a0 = x0.norm2().sqrt();
        let phase = if a0 == 0.0 { C::ONE } else { C::new(x0.re / a0, x0.im / a0) };
        let alpha = phase.scale(-norm);
        v[0] = v[0].sub(alpha);
        let vn = v.iter().map(|x| x.norm2()).sum::<f64>().sqrt();
        if vn == 0.0 {
            vs.push(vec![C::ZERO; m - j]);
            continue;
        }
        for x in v.iter_mut() {
            *x = x.scale(1.0 / vn);
        }
        // r[j.., j..] -= 2 v (v† r[j.., j..])
        for c in j..n {
            let mut dot = C::ZERO;
            for (t, vi) in v.iter().enumerate() {
                dot = dot.add(vi.conj().mul(r[(j + t) * n + c]));
            }
            let d2 = dot.scale(2.0);
            for (t, vi) in v.iter().enumerate() {
                r[(j + t) * n + c] = r[(j + t) * n + c].sub(vi.mul(d2));
            }
        }
        vs.push(v);
    }
    // Q = H₀ H₁ … H_{k−1} applied to the first k columns of the identity.
    let mut q = vec![C::ZERO; m * k];
    for i in 0..k {
        q[i * k + i] = C::ONE;
    }
    for j in (0..k).rev() {
        let v = &vs[j];
        for c in 0..k {
            let mut dot = C::ZERO;
            for (t, vi) in v.iter().enumerate() {
                dot = dot.add(vi.conj().mul(q[(j + t) * k + c]));
            }
            let d2 = dot.scale(2.0);
            for (t, vi) in v.iter().enumerate() {
                q[(j + t) * k + c] = q[(j + t) * k + c].sub(vi.mul(d2));
            }
        }
    }
    let mut rr = vec![C::ZERO; k * n];
    for i in 0..k {
        for c in i..n {
            rr[i * n + c] = r[i * n + c];
        }
    }
    (q, rr, k)
}

#[cfg(all(test, feature = "quantum_bptn"))]
mod tests {
    use super::*;

    #[test]
    fn qr_reconstructs_and_q_is_orthonormal() {
        let mut s = 7u64;
        let mut next = || {
            s ^= s << 13;
            s ^= s >> 7;
            s ^= s << 17;
            (s >> 11) as f64 / (1u64 << 53) as f64 - 0.5
        };
        for (m, n) in [(5, 3), (3, 5), (8, 8), (1, 4), (6, 1)] {
            let a: Vec<C> = (0..m * n).map(|_| C::new(next(), next())).collect();
            let (q, r, k) = qr(&a, m, n);
            for i in 0..m {
                for c in 0..n {
                    let mut x = C::ZERO;
                    for t in 0..k {
                        x = x.add(q[i * k + t].mul(r[t * n + c]));
                    }
                    let d = x.sub(a[i * n + c]);
                    assert!(d.norm2() < 1e-26, "({m}x{n}) entry {i},{c}");
                }
            }
            for a1 in 0..k {
                for b1 in 0..k {
                    let mut x = C::ZERO;
                    for i in 0..m {
                        x = x.add(q[i * k + a1].conj().mul(q[i * k + b1]));
                    }
                    let want = if a1 == b1 { 1.0 } else { 0.0 };
                    assert!((x.re - want).abs() < 1e-13 && x.im.abs() < 1e-13);
                }
            }
            for i in 0..k {
                for c in 0..i.min(n) {
                    assert_eq!(r[i * n + c].norm2(), 0.0, "R is upper triangular");
                }
            }
        }
    }
}

#[cfg(all(test, feature = "quantum_bptn"))]
mod svd_tests {
    use super::*;

    /// The converged SVD reconstructs `A = U·S·V†` and returns orthonormal
    /// singular vectors, both to rounding, on full-rank and rank-deficient
    /// matrices. (The original's looser stopping rule leaves `U` orthonormal
    /// only to about 1e-7.)
    #[test]
    fn the_converged_svd_reconstructs_with_orthonormal_vectors() {
        let mut s = 99u64;
        let mut next = || {
            s ^= s << 13;
            s ^= s >> 7;
            s ^= s << 17;
            (s >> 11) as f64 / (1u64 << 53) as f64 - 0.5
        };
        for trial in 0..200 {
            let rows = 2 + trial % 9;
            let cols = 2 + (trial / 9) % 7;
            let rank = 1 + trial % rows.min(cols);
            let b: Vec<C> = (0..rows * rank).map(|_| C::new(next(), next())).collect();
            let c: Vec<C> = (0..cols * rank).map(|_| C::new(next(), next())).collect();
            let a: Vec<C> = (0..rows * cols)
                .map(|i| {
                    let (r, q) = (i / cols, i % cols);
                    (0..rank).fold(C::ZERO, |acc, k| acc.add(b[r * rank + k].mul(c[q * rank + k])))
                })
                .collect();
            let columns: Vec<Vec<C>> = (0..cols).map(|q| (0..rows).map(|r| a[r * cols + q]).collect()).collect();
            let (u, sv, v) = jacobi_svd_strict(columns, rows, cols);
            for r in 0..rows {
                for q in 0..cols {
                    let x = (0..cols).fold(C::ZERO, |acc, k| acc.add(u[k][r].scale(sv[k]).mul(v[k][q].conj())));
                    assert!(x.sub(a[r * cols + q]).norm2() < 1e-26, "trial {trial}: reconstruction");
                }
            }
            for k1 in 0..cols {
                for k2 in 0..cols {
                    for (left, vecs, len) in [(true, &u, rows), (false, &v, cols)] {
                        if left && (sv[k1] < 1e-8 * sv[0] || sv[k2] < 1e-8 * sv[0]) {
                            continue; // a null column of U is not normalised
                        }
                        let d = (0..len).fold(C::ZERO, |acc, r| acc.add(vecs[k1][r].conj().mul(vecs[k2][r])));
                        let want = if k1 == k2 { 1.0 } else { 0.0 };
                        assert!((d.re - want).abs() < 1e-13 && d.im.abs() < 1e-13, "trial {trial}: orthonormality");
                    }
                }
            }
        }
    }
}