wai-quantum 0.3.20

//! Deterministic NOISE LEARNING — `wai.quantum.cal.noise_learn`
//! (extensions/quantum-ops § Calibration, noise model).
//!
//! Error mitigation's highest-fidelity methods (probabilistic error cancellation /
//! amplification) do not run on a fidelity number — they run on a **learned noise
//! model**: the sparse Pauli–Lindblad model, whose rates are estimated by **Cycle
//! Benchmarking** and **Cycle Error Reconstruction**. Twirling a repeated gate
//! cycle turns its noise into a Pauli channel; the decay of each Pauli's
//! expectation over cycle depth measures that Pauli's **fidelity** `f_a`; and the
//! model's rates `λ_k` are recovered from the fidelities by a linear inversion,
//! because for the sparse Pauli–Lindblad model
//!
//! ```text
//!   f_a = exp(-2 · Σ_{k : a anticommutes P_k} λ_k)   ⟺   -½ ln f_a = (M λ)_a
//! ```
//!
//! with `M[a][k] = 1` iff probe Pauli `a` anticommutes generator `P_k`.
//!
//! This module learns that model the WAI way: pure fixed-point integer numerics,
//! byte-identical on every machine, recovering a *hidden* device's rates from a
//! deterministically-simulated Cycle-Benchmarking experiment, then sealing a
//! signed `wai.quantum.calibration` receipt. Its content hash is exactly the
//! `noise_model_hash` a [`crate::quantum_ops::MitigationReceipt`] binds — so this
//! is the **producer** for the PEC hook the mitigation layer already exposes:
//! learn the model, sign it, and a probabilistic-error-cancellation of a Pauli
//! observable (`⟨P⟩_ideal = ⟨P⟩_noisy / f_P`) binds the exact model it used.
//!
//! Honest boundary: a Pauli-channel noise model on a simulated cycle (no QPU) —
//! the twirled front-end is exactly the surface CB/CER runs on. What is not
//! approximate is the reproducibility of the reconstruction and the receipts.

use crate::quantum_cal::{exp_neg_fx, seal_artifacts, CalArtifacts, CAL_FRAC, CAL_ONE};
use crate::quantum_ops::{content_hash, CalibrationReceipt, GrantRef, MitigationReceipt};
use ed25519_dalek::SigningKey;

const ONE: i64 = CAL_ONE;
const FRAC: u32 = CAL_FRAC;

#[inline]
fn fmul(a: i64, b: i64) -> i64 {
    ((a as i128 * b as i128) >> FRAC) as i64
}
#[inline]
fn fdiv(a: i64, b: i64) -> i64 {
    if b == 0 {
        return 0;
    }
    (((a as i128) << FRAC) / b as i128) as i64
}

/// Fixed-point `ln(x)` for `x ∈ (0, ONE]` via `2·artanh((x−1)/(x+1))` (≤ 0).
fn ln_fx(x: i64) -> i64 {
    if x <= 0 {
        return -20 * ONE;
    }
    let z = fdiv(x - ONE, x + ONE);
    let z2 = fmul(z, z);
    let mut term = z;
    let mut acc = 0i64;
    let mut k = 1i64;
    for _ in 0..12 {
        acc += term / k;
        term = fmul(term, z2);
        k += 2;
    }
    2 * acc
}

fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

// ===========================================================================
// Pauli operators (symplectic representation)
// ===========================================================================

/// An `n`-qubit Pauli (up to global phase) as symplectic bit-vectors: bit `i` of
/// `x` marks an X on qubit `i`, bit `i` of `z` marks a Z; `Y = XZ` sets both.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct Pauli {
    pub x: u32,
    pub z: u32,
}

impl Pauli {
    pub const fn new(x: u32, z: u32) -> Pauli {
        Pauli { x, z }
    }
    /// Do `self` and `other` anticommute? (Symplectic inner product = 1.)
    pub fn anticommutes(&self, other: &Pauli) -> bool {
        ((self.x & other.z).count_ones() + (self.z & other.x).count_ones()) & 1 == 1
    }
    pub fn weight(&self) -> u32 {
        (self.x | self.z).count_ones()
    }
    fn code(&self) -> u64 {
        ((self.x as u64) << 32) | self.z as u64
    }
}

/// All non-identity Paulis on `n` qubits (`4^n − 1` of them) — the probe set.
pub fn all_paulis(n: u32) -> Vec<Pauli> {
    let mut v = Vec::new();
    let dim = 1u32 << n;
    for x in 0..dim {
        for z in 0..dim {
            if x != 0 || z != 0 {
                v.push(Pauli::new(x, z));
            }
        }
    }
    v
}

/// The sparse Pauli–Lindblad generator set: every weight-1 Pauli (`X,Y,Z` on each
/// qubit) plus every weight-2 Pauli supported on an edge of `connectivity`. This
/// is the "mostly local" generator basis the model is built on.
pub fn sparse_generators(n: u32, edges: &[(u32, u32)]) -> Vec<Pauli> {
    let single = [(1u32, 0u32), (0, 1), (1, 1)]; // X, Z, Y  (as (x,z) on one qubit)
    let mut g = Vec::new();
    for q in 0..n {
        for &(px, pz) in &single {
            g.push(Pauli::new(px << q, pz << q));
        }
    }
    for &(i, j) in edges {
        for &(ax, az) in &single {
            for &(bx, bz) in &single {
                g.push(Pauli::new((ax << i) | (bx << j), (az << i) | (bz << j)));
            }
        }
    }
    g
}

// ===========================================================================
// Sparse Pauli–Lindblad noise model
// ===========================================================================

/// A sparse Pauli–Lindblad channel: rate `lambdas[k] ≥ 0` on generator
/// `generators[k]`. Deterministic and content-addressable.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct NoiseModel {
    pub generators: Vec<Pauli>,
    pub lambdas: Vec<i64>,
}

impl NoiseModel {
    /// The Pauli fidelity `f_a = exp(-2 Σ_{k: a anti P_k} λ_k)` for probe `a`.
    pub fn fidelity(&self, a: &Pauli) -> i64 {
        let mut s = 0i64;
        for (k, g) in self.generators.iter().enumerate() {
            if a.anticommutes(g) {
                s += self.lambdas[k];
            }
        }
        exp_neg_fx(2 * s)
    }

    /// Canonical bytes for content-hashing — the pinned model a mitigation binds.
    pub fn bytes(&self) -> Vec<u8> {
        let mut b = Vec::with_capacity(self.generators.len() * 16 + 8);
        b.extend_from_slice(b"wai:qc-noise-model\x01");
        // sort by generator code for a canonical order
        let mut idx: Vec<usize> = (0..self.generators.len()).collect();
        idx.sort_by_key(|&i| self.generators[i].code());
        for &i in &idx {
            b.extend_from_slice(&self.generators[i].code().to_le_bytes());
            b.extend_from_slice(&self.lambdas[i].to_le_bytes());
        }
        b
    }
    pub fn model_hash(&self) -> [u8; 32] {
        content_hash(&self.bytes())
    }
    /// Total error rate Σ λ_k (a scalar summary of noise strength).
    pub fn total_rate(&self) -> i64 {
        self.lambdas.iter().sum()
    }
}

// ===========================================================================
// Cycle Benchmarking simulation + reconstruction (CER)
// ===========================================================================

/// The learned model plus the hidden truth it was recovered from.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LearnResult {
    pub learned: NoiseModel,
    pub truth: NoiseModel,
    /// The measured Pauli fidelities per probe (from the CB decays).
    pub probes: Vec<Pauli>,
    pub fidelity_meas: Vec<i64>,
    /// Mean absolute error between learned and true rates (fixed-point).
    pub rate_error_fx: i64,
    /// Cycle depths used in the benchmarking decays.
    pub depths: Vec<u32>,
}

/// Measure one probe Pauli's fidelity by a Cycle-Benchmarking decay: the
/// expectation at depth `m` is `f^m`; we sample it with `shots` and recover `f`
/// from a log-linear fit of `ln E_m` vs `m` (slope = `ln f`). Deterministic.
fn measure_fidelity(f_true: i64, depths: &[u32], shots: u32, seed: u64) -> i64 {
    // log-linear least squares: slope = Σ(m-m̄)(y-ȳ) / Σ(m-m̄)²,  y = ln E_m
    let mut xs = Vec::new();
    let mut ys = Vec::new();
    for (di, &m) in depths.iter().enumerate() {
        // E_m = f_true^m
        let mut e = ONE;
        for _ in 0..m {
            e = fmul(e, f_true);
        }
        // sample P(+1) = (1+E)/2
        let mut st = seed.wrapping_mul(0x100_0001).wrapping_add(di as u64 + 1);
        let p_plus = (ONE + e) / 2;
        let mut hits = 0u64;
        for _ in 0..shots {
            let r = (splitmix64(&mut st) >> (64 - FRAC)) as i64;
            if r < p_plus {
                hits += 1;
            }
        }
        let e_hat = (2 * (hits as i128 * ONE as i128) / shots.max(1) as i128 - ONE as i128) as i64;
        let e_hat = e_hat.clamp(ONE / 1000, ONE);
        xs.push(m as i64 * ONE);
        ys.push(ln_fx(e_hat)); // ≤ 0
    }
    let n = xs.len() as i128;
    let sx: i128 = xs.iter().map(|&v| v as i128).sum();
    let sy: i128 = ys.iter().map(|&v| v as i128).sum();
    let mut sxx = 0i128;
    let mut sxy = 0i128;
    for i in 0..xs.len() {
        sxx += xs[i] as i128 * xs[i] as i128;
        sxy += xs[i] as i128 * ys[i] as i128;
    }
    let denom = n * sxx - sx * sx;
    if denom == 0 {
        return ONE;
    }
    let slope = (((n * sxy - sx * sy) << FRAC) / denom) as i64; // ln f (≤0)
    // f = exp(slope) = exp_neg(-slope)
    exp_neg_fx(-slope).clamp(0, ONE)
}

/// Learn a sparse Pauli–Lindblad model from a Cycle-Benchmarking experiment on a
/// hidden `truth`, using generator set `generators` and probe set `probes`. The
/// rates are recovered by positivity-preserving multiplicative NNLS on the linear
/// system `M λ = y`, `y_a = -½ ln f_a`, `M[a][k] = [a anti P_k]`.
pub fn learn_noise_model(
    truth: &NoiseModel,
    generators: &[Pauli],
    probes: &[Pauli],
    depths: &[u32],
    shots: u32,
    iters: u32,
    seed: u64,
) -> LearnResult {
    // Anticommutation matrix M (probes × generators), stored row-major as bools.
    let ng = generators.len();
    let np = probes.len();
    let mut m = vec![false; np * ng];
    for (ai, a) in probes.iter().enumerate() {
        for (k, g) in generators.iter().enumerate() {
            m[ai * ng + k] = a.anticommutes(g);
        }
    }

    // "Measured" targets y_a = -½ ln f_a from the CB decays.
    let mut y = Vec::with_capacity(np);
    let mut fmeas = Vec::with_capacity(np);
    for (ai, a) in probes.iter().enumerate() {
        let f_true = truth.fidelity(a);
        let f_hat = measure_fidelity(f_true, depths, shots, seed.wrapping_add(ai as u64 * 0x9E37));
        fmeas.push(f_hat);
        y.push((-ln_fx(f_hat)) / 2); // ≥ 0
    }

    // Mᵀ y (per generator) — constant across iterations.
    let mut mty = vec![0i64; ng];
    for ai in 0..np {
        for k in 0..ng {
            if m[ai * ng + k] {
                mty[k] += y[ai];
            }
        }
    }

    // Multiplicative NNLS: λ_k ← λ_k · (Mᵀy)_k / (Mᵀ M λ)_k
    let mut lam = vec![ONE / 100; ng]; // small positive start
    for _ in 0..iters {
        // predicted (M λ)_a
        let mut ml = vec![0i64; np];
        for ai in 0..np {
            let mut s = 0i64;
            for k in 0..ng {
                if m[ai * ng + k] {
                    s += lam[k];
                }
            }
            ml[ai] = s;
        }
        // MᵀMλ per generator
        for k in 0..ng {
            let mut denom = 0i64;
            for ai in 0..np {
                if m[ai * ng + k] {
                    denom += ml[ai];
                }
            }
            if denom > 0 {
                lam[k] = fmul(lam[k], fdiv(mty[k], denom));
            }
        }
    }

    let learned = NoiseModel { generators: generators.to_vec(), lambdas: lam };
    let rate_error_fx = {
        let s: i64 = truth
            .lambdas
            .iter()
            .zip(&learned.lambdas)
            .map(|(&a, &b)| (a - b).abs())
            .sum();
        s / ng.max(1) as i64
    };

    LearnResult {
        learned,
        truth: truth.clone(),
        probes: probes.to_vec(),
        fidelity_meas: fmeas,
        rate_error_fx,
        depths: depths.to_vec(),
    }
}

impl LearnResult {
    pub fn hash(&self) -> [u8; 32] {
        self.learned.model_hash()
    }
    pub fn artifacts(&self) -> CalArtifacts {
        CalArtifacts {
            target: "noise-model:cycle".into(),
            config_bytes: self.learned.bytes(),
            evidence_kind: "cycle_error_reconstruction".into(),
            evidence_bytes: {
                let mut ev = Vec::new();
                for &f in &self.fidelity_meas {
                    ev.extend_from_slice(&f.to_le_bytes());
                }
                ev
            },
            summary: format!(
                "sparse Pauli–Lindblad model, {} generators, total rate {}",
                self.learned.generators.len(),
                self.learned.total_rate()
            ),
        }
    }
    /// Seal the learned model as a signed calibration receipt.
    pub fn seal(&self, signer: &SigningKey, signer_id: &str, joules_micro: u64, grant: GrantRef) -> CalibrationReceipt {
        seal_artifacts(signer, signer_id, "sim:transmon:cycle", &self.artifacts(), joules_micro, grant, None)
    }
}

// ===========================================================================
// The loop close: PEC of a Pauli observable using the learned model
// ===========================================================================

/// Probabilistic error cancellation of a Pauli observable under the learned Pauli
/// channel. The twirled noise attenuates `⟨P⟩` by exactly `f_P`, so the ideal
/// value is `⟨P⟩_noisy / f_P` with the learned fidelity. Returns
/// `(raw_noisy_fx, mitigated_fx)`.
pub fn pec_pauli(model: &NoiseModel, observable: &Pauli, noisy_expectation_fx: i64) -> (i64, i64) {
    let f_p = model.fidelity(observable);
    let mitigated = fdiv(noisy_expectation_fx, f_p.max(1));
    (noisy_expectation_fx, mitigated.clamp(-ONE, ONE))
}

/// Seal the PEC step as a mitigation receipt that binds the learned model's hash —
/// the loop close: the mitigation names exactly which learned model produced its
/// correction.
#[allow(clippy::too_many_arguments)]
pub fn seal_pec(
    signer: &SigningKey,
    signer_id: &str,
    model: &NoiseModel,
    observable_label: &str,
    raw_input_bytes: &[u8],
    raw_fx: i64,
    mitigated_fx: i64,
    error_bar_fx: i64,
    shots: u64,
    joules_micro: u64,
    grant: GrantRef,
) -> MitigationReceipt {
    MitigationReceipt::seal(
        signer,
        signer_id,
        "sim:transmon:cycle",
        "pec.pauli_lindblad",
        observable_label,
        content_hash(raw_input_bytes),
        Some(model.model_hash()),
        raw_fx,
        mitigated_fx,
        error_bar_fx,
        shots,
        joules_micro,
        grant,
        None,
    )
}

/// A representative hidden device model on `n` qubits with the given `edges`: a
/// deterministic spread of rates over the sparse generators.
pub fn hidden_model(n: u32, edges: &[(u32, u32)], seed: u64) -> NoiseModel {
    let generators = sparse_generators(n, edges);
    let mut st = seed.wrapping_mul(0xD1B5_4A32).wrapping_add(0x1234_5678);
    let lambdas = generators
        .iter()
        .map(|g| {
            let r = (splitmix64(&mut st) % 1000) as i64; // [0, 1000)
            // weight-1 generators noisier than weight-2 (realistic), small rates
            let base = if g.weight() == 1 { ONE / 80 } else { ONE / 300 };
            base + fmul(base, r * ONE / 1000)
        })
        .collect();
    NoiseModel { generators, lambdas }
}

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

    fn key(s: u8) -> SigningKey {
        SigningKey::from_bytes(&[s; 32])
    }
    fn line3() -> (u32, Vec<(u32, u32)>) {
        (3, vec![(0, 1), (1, 2)])
    }
    fn f(v: i64) -> f64 {
        v as f64 / ONE as f64
    }

    #[test]
    fn pauli_anticommutation() {
        let x0 = Pauli::new(0b001, 0);
        let z0 = Pauli::new(0, 0b001);
        let x1 = Pauli::new(0b010, 0);
        assert!(x0.anticommutes(&z0)); // X and Z on same qubit anticommute
        assert!(!x0.anticommutes(&x1)); // disjoint support commute
        assert!(!x0.anticommutes(&x0)); // self commutes
    }

    #[test]
    fn fidelity_is_below_one_and_multiplicative() {
        let (n, e) = line3();
        let m = hidden_model(n, &e, 1);
        for a in all_paulis(n).iter().take(10) {
            let fa = m.fidelity(a);
            assert!(fa > 0 && fa <= ONE, "fidelity in (0,1]: {}", f(fa));
        }
    }

    #[test]
    fn cer_recovers_hidden_rates() {
        // The heart of it: reconstruct a hidden sparse Pauli–Lindblad model from a
        // simulated Cycle-Benchmarking experiment.
        let (n, e) = line3();
        let truth = hidden_model(n, &e, 7);
        let gens = sparse_generators(n, &e);
        let probes = all_paulis(n);
        let depths = [1u32, 2, 4, 8, 16];
        let r = learn_noise_model(&truth, &gens, &probes, &depths, 60_000, 200, 0xCAFE);
        // learned rates track the truth: mean abs error well below the mean rate
        let mean_rate = truth.total_rate() / truth.lambdas.len() as i64;
        assert!(
            r.rate_error_fx < mean_rate,
            "rate error {} should be below mean rate {}",
            f(r.rate_error_fx), f(mean_rate)
        );
        // and the reconstructed total rate is close to the true total
        let dt = (r.learned.total_rate() - truth.total_rate()).abs();
        assert!(dt < truth.total_rate() / 3, "total rate close: Δ {}", f(dt));
    }

    #[test]
    fn learning_is_deterministic() {
        let (n, e) = line3();
        let truth = hidden_model(n, &e, 3);
        let gens = sparse_generators(n, &e);
        let probes = all_paulis(n);
        let d = [1u32, 2, 4, 8];
        let a = learn_noise_model(&truth, &gens, &probes, &d, 20_000, 100, 42);
        let b = learn_noise_model(&truth, &gens, &probes, &d, 20_000, 100, 42);
        assert_eq!(a, b);
    }

    #[test]
    fn pec_recovers_ideal_pauli_expectation() {
        // A Pauli observable attenuated by the channel is restored by PEC using
        // the learned model.
        let (n, e) = line3();
        let truth = hidden_model(n, &e, 5);
        let gens = sparse_generators(n, &e);
        let probes = all_paulis(n);
        let r = learn_noise_model(&truth, &gens, &probes, &[1, 2, 4, 8, 16], 60_000, 200, 9);

        let obs = Pauli::new(0b111, 0); // X0 X1 X2
        let ideal = (0.9 * ONE as f64) as i64;
        let noisy = fmul(ideal, truth.fidelity(&obs)); // what the device would report
        let (raw, mit) = pec_pauli(&r.learned, &obs, noisy);
        assert_eq!(raw, noisy);
        assert!((mit - ideal).abs() < (0.05 * ONE as f64) as i64, "PEC ~ ideal: {} vs {}", f(mit), f(ideal));
    }

    #[test]
    fn model_seals_and_mitigation_binds_it() {
        // The loop close: seal the learned model, then a PEC receipt that binds
        // its hash. Both verify, and the binding checks out.
        let (n, e) = line3();
        let truth = hidden_model(n, &e, 11);
        let gens = sparse_generators(n, &e);
        let probes = all_paulis(n);
        let r = learn_noise_model(&truth, &gens, &probes, &[1, 2, 4, 8], 40_000, 150, 1);

        let model_rec = r.seal(&key(1), "did:key:lab", 2_000_000, GrantRef::unbounded("quantum.calibrate"));
        assert!(model_rec.verify());

        let obs = Pauli::new(0, 0b111); // Z0 Z1 Z2
        let raw_bytes = b"<twirled counts for Z0Z1Z2>";
        let mit = seal_pec(
            &key(1), "did:key:lab", &r.learned, "Z0Z1Z2",
            raw_bytes, ONE / 2, fdiv(ONE / 2, r.learned.fidelity(&obs)), ONE / 100,
            60_000, 500_000, GrantRef::unbounded("quantum.mitigate"),
        );
        assert!(mit.verify());
        // the mitigation names exactly the learned model
        assert!(mit.noise_model_matches(&r.learned.bytes()));
        assert_eq!(mit.noise_model_hash, Some(r.learned.model_hash()));
    }
}