open-standards/wai/wai-quantum/examples/qec_memory/pairs.rs

//! This crate's decoder against each published decoder of a record, shot by
//! shot: exact McNemar tests on the shots only one of the two got wrong, and
//! a paired bootstrap of the fitted rates and of `Λ`.
//!
//! ```text
//! qec_memory pairs SET [--replicates N] [--seed S] [--threads T] [email protected]...
//! ```
//!
//! Each `KEY@RECEIPT` names this crate's decoder (`decoder|prior`) in a
//! receipt `survey` sealed with its discordant counts; the receipts' fitted
//! experiments are pooled (a receipt for each distance, say). For every
//! published decoder scored in all of them:
//!
//! - per distance, and over every distance given, the failures of each and
//!   the discordant shots `b` (this decoder wrong, the published one right)
//!   and `c` (the reverse), with the exact two-sided McNemar p-value: under
//!   the null that the two fail alike, each discordant shot is either way
//!   with probability one half, `p = min(1, 2·P(X ≤ min(b, c)))`,
//!   `X ~ Bin(b + c, 1/2)`;
//! - the fitted `ε` per distance and `Λ` of each, as `survey` fits them;
//! - a paired bootstrap: in each replicate every fitted experiment's shots are
//!   drawn again with replacement, the two decoders' outcomes on a shot kept
//!   together (the experiment's joint table of right and wrong is all a
//!   draw needs), and both are fitted again. The 2.5 and 97.5 percentiles of
//!   `Λ_this − Λ_published` and of `ε_this/ε_published` per distance, and the
//!   fraction of replicates where this decoder's `Λ` is not above the
//!   published one's. Replicate `r` draws from its own generator, seeded
//!   from `S` and `r`, so the result does not depend on the threads.

use crate::data::Set;
use serde_json::{Value, json};
use std::collections::BTreeMap;
use wai_quantum::quantum_memory::{RoundPoint, fit_decay, fit_lambda, mean_rate};

/// One fitted experiment, against one published decoder: its series
/// (distance, patch, basis), rounds and shots, and the joint table
/// `[both wrong, only this decoder, only the published one]` (both right
/// is the rest).
#[derive(Clone)]
struct Joint {
    series: (u32, String, String),
    rounds: u32,
    shots: u64,
    both: u64,
    only_ours: u64,
    only_theirs: u64,
}

/// The exact two-sided McNemar p-value for `b` and `c` discordant shots.
pub fn mcnemar(b: u64, c: u64) -> f64 {
    let n = b + c;
    if n == 0 {
        return 1.0;
    }
    let k = b.min(c);
    // P(X ≤ k) as a sum of terms in log space, from the largest:
    // ln(C(n, i)/2ⁿ), each from the one before.
    let mut terms = Vec::with_capacity(k as usize + 1);
    let mut t = -(n as f64) * std::f64::consts::LN_2;
    for i in 0..=k {
        terms.push(t);
        t += ((n - i) as f64).ln() - ((i + 1) as f64).ln();
    }
    let top = terms.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let tail = top.exp() * terms.iter().map(|t| (t - top).exp()).sum::<f64>();
    (2.0 * tail).min(1.0)
}

/// Fitted ε per distance and Λ from per-experiment failures.
fn fit(joints: &[Joint], errors: impl Fn(&Joint) -> u64) -> (BTreeMap<u32, f64>, Option<f64>) {
    let mut groups: BTreeMap<&(u32, String, String), Vec<RoundPoint>> = BTreeMap::new();
    for j in joints {
        groups.entry(&j.series).or_default().push(RoundPoint { rounds: j.rounds, shots: j.shots, errors: errors(j) });
    }
    let mut per_d: BTreeMap<u32, Vec<_>> = BTreeMap::new();
    for (series, pts) in &groups {
        if let Some(f) = fit_decay(pts, 0) {
            per_d.entry(series.0).or_default().push(f);
        }
    }
    let mut eps = BTreeMap::new();
    let mut lam_in = Vec::new();
    for (d, fits) in &per_d {
        let m = mean_rate(fits).unwrap();
        eps.insert(*d, m.eps);
        lam_in.push((*d, m.eps, m.sigma_independent));
    }
    (eps, fit_lambda(&lam_in).map(|l| l.0))
}

struct SplitMix(u64);

impl SplitMix {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        z ^ (z >> 31)
    }
    /// Uniform below `n` (Lemire's multiply, the bias below 2⁻⁴⁰ for the
    /// shot counts here).
    fn below(&mut self, n: u64) -> u64 {
        ((u128::from(self.next()) * u128::from(n)) >> 64) as u64
    }
}

/// One replicate: every experiment's shots drawn again.
fn redraw(joints: &[Joint], rng: &mut SplitMix) -> Vec<Joint> {
    joints
        .iter()
        .map(|j| {
            let (mut both, mut ours, mut theirs) = (0, 0, 0);
            let (a, b, c) = (j.both, j.both + j.only_ours, j.both + j.only_ours + j.only_theirs);
            for _ in 0..j.shots {
                let r = rng.below(j.shots);
                if r < a {
                    both += 1;
                } else if r < b {
                    ours += 1;
                } else if r < c {
                    theirs += 1;
                }
            }
            Joint { both, only_ours: ours, only_theirs: theirs, ..j.clone() }
        })
        .collect()
}

fn quantile(v: &mut [f64], q: f64) -> f64 {
    v.sort_by(|a, b| a.total_cmp(b));
    v[((v.len() as f64 - 1.0) * q).round() as usize]
}

pub fn pairs(a: &[String]) {
    let set = Set::parse(a.first().map(String::as_str).unwrap_or("")).unwrap_or_else(|| crate::usage());
    let (mut replicates, mut seed, mut threads) = (2000usize, 1u64, 8usize);
    let mut inputs: Vec<(String, String)> = Vec::new();
    let mut i = 1;
    while i < a.len() {
        match a[i].as_str() {
            "--replicates" | "--seed" | "--threads" => {
                let v = a.get(i + 1).unwrap_or_else(|| crate::usage());
                match a[i].as_str() {
                    "--replicates" => replicates = v.parse().unwrap(),
                    "--seed" => seed = v.parse().unwrap(),
                    _ => threads = v.parse().unwrap(),
                }
                i += 2;
            }
            x => {
                let (k, p) = x.rsplit_once('@').unwrap_or_else(|| crate::usage());
                inputs.push((k.to_string(), p.to_string()));
                i += 1;
            }
        }
    }
    // Every fitted experiment: this crate's key, its counts and its discordant table.
    let mut exps: Vec<(String, Value)> = Vec::new();
    for (key, path) in &inputs {
        let doc: Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap_or_else(|e| panic!("{path}: {e}"))).unwrap();
        assert!(wai_quantum::engine::receipt::verify(&doc["request"], &doc["result"], &doc["receipt"]), "{path}: the receipt does not verify");
        for e in doc["result"]["experiments"].as_array().unwrap() {
            if set.fitted(e["rounds"].as_u64().unwrap() as u32) {
                assert!(!e["discordant"][key.as_str()].is_null(), "{path}: no discordant counts for {key} (sealed before survey kept them?)");
                exps.push((key.clone(), e.clone()));
            }
        }
    }
    let mut published: Vec<String> = exps[0].1["discordant"][exps[0].0.as_str()].as_object().unwrap().keys().cloned().collect();
    published.retain(|p| exps.iter().all(|(k, e)| !e["discordant"][k.as_str()][p.as_str()].is_null()));
    let mut out = Vec::new();
    for p in &published {
        let joints: Vec<Joint> = exps
            .iter()
            .map(|(k, e)| {
                let n = |key: &str| e["errors"][key].as_u64().unwrap();
                let d = &e["discordant"][k.as_str()][p.as_str()];
                let (b, c) = (d[0].as_u64().unwrap(), d[1].as_u64().unwrap());
                let both = n(k) - b;
                assert_eq!(both, n(p) - c, "{}: the joint table adds up", e["experiment"]);
                Joint {
                    series: (e["distance"].as_u64().unwrap() as u32, e["patch"].as_str().unwrap().into(), e["basis"].as_str().unwrap().into()),
                    rounds: e["rounds"].as_u64().unwrap() as u32,
                    shots: e["shots"].as_u64().unwrap(),
                    both,
                    only_ours: b,
                    only_theirs: c,
                }
            })
            .collect();
        let ours = |j: &Joint| j.both + j.only_ours;
        let theirs = |j: &Joint| j.both + j.only_theirs;
        let mut tests = serde_json::Map::new();
        let distances: Vec<u32> = {
            let mut v: Vec<u32> = joints.iter().map(|j| j.series.0).collect();
            v.sort_unstable();
            v.dedup();
            v
        };
        let mut pooled = |label: String, js: Vec<&Joint>| {
            let (fo, ft) = (js.iter().map(|j| ours(j)).sum::<u64>(), js.iter().map(|j| theirs(j)).sum::<u64>());
            let (b, c) = (js.iter().map(|j| j.only_ours).sum::<u64>(), js.iter().map(|j| j.only_theirs).sum::<u64>());
            let shots: u64 = js.iter().map(|j| j.shots).sum();
            println!("{p} {label}: {} experiments, {shots} shots; failures {fo} (this) against {ft}; discordant {b} against {c}; McNemar p = {:.3e}", js.len(), mcnemar(b, c));
            tests.insert(label, json!({ "experiments": js.len(), "shots": shots, "failures": [fo, ft], "discordant": [b, c], "mcnemar_p": mcnemar(b, c) }));
        };
        for &d in &distances {
            pooled(format!("d{d}"), joints.iter().filter(|j| j.series.0 == d).collect());
        }
        if distances.len() > 1 {
            pooled("all".into(), joints.iter().collect());
        }
        let (eo, lo) = fit(&joints, ours);
        let (et, lt) = fit(&joints, theirs);
        println!("{p}: eps this {eo:?} Lambda {lo:?}; eps published {et:?} Lambda {lt:?}");
        // The bootstrap, replicates spread over threads, each from its own seed.
        let reps: Vec<(BTreeMap<u32, f64>, Option<f64>)> = {
            let next = std::sync::atomic::AtomicUsize::new(0);
            let slots = std::sync::Mutex::new(vec![None; replicates]);
            std::thread::scope(|s| {
                for _ in 0..threads.max(1) {
                    s.spawn(|| {
                        loop {
                            let r = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                            if r >= replicates {
                                break;
                            }
                            let mut rng = SplitMix(seed ^ (r as u64).wrapping_mul(0xd1b5_4a32_d192_ed03));
                            let js = redraw(&joints, &mut rng);
                            let (a, la) = fit(&js, ours);
                            let (b, lb) = fit(&js, theirs);
                            let ratio: BTreeMap<u32, f64> = a.iter().filter_map(|(d, x)| b.get(d).map(|y| (*d, x / y))).collect();
                            let dl = la.zip(lb).map(|(x, y)| x - y);
                            slots.lock().unwrap()[r] = Some((ratio, dl));
                        }
                    });
                }
            });
            slots.into_inner().unwrap().into_iter().map(Option::unwrap).collect()
        };
        let mut boot = serde_json::Map::new();
        for &d in &distances {
            let mut v: Vec<f64> = reps.iter().filter_map(|r| r.0.get(&d).copied()).collect();
            if v.is_empty() {
                continue;
            }
            let above = v.iter().filter(|&&x| x >= 1.0).count() as f64 / v.len() as f64;
            let (lo95, hi95) = (quantile(&mut v, 0.025), quantile(&mut v, 0.975));
            println!("{p} d{d}: eps this/published, 95% interval [{lo95:.4}, {hi95:.4}], ratio at least one in {:.4} of {} replicates", above, v.len());
            boot.insert(format!("eps_ratio_d{d}"), json!({ "p2_5": lo95, "p97_5": hi95, "fraction_at_least_one": above, "replicates": v.len() }));
        }
        let mut dl: Vec<f64> = reps.iter().filter_map(|r| r.1).collect();
        if !dl.is_empty() {
            let not_above = dl.iter().filter(|&&x| x <= 0.0).count() as f64 / dl.len() as f64;
            let (lo95, hi95) = (quantile(&mut dl, 0.025), quantile(&mut dl, 0.975));
            println!("{p}: Lambda this − published, 95% interval [{lo95:.4}, {hi95:.4}], not above zero in {not_above:.4} of {} replicates", dl.len());
            boot.insert("lambda_difference".into(), json!({ "p2_5": lo95, "p97_5": hi95, "fraction_not_above_zero": not_above, "replicates": dl.len() }));
        }
        out.push(json!({
            "published": p, "tests": tests,
            "eps_this": eo, "lambda_this": lo, "eps_published": et, "lambda_published": lt,
            "bootstrap": boot,
        }));
    }
    let doc = json!({
        "set": set.name(), "inputs": inputs.iter().map(|(k, p)| format!("{k}@{p}")).collect::<Vec<_>>(),
        "replicates": replicates, "seed": seed, "pairs": out,
    });
    println!("{doc}");
}