wai-quantum 0.3.30

//! Variational quantum machine learning — interference on classical hardware
//! (`wai.quantum.vml`).
//!
//! A Boltzmann/thermodynamic sampler draws from a real, non-negative Gibbs
//! distribution `P(x) ∝ e^{-E(x)}`. What such a sampler cannot represent is
//! **interference**: a *complex* amplitude `ψ(x)` whose wrong answers
//! **destructively cancel**. That is the genuinely-quantum resource, and it runs on
//! classical hardware via a state-vector — exactly what this module does.
//!
//! Two capabilities, both deployable and both honest:
//! - **Quantum Circuit Born Machine** — a parameterized circuit as a generative
//!   model, `P_θ(x) = |⟨x|ψ_θ⟩|²`, trained by the **exact parameter-shift gradient**
//!   (no finite differences) to match a target distribution. Interference-based
//!   generation, learned on a laptop.
//! - **Barren-plateau probe** — the field's central *negative* result made
//!   measurable: the variance of a cost gradient vanishes ~exponentially in qubit
//!   count for deep random circuits. The tool that tells substance from hype.
//!
//! Determinism contract (like the MPS backend): continuous rotation angles are
//! irrational, so this is deterministic `f64` — reproducible via IEEE-754 strict
//! ops, not byte-exact. That is the correct, honest boundary for variational ML.

/// Complex amplitude (deterministic f64: only + − × ÷ √ cos sin).
#[derive(Clone, Copy)]
struct C {
    re: f64,
    im: f64,
}
impl C {
    const ZERO: C = C { re: 0.0, im: 0.0 };
    #[inline]
    fn add(self, o: C) -> C {
        C { re: self.re + o.re, im: self.im + o.im }
    }
    #[inline]
    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]
    fn scale(self, s: f64) -> C {
        C { re: self.re * s, im: self.im * s }
    }
    #[inline]
    fn norm2(self) -> f64 {
        self.re * self.re + self.im * self.im
    }
}

fn splitmix64(s: &mut u64) -> u64 {
    *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *s;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}
/// Deterministic uniform in `[-π, π)`.
fn rand_angle(s: &mut u64) -> f64 {
    let u = (splitmix64(s) >> 11) as f64 / (1u64 << 53) as f64;
    (u * 2.0 - 1.0) * std::f64::consts::PI
}

// ===========================================================================
// Parameterized state vector (RY, RZ, CX)
// ===========================================================================

#[inline]
fn apply_ry(amps: &mut [C], n: u8, q: u8, theta: f64) {
    let (s, c) = crate::repro::sin_cos(theta * 0.5);
    let bit = 1usize << q;
    let dim = 1usize << n;
    let mut i = 0;
    while i < dim {
        if i & bit == 0 {
            let j = i | bit;
            let (a, b) = (amps[i], amps[j]);
            amps[i] = a.scale(c).add(b.scale(-s));
            amps[j] = a.scale(s).add(b.scale(c));
        }
        i += 1;
    }
}
#[inline]
fn apply_rz(amps: &mut [C], q: u8, theta: f64) {
    let (sm, cm) = crate::repro::sin_cos(theta * 0.5);
    let e0 = C { re: cm, im: -sm }; // e^{-iθ/2}
    let e1 = C { re: cm, im: sm }; // e^{+iθ/2}
    let bit = 1usize << q;
    for (i, a) in amps.iter_mut().enumerate() {
        *a = a.mul(if i & bit == 0 { e0 } else { e1 });
    }
}
#[inline]
fn apply_cx(amps: &mut [C], n: u8, c: u8, t: u8) {
    let (cb, tb) = (1usize << c, 1usize << t);
    let dim = 1usize << n;
    for i in 0..dim {
        if (i & cb) != 0 && (i & tb) == 0 {
            amps.swap(i, i | tb);
        }
    }
}

/// The **hardware-efficient ansatz**: `layers` blocks of (RY, RZ on each qubit) +
/// a CX ring. `params.len() == layers * 2 * n`. The state whose Born distribution
/// is the generative model.
fn ansatz_state(n: u8, layers: u8, params: &[f64]) -> Vec<C> {
    let dim = 1usize << n;
    let mut amps = vec![C::ZERO; dim];
    amps[0] = C { re: 1.0, im: 0.0 };
    let mut p = 0;
    for _ in 0..layers {
        for q in 0..n {
            apply_ry(&mut amps, n, q, params[p]);
            p += 1;
            apply_rz(&mut amps, q, params[p]);
            p += 1;
        }
        for q in 0..n {
            apply_cx(&mut amps, n, q, (q + 1) % n);
        }
    }
    amps
}

/// Number of parameters for `(n, layers)`.
pub fn num_params(n: u8, layers: u8) -> usize {
    layers as usize * 2 * n as usize
}

/// The Born distribution `P_θ(x) = |ψ_θ(x)|²`.
pub fn born_probs(n: u8, layers: u8, params: &[f64]) -> Vec<f64> {
    ansatz_state(n, layers, params).iter().map(|a| a.norm2()).collect()
}

// ===========================================================================
// Born machine training  (exact parameter-shift gradient of KL)
// ===========================================================================

/// The outcome of training a Born machine.
pub struct BornFit {
    pub params: Vec<f64>,
    pub loss_curve: Vec<f64>,
    pub final_probs: Vec<f64>,
    pub final_kl: f64,
}

fn kl(target: &[f64], model: &[f64]) -> f64 {
    let mut s = 0.0;
    for (&q, &p) in target.iter().zip(model) {
        if q > 0.0 {
            s += q * crate::repro::ln(q / p.max(1e-12));
        }
    }
    s
}

/// Train a Born machine to match `target` (a distribution over `2ⁿ` outcomes) by
/// gradient descent on `KL(target ‖ P_θ)`, using the **exact parameter-shift**
/// gradient `∂P_θ(x)/∂θ = [P_{θ+π/2}(x) − P_{θ−π/2}(x)]/2`. Deterministic in `seed`.
pub fn train_born(target: &[f64], n: u8, layers: u8, steps: u32, lr: f64, seed: u64) -> BornFit {
    let np = num_params(n, layers);
    let mut st = seed.wrapping_mul(0xA24B_AED4).wrapping_add(1);
    let mut params: Vec<f64> = (0..np).map(|_| rand_angle(&mut st)).collect();
    let shift = std::f64::consts::FRAC_PI_2;
    let mut loss_curve = Vec::with_capacity(steps as usize + 1);
    for _ in 0..steps {
        let probs = born_probs(n, layers, &params);
        loss_curve.push(kl(target, &probs));
        // gradient of KL = -Σ_x q(x)/p(x) · ∂p(x)/∂θ_i
        let mut grad = vec![0.0; np];
        for i in 0..np {
            let saved = params[i];
            params[i] = saved + shift;
            let pp = born_probs(n, layers, &params);
            params[i] = saved - shift;
            let pm = born_probs(n, layers, &params);
            params[i] = saved;
            let mut g = 0.0;
            for x in 0..target.len() {
                if target[x] > 0.0 {
                    let dp = 0.5 * (pp[x] - pm[x]);
                    g -= target[x] / probs[x].max(1e-12) * dp;
                }
            }
            grad[i] = g;
        }
        for i in 0..np {
            params[i] -= lr * grad[i];
        }
    }
    let final_probs = born_probs(n, layers, &params);
    let final_kl = kl(target, &final_probs);
    loss_curve.push(final_kl);
    BornFit { params, loss_curve, final_probs, final_kl }
}

// ===========================================================================
// Barren-plateau probe  (the honest limit)
// ===========================================================================

/// `⟨Z_0⟩` for a random ansatz — the cost whose gradient we probe.
fn cost_z0(n: u8, layers: u8, params: &[f64]) -> f64 {
    let probs = born_probs(n, layers, params);
    let mut acc = 0.0;
    for (i, p) in probs.iter().enumerate() {
        acc += if i & 1 == 0 { *p } else { -*p };
    }
    acc
}

/// The **barren-plateau** signature: the variance of `∂⟨Z_0⟩/∂θ_0` over `samples`
/// random parameter sets, at `(n, layers)`. It falls ~exponentially in `n` for deep
/// circuits — the reason variational QML gradients vanish, and the honest bound on
/// when it can ever beat classical. Deterministic in `seed`.
pub fn barren_plateau_variance(n: u8, layers: u8, samples: u32, seed: u64) -> f64 {
    let np = num_params(n, layers);
    let mut st = seed.wrapping_mul(0x2545_F491).wrapping_add(7);
    let shift = std::f64::consts::FRAC_PI_2;
    let mut grads = Vec::with_capacity(samples as usize);
    for _ in 0..samples {
        let mut params: Vec<f64> = (0..np).map(|_| rand_angle(&mut st)).collect();
        let saved = params[0];
        params[0] = saved + shift;
        let cp = cost_z0(n, layers, &params);
        params[0] = saved - shift;
        let cm = cost_z0(n, layers, &params);
        grads.push(0.5 * (cp - cm));
    }
    let mean = grads.iter().sum::<f64>() / grads.len() as f64;
    grads.iter().map(|g| (g - mean) * (g - mean)).sum::<f64>() / grads.len() as f64
}

/// A `2ⁿ`-outcome target that only a circuit with **interference** can match well:
/// the parity distribution — uniform over even-parity strings, zero on odd — which
/// requires amplitudes to cancel on the odd sector.
pub fn parity_target(n: u8) -> Vec<f64> {
    let dim = 1usize << n;
    let even: Vec<usize> = (0..dim).filter(|x| (x.count_ones() % 2) == 0).collect();
    let w = 1.0 / even.len() as f64;
    let mut t = vec![0.0; dim];
    for x in even {
        t[x] = w;
    }
    t
}

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

    #[test]
    fn born_probs_normalized() {
        let n = 4;
        let l = 3;
        let mut st = 1u64;
        let params: Vec<f64> = (0..num_params(n, l)).map(|_| rand_angle(&mut st)).collect();
        let p = born_probs(n, l, &params);
        assert!((p.iter().sum::<f64>() - 1.0).abs() < 1e-9, "Born probs must sum to 1");
    }

    #[test]
    fn born_machine_learns_a_bell_target() {
        // target: 50% |00>, 50% |11> — a distribution requiring correlation.
        let target = vec![0.5, 0.0, 0.0, 0.5];
        let fit = train_born(&target, 2, 3, 120, 0.3, 7);
        assert!(fit.loss_curve[0] > fit.final_kl, "training must reduce KL");
        assert!(fit.final_kl < 0.05, "should fit the target well: KL={}", fit.final_kl);
        // mass concentrates on |00> and |11>
        assert!(fit.final_probs[0] + fit.final_probs[3] > 0.9, "mass on the correlated pair");
    }

    #[test]
    fn born_machine_learns_the_parity_target() {
        // the interference target: even-parity uniform, zero on odd.
        let n = 3;
        let target = parity_target(n);
        let fit = train_born(&target, n, 4, 200, 0.2, 3);
        assert!(fit.final_kl < 0.15, "parity target learned: KL={}", fit.final_kl);
        // odd-parity strings suppressed (destructive interference)
        let odd_mass: f64 = (0..(1 << n)).filter(|x: &usize| x.count_ones() % 2 == 1).map(|x| fit.final_probs[x]).sum();
        assert!(odd_mass < 0.2, "odd sector suppressed by interference: {odd_mass}");
    }

    #[test]
    fn barren_plateau_variance_decays_with_qubits() {
        // the honest limit: deeper/wider random circuits → exponentially smaller
        // gradients. Variance should fall as n grows.
        let v_small = barren_plateau_variance(2, 6, 400, 1);
        let v_mid = barren_plateau_variance(4, 6, 400, 1);
        let v_big = barren_plateau_variance(6, 6, 400, 1);
        assert!(v_mid < v_small, "gradient variance shrinks with n: {v_mid} < {v_small}");
        assert!(v_big < v_mid, "…and keeps shrinking: {v_big} < {v_mid}");
        // roughly geometric decay is the barren-plateau signature
        assert!(v_big < v_small * 0.4, "clear exponential-ish decay");
    }

    #[test]
    fn deterministic() {
        let a = train_born(&[0.5, 0.0, 0.0, 0.5], 2, 2, 30, 0.3, 42).final_kl;
        let b = train_born(&[0.5, 0.0, 0.0, 0.5], 2, 2, 30, 0.3, 42).final_kl;
        assert_eq!(a.to_bits(), b.to_bits(), "training must be bit-reproducible");
    }
}