stabilizer_qec 2.1.0

//! The coherent sampler's identity, checked exhaustively: every random choice a shot can make
//! (the frame's randomisation, every noise branch, every location firing or not), each with its
//! twirl probability, weighted by the whole spacetime kernel. The weighted outcomes must be the
//! state vector's to rounding, for any small circuit. This is where the signs are proved.

use std::collections::BTreeMap;

use super::shot::Source;
use super::{Element, Program, Slot};

/// Every branch of a shot, by replaying it with its choices counted like an odometer.
struct Odometer {
    choices: Vec<usize>,
    probs: Vec<Vec<f64>>,
    pos: usize,
    p: f64,
}

impl Odometer {
    fn choose(&mut self, probs: Vec<f64>) -> usize {
        if self.pos == self.choices.len() {
            let k = probs.iter().position(|&p| p > 0.0).unwrap_or(0);
            self.choices.push(k);
            self.probs.push(probs);
        }
        let k = self.choices[self.pos];
        self.p *= self.probs[self.pos][k];
        self.pos += 1;
        k
    }

    /// The next branch; false when every one has run.
    fn advance(&mut self) -> bool {
        while let Some(&k) = self.choices.last() {
            let probs = self.probs.last().unwrap();
            if let Some(next) = (k + 1..probs.len()).find(|&j| probs[j] > 0.0) {
                *self.choices.last_mut().unwrap() = next;
                return true;
            }
            self.choices.pop();
            self.probs.pop();
        }
        false
    }
}

impl Source for Odometer {
    fn bern(&mut self, p: f64) -> bool {
        self.choose(vec![1.0 - p, p]) == 1
    }

    fn pick(&mut self, probs: &[f64]) -> Option<usize> {
        let mut v = probs.to_vec();
        v.push(1.0 - probs.iter().sum::<f64>());
        let k = self.choose(v);
        (k < probs.len()).then_some(k)
    }

    fn index(&mut self, p: f64, n: u64) -> Option<u64> {
        let mut v = vec![1.0 - p];
        v.extend((0..n).map(|_| p / n as f64));
        let k = self.choose(v);
        (k > 0).then(|| k as u64 - 1)
    }

    fn coin(&mut self) -> bool {
        self.choose(vec![0.5, 0.5]) == 1
    }
}

impl Program {
    /// The whole spacetime kernel, by trying every set of locations (a few dozen locations at
    /// most).
    pub fn full_kernel(&self) -> Result<Vec<Element>, String> {
        let n = self.locations.len();
        if n > 20 {
            return Err(format!("{n} locations: the whole kernel is for small circuits"));
        }
        let solver = super::GaugeSolver::new(self);
        let mut out = Vec::new();
        for mask in 1u32..(1 << n) {
            let members: Vec<usize> = (0..n).filter(|&g| mask >> g & 1 == 1).collect();
            out.extend(self.element(&members, &solver));
        }
        Ok(out)
    }

    /// Every outcome's probability by the weighted twirl, summed over every branch. With `raw`,
    /// detectors and observables are the record's parities themselves rather than compared with
    /// the reference run's.
    pub fn exact_distribution(&self, kernel: &[Element], raw: bool, max_branches: usize) -> Result<BTreeMap<(Vec<bool>, u64), f64>, String> {
        self.exact_distribution_by(&mut |shot| self.weight(shot, kernel), raw, max_branches)
    }

    /// The same with any per-shot weight (a truncated kernel's, to measure its bias exactly).
    pub fn exact_distribution_by(&self, weight: &mut dyn FnMut(&super::Shot) -> f64, raw: bool, max_branches: usize) -> Result<BTreeMap<(Vec<bool>, u64), f64>, String> {
        let reference: Vec<bool> = self.slots.iter().zip(&self.reference).filter(|(s, _)| matches!(s, Slot::Record(_))).map(|(_, &v)| v).collect();
        let mut out = BTreeMap::new();
        let mut odo = Odometer { choices: Vec::new(), probs: Vec::new(), pos: 0, p: 1.0 };
        let mut branches = 0usize;
        loop {
            odo.pos = 0;
            odo.p = 1.0;
            let shot = self.run_shot(&mut odo);
            branches += 1;
            if branches > max_branches {
                return Err(format!("more than {max_branches} branches"));
            }
            let rec: Vec<bool> = if raw { shot.records.iter().zip(&reference).map(|(a, b)| a ^ b).collect() } else { shot.records.clone() };
            let w = weight(&shot);
            *out.entry(self.outcome(&rec)).or_insert(0.0) += odo.p * w;
            if !odo.advance() {
                break;
            }
        }
        Ok(out)
    }
}

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

    /// The state vector's distribution of raw parities.
    fn truth(c: &Circuit) -> BTreeMap<(Vec<bool>, u64), f64> {
        let sv = crate::statevec::Program::new(c).unwrap();
        let reference = sv.reference_parities().to_vec();
        let nd = sv.num_detectors();
        let mut out = BTreeMap::new();
        for ((dets, obs), p) in sv.distribution(1 << 22).unwrap() {
            let d: Vec<bool> = dets.iter().zip(&reference).map(|(a, b)| a ^ b).collect();
            let o = (0..sv.num_observables()).fold(0u64, |acc, k| acc | (((obs >> k & 1 == 1) ^ reference[nd + k]) as u64) << k);
            *out.entry((d, o)).or_insert(0.0) += p;
        }
        out
    }

    fn same(text: &str, what: &str, ours: &BTreeMap<(Vec<bool>, u64), f64>, theirs: &BTreeMap<(Vec<bool>, u64), f64>) {
        let keys: std::collections::BTreeSet<_> = ours.keys().chain(theirs.keys()).cloned().collect();
        for k in keys {
            let (a, b) = (ours.get(&k).copied().unwrap_or(0.0), theirs.get(&k).copied().unwrap_or(0.0));
            assert!((a - b).abs() < 1e-10, "{text}\n{what} {k:?}: weighted twirl {a}, state vector {b}");
        }
    }

    /// The weighted twirl's distribution is the state vector's, both as raw parities (the
    /// weights' identity) and as detection events (what the sampler reports: flips from the
    /// reference run, which must be the state vector's reference too; a detector on a random
    /// outcome tells the two apart).
    fn check(text: &str) {
        let c = Circuit::parse(text).unwrap();
        let p = Program::new(&c).unwrap();
        let kernel = p.full_kernel().unwrap();
        same(text, "raw", &p.exact_distribution(&kernel, true, 1 << 22).unwrap(), &truth(&c));
        let sv = crate::statevec::Program::new(&c).unwrap();
        same(text, "events", &p.exact_distribution(&kernel, false, 1 << 22).unwrap(), &sv.distribution(1 << 22).unwrap());
    }

    #[test]
    fn one_rotation() {
        check("R 0\nI_ERROR[R_X(theta=0.3)] 0\nM 0\nDETECTOR rec[-1]");
    }

    #[test]
    fn rotations_that_add() {
        check("R 0\nI_ERROR[R_X(theta=0.3)] 0\nI_ERROR[R_X(theta=0.2)] 0\nM 0\nDETECTOR rec[-1]");
        check("R 0\nI_ERROR[R_X(theta=0.3)] 0\nI_ERROR[R_X(theta=-0.2)] 0\nM 0\nDETECTOR rec[-1]");
    }

    #[test]
    fn rotations_through_gates() {
        check("R 0\nH 0\nI_ERROR[R_Z(theta=0.3)] 0\nH 0\nI_ERROR[R_X(theta=0.2)] 0\nM 0\nDETECTOR rec[-1]");
        check("R 0 1\nI_ERROR[R_Y(theta=0.3)] 0\nCX 0 1\nI_ERROR[R_X(theta=0.25)] 1\nS 1\nI_ERROR[R_Y(theta=0.1)] 1\nM 0 1\nDETECTOR rec[-1]\nDETECTOR rec[-2]");
        check("R 0 1\nH 0 1\nII_ERROR[R_ZZ(theta=0.4)] 0 1\nCZ 0 1\nI_ERROR[R_Z(theta=0.2)] 1\nH 0 1\nM 0 1\nDETECTOR rec[-1]\nDETECTOR rec[-2]");
    }

    #[test]
    fn rotations_with_noise_and_mid_circuit_measurement() {
        check("R 0 1\nH 0\nI_ERROR[R_X(theta=0.3)] 1\nCX 0 1\nDEPOLARIZE1(0.1) 1\nM(0.05) 1\nI_ERROR[R_X(theta=0.2)] 1\nM 1\nDETECTOR rec[-1] rec[-2]\nOBSERVABLE_INCLUDE(0) rec[-1]");
        check("R 0\nI_ERROR[R_X(theta=0.3)] 0\nMR 0\nI_ERROR[R_X(theta=0.2)] 0\nX_ERROR(0.1) 0\nM 0\nDETECTOR rec[-1]\nDETECTOR rec[-2]");
        check("RX 0\nI_ERROR[R_Z(theta=0.3)] 0\nMX 0\nRX 0\nI_ERROR[R_Z(theta=0.2)] 0\nMX 0\nDETECTOR rec[-1]\nDETECTOR rec[-2]");
    }

    #[test]
    fn repetition_code_rounds() {
        check(
            "R 0 1 2\nI_ERROR[R_X(theta=0.2)] 0 2\nCX 0 1\nCX 2 1\nMR 1\nI_ERROR[R_X(theta=0.15)] 0 2\nCX 0 1\nCX 2 1\nMR 1\n\
             DETECTOR rec[-2]\nDETECTOR rec[-1] rec[-2]\nM 0 2\nDETECTOR rec[-1] rec[-2] rec[-3]\nOBSERVABLE_INCLUDE(0) rec[-1]",
        );
        // The same in X: rotations about Z on |+⟩ data, X-type checks.
        check(
            "RX 0 2\nR 1\nI_ERROR[R_Z(theta=0.2)] 0 2\nH 1\nCZ 1 0\nCZ 1 2\nH 1\nMR 1\nI_ERROR[R_Z(theta=-0.3)] 0 2\nH 1\nCZ 1 0\nCZ 1 2\nH 1\nMR 1\n\
             DETECTOR rec[-2]\nDETECTOR rec[-1] rec[-2]\nMX 0 2\nDETECTOR rec[-1] rec[-2] rec[-3]\nOBSERVABLE_INCLUDE(0) rec[-1]",
        );
    }

    /// A detector on a random outcome, with random outcomes on other qubits before it: the
    /// adversarial review's circuits, whose detector the 1.9 sampler reported complemented
    /// whenever its seeded reference run read the measurement as 1.
    #[test]
    fn detectors_on_random_outcomes() {
        check("R 0 1 2\nI_ERROR[R_Y(theta=0.24)] 0\nMX 1 2\nMX 0\nDETECTOR rec[-1]");
        check("R 0 1 2\nI_ERROR[R_Y(theta=0.24)] 0\nMX 1\nMX 0\nDETECTOR rec[-1]");
        check("R 0 1 2 3\nI_ERROR[R_X(theta=0.15)] 0 2\nCX 0 1 2 1\nM 1\nI_ERROR[R_Y(theta=0.2)] 3\nMX 0 2\nMX 3\nDETECTOR rec[-1]");
        check("RX 0 1 2\nI_ERROR[R_Z(theta=0.3)] 0\nMX 1\nMY 0\nDETECTOR rec[-1]");
    }

    /// The review's family: a rotated qubit read out where its outcome is random, after k
    /// other qubits with random outcomes nobody compares (read in X, read and reset in X, H
    /// then M, or prepared with RX and never read), k = 0..4; the detector on the random
    /// outcome, and the whole record's events besides.
    #[test]
    fn random_outcomes_elsewhere() {
        let line = |op: &str, qs: &[usize]| if qs.is_empty() { None } else { Some(format!("{op} {}", qs.iter().map(|q| q.to_string()).collect::<Vec<_>>().join(" "))) };
        for k in 0..=4usize {
            let others: Vec<usize> = (1..=k).collect();
            let all: Vec<usize> = (0..=k).collect();
            for (prep, rotation, readout) in [("R", "I_ERROR[R_Y(theta=0.24)] 0", "MX 0"), ("RX", "I_ERROR[R_Z(theta=0.3)] 0", "MY 0"), ("R", "I_ERROR[R_X(theta=0.2)] 0", "MY 0")] {
                for elsewhere in ["MX", "MRX", "H+M", "RX"] {
                    // Each MRX doubles the branches twice over; four of them cost seconds.
                    if elsewhere == "MRX" && k == 4 {
                        continue;
                    }
                    let mut lines: Vec<String> = vec![line(prep, &all).unwrap(), rotation.to_string()];
                    match elsewhere {
                        "H+M" => {
                            lines.extend(line("H", &others));
                            lines.extend(line("M", &others));
                        }
                        "RX" => lines.extend(line("RX", &others)),
                        op => lines.extend(line(op, &others)),
                    }
                    lines.push(readout.to_string());
                    lines.push("DETECTOR rec[-1]".to_string());
                    if elsewhere != "RX" && k > 0 {
                        lines.push("DETECTOR rec[-2]".to_string());
                        lines.push(format!("OBSERVABLE_INCLUDE(0) rec[-1] rec[-{}]", k + 1));
                    }
                    check(&lines.join("\n"));
                }
            }
        }
    }

    /// Random small circuits of every kind of step.
    #[test]
    fn random_circuits() {
        let mut rng = crate::surface_code::Xorshift::new(2026);
        let ops1 = ["H", "S", "X", "Y", "Z", "S_DAG", "SQRT_X", "I_ERROR[R_X(theta=0.3)]", "I_ERROR[R_Y(theta=-0.25)]", "I_ERROR[R_Z(theta=0.4)]", "X_ERROR(0.1)", "Z_ERROR(0.15)", "DEPOLARIZE1(0.1)", "R", "RX", "M", "MX", "MR", "M(0.1)"];
        let ops2 = ["CX", "CZ", "CY", "SWAP", "II_ERROR[R_ZZ(theta=0.35)]", "II_ERROR[R_XX(theta=-0.2)]", "DEPOLARIZE2(0.1)"];
        for case in 0..150 {
            let n = 1 + (rng.next_u64() % 3) as usize;
            let mut lines = Vec::new();
            let (mut rotations, mut randomness, mut measured, mut noise) = (0, 0, 0usize, 0);
            while lines.len() < 12 {
                let two = n > 1 && rng.next_u64().is_multiple_of(3);
                let op = if two { ops2[(rng.next_u64() % ops2.len() as u64) as usize] } else { ops1[(rng.next_u64() % ops1.len() as u64) as usize] };
                if op.contains("ERROR(") || op.contains("DEPOLARIZE") || op.contains("M(") {
                    if noise == 2 {
                        continue;
                    }
                    noise += 1;
                }
                if op.contains("theta") {
                    if rotations == 4 {
                        continue;
                    }
                    rotations += 1;
                }
                if op.starts_with('M') || op.starts_with('R') && !op.starts_with("R_") {
                    if randomness == 4 {
                        continue;
                    }
                    randomness += 1;
                    measured += op.starts_with('M') as usize;
                }
                let a = (rng.next_u64() % n as u64) as usize;
                if two {
                    let b = (a + 1 + (rng.next_u64() % (n as u64 - 1)) as usize) % n;
                    lines.push(format!("{op} {a} {b}"));
                } else {
                    lines.push(format!("{op} {a}"));
                }
            }
            let finals: Vec<String> = (0..n).map(|q| q.to_string()).collect();
            lines.push(format!("M {}", finals.join(" ")));
            let total = measured + n;
            for k in 0..3usize.min(total) {
                let a = 1 + (rng.next_u64() % total as u64) as usize;
                let b = 1 + (rng.next_u64() % total as u64) as usize;
                lines.push(if a == b || k == 0 { format!("DETECTOR rec[-{a}]") } else { format!("DETECTOR rec[-{a}] rec[-{b}]") });
            }
            lines.push(format!("OBSERVABLE_INCLUDE(0) rec[-{}]", 1 + (rng.next_u64() % total as u64)));
            let text = lines.join("\n");
            let _ = case;
            check(&text);
        }
    }
}