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

//! Surface-code memory experiments decoded from the processors' own records,
//! with this crate's decoders, and the published decoders' predictions on the
//! same shots scored the same way.
//!
//! ```text
//! qec_memory survey SET DIR --decoders mwpm,uf,corr --prior PRIOR [--threads N] [--shots N] [--skip N] [--only SUBSTR] [--predictions DIR] --out FILE.jsonl [--receipt FILE.json]
//! qec_memory fit SET FILE.jsonl...
//! qec_memory pairs SET [--replicates N] [--seed S] [--threads T] [email protected]...
//! ```
//!
//! `--shots N --skip K` decodes shots `K` to `K + N` of each experiment (`K`
//! even), and scores the published decoders on the same shots. With
//! `--predictions DIR` every decoder's predictions are written there, one
//! byte a shot, with the published decoders' (`_published_NAME.b8`) and the
//! record's own flips (`_actual.b8`). A receipt carries, for each of this
//! crate's decoders against each published one, the shots only one of the
//! two got wrong; `pairs` tests those (see `pairs.rs`).
//!
//! `SET` is `2023` (doi:10.5281/zenodo.6804040, the unpacked archive of
//! the 2023 distance-3 and distance-5 experiment) or `2024`
//! (doi:10.5281/zenodo.13273331, the unpacked archive of the 105-qubit
//! distance-3, 5 and 7 experiment); both are CC BY 4.0, credited in
//! `tests/data/qec_memory/README.md`. `DIR` is the
//! directory the archive unpacked into. Priors: `pij` (2023: the models fitted
//! to the data, each decoding the other half of the shots), `circuit` (2023:
//! the noisy circuit's model), `si1000` and `rl` (2024: the two published
//! decoder priors), `fit` (this crate's own fit to the detection events, each
//! half of the shots decoded with the model fitted to the other half).

mod data;
mod pairs;
mod receipt;

use data::Set;
use serde_json::{Value, json};
use std::path::{Path, PathBuf};
use std::time::Instant;
use wai_quantum::quantum_frame::ErrorModel;
use wai_quantum::quantum_match::{CorrelatedMatching, MatchingGraph};
use wai_quantum::quantum_memory::{BeliefMatcher, CorrMatcher, DecayFit, Matcher, MlDecoder, RoundPoint, fit_decay, fit_lambda, mean_rate, predict};
use wai_quantum::quantum_memory_stream::{Inner, ParallelWindows, Plan};
use wai_quantum::quantum_memory_tn::{SparseMl, TnDecoder};
use wai_quantum::quantum_uf::UnionFind;

fn usage() -> ! {
    eprintln!("usage: qec_memory survey SET DIR --decoders LIST --prior PRIOR [--threads N] [--shots N] [--skip N] [--only SUBSTR] [--predictions DIR] --out FILE.jsonl [--receipt FILE.json]\n       qec_memory fit SET FILE.jsonl...\n       qec_memory pairs SET [--replicates N] [--seed S] [--threads T] [email protected]...");
    std::process::exit(2)
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    match a.get(1).map(String::as_str) {
        Some("survey") => survey(&a[2..]),
        Some("fit") => fit(&a[2..]),
        Some("pairs") => pairs::pairs(&a[2..]),
        _ => usage(),
    }
}

/// One decoder, built for one model.
enum Dec {
    Mwpm(Matcher),
    Uf(UnionFind),
    Corr0(CorrelatedMatching),
    Corr(CorrMatcher),
    Bm(Box<BeliefMatcher>),
    Ml(MlDecoder),
    Tn(Box<TnDecoder>),
    Pw(Box<ParallelWindows>),
    Sp(Box<SparseMl>),
}

/// The largest table the exact decoder may use, in bits of index.
const ML_BITS: u32 = 22;

impl Dec {
    fn build(name: &str, m: &ErrorModel, distance: u32) -> Result<Dec, String> {
        Ok(match name {
            // ufw, mwpmw: parallel windows (commit = buffer = gap = d) of union-find or exact matching;
            // ufw<B>, mwpmw<B>: the same with buffer = gap = B rounds.
            w if w.starts_with("ufw") || w.starts_with("mwpmw") => {
                let (inner, rest) = if let Some(r) = w.strip_prefix("ufw") { (Inner::UnionFind, r) } else { (Inner::Matching, &w[5..]) };
                let mut plan = Plan::of_distance(distance);
                if !rest.is_empty() {
                    let b: u32 = rest.parse().map_err(|_| format!("{w}: a buffer in rounds"))?;
                    plan.buffer = b;
                    plan.gap = b;
                }
                Dec::Pw(Box::new(ParallelWindows::from_model(m, plan, inner)?))
            }
            "mwpm" => Dec::Mwpm(Matcher::from_model(m).expect("a graphlike model")),
            "uf" => Dec::Uf(UnionFind::new(&MatchingGraph::from_model(m).expect("a graphlike model"))),
            "corr0" => Dec::Corr0(CorrelatedMatching::from_model(m).expect("a graphlike model")),
            "corr" => Dec::Corr(CorrMatcher::from_model(m).expect("a graphlike model")),
            "corr2" => {
                let mut c = CorrMatcher::from_model(m).expect("a graphlike model");
                c.always_second = true;
                Dec::Corr(c)
            }
            b if b.starts_with("bm") => {
                // bm[ITERATIONS][h]: h returns propagation's own answer when it explains the syndrome.
                let hard = b.ends_with('h');
                let digits = b.trim_start_matches("bm").trim_end_matches('h');
                let mut d = BeliefMatcher::from_model(m, if digits.is_empty() { 20 } else { digits.parse().unwrap() }).expect("a graphlike model");
                d.trust_converged = hard;
                Dec::Bm(Box::new(d))
            }
            "ml" => Dec::Ml(MlDecoder::from_model(m, ML_BITS)?),
            // tn<CHI>[e<K>]: the boundary matrix-product state at bond CHI, singular values
            // below 10^-K of each cut's largest dropped (K = 14 when not given).
            // sp<CHI>[e<K>]: the same elimination with the table kept sparse, cut to its CHI
            // largest entries, entries below 10^-K of the largest dropped at each close (K = 12).
            // A trailing b on either: the truncation weighed by belief propagation's posteriors
            // (20 sweeps) instead of the priors.
            t if t.starts_with("sp") => {
                let (body, bp) = t.strip_suffix('b').map_or((t, false), |x| (x, true));
                let (chi, k) = body[2..].split_once('e').unwrap_or((&body[2..], "12"));
                let bad = |_| format!("{t}: sp, a number of entries, and e with an exponent");
                let mut d = SparseMl::from_model(m, chi.parse().map_err(bad)?)?;
                d.cutoff = format!("1e-{}", k.parse::<u32>().map_err(bad)?).parse().unwrap();
                if bp {
                    d.use_belief_propagation(m, 20)?;
                }
                Dec::Sp(Box::new(d))
            }
            t if t.starts_with("tn") => {
                let (body, bp) = t.strip_suffix('b').map_or((t, false), |x| (x, true));
                let (chi, k) = body[2..].split_once('e').unwrap_or((&body[2..], "14"));
                let bad = |_| format!("{t}: tn, a bond dimension, and e with an exponent");
                let mut d = TnDecoder::from_model(m, chi.parse().map_err(bad)?)?;
                // Read as decimal text, correctly rounded: the same double everywhere.
                d.cutoff = format!("1e-{}", k.parse::<u32>().map_err(bad)?).parse().unwrap();
                if bp {
                    d.use_belief_propagation(m, 20)?;
                }
                Dec::Tn(Box::new(d))
            }
            other => panic!("no decoder {other}"),
        })
    }
    fn describe(&self) -> Value {
        match self {
            Dec::Ml(d) => json!({ "width": d.width(), "mechanisms": d.mechanisms(), "work_per_shot": d.work_per_shot() }),
            Dec::Bm(b) => json!({ "iterations": b.iterations, "trust_converged": b.trust_converged }),
            Dec::Pw(p) => {
                let (a, b) = p.windows();
                json!({ "plan": format!("{:?}", p.plan()), "layer_a": a.len(), "layer_b": b.len(), "rounds_decoded": p.rounds_decoded(), "rounds": p.rounds() })
            }
            Dec::Sp(t) => json!({ "chi": t.chi, "cutoff": t.cutoff, "width": t.width(), "mechanisms": t.mechanisms() }),
            Dec::Tn(t) => json!({ "chi": t.chi, "cutoff": t.cutoff, "width": t.width(), "mechanisms": t.mechanisms(), "layout": format!("{:?}", t.layout()), "mean_span": t.mean_span() }),
            _ => Value::Null,
        }
    }
}

enum Work {
    M(wai_quantum::quantum_memory::MatchWork),
    U(wai_quantum::quantum_uf::Workspace),
    C(wai_quantum::quantum_memory::CorrWork),
    B(Box<wai_quantum::quantum_memory::BeliefWork>),
    L(wai_quantum::quantum_memory::MlWork),
    T(Box<wai_quantum::quantum_memory_tn::TnWork>),
    P(Box<wai_quantum::quantum_memory_stream::StreamWork>),
    S(Box<wai_quantum::quantum_memory_tn::SparseWork>),
    N,
}

impl Dec {
    fn work(&self) -> Work {
        match self {
            Dec::Mwpm(m) => Work::M(m.workspace()),
            Dec::Uf(u) => Work::U(u.workspace()),
            Dec::Corr0(_) => Work::N,
            Dec::Corr(c) => Work::C(c.workspace()),
            Dec::Bm(b) => Work::B(Box::new(b.workspace())),
            Dec::Ml(l) => Work::L(l.workspace()),
            Dec::Tn(t) => Work::T(Box::new(t.workspace())),
            Dec::Pw(p) => Work::P(Box::new(p.workspace())),
            Dec::Sp(t) => Work::S(Box::new(t.workspace())),
        }
    }
    fn decode(&self, w: &mut Work, fired: &[u32]) -> Result<u64, String> {
        match (self, w) {
            (Dec::Mwpm(m), Work::M(ws)) => m.decode(fired, ws).map_err(|e| format!("{e:?}")),
            (Dec::Uf(u), Work::U(ws)) => u.decode(fired, ws).map_err(|e| format!("{e:?}")),
            (Dec::Corr0(c), _) => c.decode(fired).map_err(|e| format!("{e:?}")),
            (Dec::Corr(c), Work::C(ws)) => c.decode(fired, ws).map_err(|e| format!("{e:?}")),
            (Dec::Bm(b), Work::B(ws)) => b.decode(fired, ws).map_err(|e| format!("{e:?}")),
            (Dec::Ml(l), Work::L(ws)) => l.decode(fired, ws),
            (Dec::Tn(t), Work::T(ws)) => {
                let r = t.decode(fired, ws);
                note(ws.last_forced, ws.last_max_bond as u64);
                r
            }
            (Dec::Pw(p), Work::P(ws)) => p.decode(fired, ws).map_err(|e| format!("{e:?}")),
            (Dec::Sp(t), Work::S(ws)) => {
                let r = t.decode(fired, ws);
                note(ws.last_forced, ws.last_max_entries as u64);
                r
            }
            _ => unreachable!(),
        }
    }
}

/// Over one decoder's run: shots whose truncation dropped every explanation of
/// some detector (closed against its bit), and the largest bond or table.
static FORCED_SHOTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
static LARGEST: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

fn note(forced: u32, size: u64) {
    use std::sync::atomic::Ordering::Relaxed;
    if forced > 0 {
        FORCED_SHOTS.fetch_add(1, Relaxed);
    }
    LARGEST.fetch_max(size, Relaxed);
}

struct Args {
    decoders: Vec<String>,
    prior: String,
    threads: usize,
    shots: Option<usize>,
    only: Vec<String>,
    out: PathBuf,
    receipt: Option<PathBuf>,
    /// Fit priors to every shot of the record, even when fewer are decoded.
    fit_all: bool,
    /// Write each decoder's predictions here, `b8`, one file an experiment.
    predictions: Option<PathBuf>,
    /// The first shot decoded (shots before it are left out); even.
    skip: usize,
}

fn parse_args(a: &[String]) -> Args {
    let mut r = Args { decoders: vec!["mwpm".into()], prior: String::new(), threads: 8, shots: None, only: Vec::new(), out: PathBuf::new(), receipt: None, fit_all: false, predictions: None, skip: 0 };
    let mut i = 0;
    while i < a.len() {
        if a[i] == "--fit-all" {
            r.fit_all = true;
            i += 1;
            continue;
        }
        let v = a.get(i + 1).cloned().unwrap_or_else(|| usage());
        match a[i].as_str() {
            "--decoders" => r.decoders = v.split(',').map(str::to_string).collect(),
            "--prior" => r.prior = v,
            "--threads" => r.threads = v.parse().unwrap(),
            "--shots" => r.shots = Some(v.parse().unwrap()),
            "--only" => r.only = v.split(',').map(str::to_string).collect(),
            "--out" => r.out = PathBuf::from(v),
            "--receipt" => r.receipt = Some(PathBuf::from(v)),
            "--predictions" => r.predictions = Some(PathBuf::from(v)),
            "--skip" => r.skip = v.parse().unwrap_or_else(|_| usage()),
            _ => usage(),
        }
        i += 2;
    }
    if r.prior.is_empty() || r.out.as_os_str().is_empty() {
        usage();
    }
    if r.skip % 2 == 1 {
        // A shot's parity picks its model (the 2023 fitted models, and the
        // fits here, are cross-fitted by parity), so a selection keeps it.
        eprintln!("--skip: an even number of shots");
        std::process::exit(2);
    }
    r
}

/// One experiment's rows: each decoder's, then each published decoder's.
fn survey_one(set: Set, e: &data::Experiment, args: &Args) -> Vec<Value> {
    let mut rows = Vec::new();
    let all = data::shots(set, e);
    let skip = args.skip;
    let shots = match (skip, args.shots) {
        (0, None) => all.clone(),
        (0, Some(n)) => all.first(n),
        (k, n) => all.select(|s| s >= k && n.is_none_or(|n| s < k + n)),
    };
    let (events, rate) = shots.event_rate();
    let base = json!({
        "set": set.name(), "experiment": e.name, "distance": e.distance, "basis": e.basis, "rounds": e.rounds, "patch": e.patch,
        "detectors": shots.detectors(), "shots": shots.len(), "events": events, "event_rate": rate, "shots_blake3": hex(&shots.digest()),
    });
    let published = data::published(set, e, skip, shots.len());
    let wrong = |pred: &[u64], s: usize| pred[s] != shots.flips(s);
    if let Some(dir) = &args.predictions {
        // The record's own flips and each published decoder's predictions on
        // the same shots, beside this crate's, so every pair can be compared
        // shot by shot.
        std::fs::create_dir_all(dir).unwrap();
        let stem = e.name.replace('/', "_");
        let actual: Vec<u8> = (0..shots.len()).map(|s| shots.flips(s) as u8).collect();
        std::fs::write(dir.join(format!("{stem}_actual.b8")), actual).unwrap();
        for (name, pred) in &published {
            std::fs::write(dir.join(format!("{stem}_published_{name}.b8")), pred.iter().map(|&f| f as u8).collect::<Vec<u8>>()).unwrap();
        }
    }
    // Fitted priors are cross-fitted by shot parity, so a decoded shot is
    // never among the shots its model was fitted to.
    let models = data::models(set, e, &args.prior, if args.fit_all { &all } else { &shots });
    for dname in &args.decoders {
        let t0 = Instant::now();
        let decs: Result<Vec<Dec>, String> = models.iter().map(|m| Dec::build(dname, m, e.distance)).collect();
        let decs = match decs {
            Ok(d) => d,
            Err(why) => {
                eprintln!("{} {dname}: not run: {why}", e.name);
                continue;
            }
        };
        let build_s = t0.elapsed().as_secs_f64();
        FORCED_SHOTS.store(0, std::sync::atomic::Ordering::Relaxed);
        LARGEST.store(0, std::sync::atomic::Ordering::Relaxed);
        let cpu0 = receipt::cpu_ns();
        let t1 = Instant::now();
        // Shot `s` of the selection is shot `skip + s` of the record, whose
        // parity picks its model (`skip` is even).
        let pr = predict(&shots, args.threads, || decs.iter().map(Dec::work).collect::<Vec<_>>(), |w, s, fired| {
            let k = (skip + s) % decs.len();
            decs[k].decode(&mut w[k], fired)
        })
        .unwrap_or_else(|(s, err)| panic!("{} shot {s}: {err}", e.name));
        let wall = t1.elapsed().as_secs_f64();
        if matches!(decs[0], Dec::Tn(_)) {
            // Where the decomposition time went, summed over every thread since the process began.
            let (ph, rot) = wai_quantum::quantum_svd::phases();
            eprintln!("svd phases so far (all threads): {ph:?}, {rot} rotations, {} Jacobi fallbacks", wai_quantum::quantum_memory_tn::jacobi_fallbacks());
        }
        // The process's CPU time over the decode, all threads: the decoder's
        // cost a shot whatever else shares the host.
        let cpu_per_shot = match (cpu0, receipt::cpu_ns()) {
            (Some(a), Some(b)) => Some(b.saturating_sub(a) as f64 / shots.len() as f64),
            _ => None,
        };
        let errors = pr.errors(&shots);
        if let Some(dir) = &args.predictions {
            std::fs::create_dir_all(dir).unwrap();
            let bytes: Vec<u8> = pr.flips.iter().map(|&f| f as u8).collect();
            std::fs::write(dir.join(format!("{}_{}_{dname}.b8", e.name.replace('/', "_"), args.prior)), bytes).unwrap();
        }
        let mut t = pr.nanos.clone();
        t.sort_unstable();
        let q = |f: f64| t[((t.len() as f64 - 1.0) * f).round() as usize];
        let mean = t.iter().sum::<u64>() as f64 / t.len() as f64;
        let mut row = base.clone();
        let extra = json!({
            "decoder": dname, "prior": args.prior, "errors": errors, "predictions_blake3": hex(&pr.digest()),
            "ns_mean": mean, "ns_p50": q(0.5), "ns_p99": q(0.99), "ns_max": t[t.len() - 1],
            "ns_per_round_mean": mean / f64::from(e.rounds), "build_s": build_s, "wall_s": wall, "threads": args.threads,
            "cpu_ns_per_shot": cpu_per_shot,
            "decoder_shape": decs[0].describe(),
        });
        for (k, v) in extra.as_object().unwrap() {
            row[k] = v.clone();
        }
        // Against each published decoder on the same shots: the shots only
        // this decoder got wrong, and the shots only the published one did.
        let mut discordant = serde_json::Map::new();
        for (name, pred) in &published {
            let only_ours = (0..shots.len()).filter(|&s| wrong(&pr.flips, s) && !wrong(pred, s)).count();
            let only_theirs = (0..shots.len()).filter(|&s| !wrong(&pr.flips, s) && wrong(pred, s)).count();
            discordant.insert(format!("published:{name}|published"), json!([only_ours, only_theirs]));
        }
        row["discordant"] = Value::Object(discordant);
        if matches!(decs[0], Dec::Tn(_) | Dec::Sp(_)) {
            row["decoder_shape"]["forced_shots"] = json!(FORCED_SHOTS.load(std::sync::atomic::Ordering::Relaxed));
            row["decoder_shape"]["largest"] = json!(LARGEST.load(std::sync::atomic::Ordering::Relaxed));
        }
        eprintln!("{} {} {}: {} errors of {} ({:.4}), {:.1} µs/shot mean, {:.2} µs/round, {:.1}s {}", e.name, dname, args.prior, errors, shots.len(), errors as f64 / shots.len() as f64, mean / 1e3, mean / 1e3 / f64::from(e.rounds), wall, row["decoder_shape"]);
        rows.push(row);
    }
    for (name, pred) in &published {
        let mut row = base.clone();
        row["decoder"] = json!(format!("published:{name}"));
        row["prior"] = json!("published");
        row["errors"] = json!((0..shots.len()).filter(|&s| wrong(pred, s)).count());
        rows.push(row);
    }
    rows
}

fn survey(a: &[String]) {
    if a.len() < 2 {
        usage();
    }
    let set = Set::parse(&a[0]).unwrap_or_else(|| usage());
    let dir = Path::new(&a[1]);
    let args = parse_args(&a[2..]);
    let mut exps = data::list(set, dir);
    if !args.only.is_empty() {
        exps.retain(|e| args.only.iter().any(|o| e.name.contains(o.as_str())));
    }
    eprintln!("{} experiments", exps.len());
    let mut out = std::fs::File::create(&args.out).unwrap();
    use std::io::Write;
    let (rows, wall, energy) = receipt::metered(|| {
        let mut all = Vec::new();
        for e in &exps {
            let rows = survey_one(set, e, &args);
            for row in &rows {
                writeln!(out, "{row}").unwrap();
            }
            all.push(rows);
        }
        all
    });
    eprintln!("done in {wall:.1}s");
    if let Some(path) = &args.receipt {
        let experiments: Vec<Value> = rows.iter().map(|r| compact(r)).collect();
        let flat: Vec<Value> = rows.into_iter().flatten().collect();
        let (fits, _) = fit_rows(set, &flat);
        let request = json!({
            "experiment": set.citation(),
            "source": set.source(),
            "set": set.name(),
            "prior": args.prior,
            "decoders": args.decoders,
            "shots_per_experiment": args.shots,
            "first_shot": args.skip,
            "experiments": exps.iter().map(|e| e.name.clone()).collect::<Vec<_>>(),
            "fit": set.fit_rule(),
            "method": {
                "records": "detection_events.b8 and the actual observable flips, read as published; each experiment's shots hashed (shots_blake3: quantum_memory::Shots::digest)",
                "decoders": "quantum_memory: mwpm = Matcher (exact minimum-weight perfect matching), corr = CorrMatcher, bm = BeliefMatcher (sum-product propagation, then matching), ml = MlDecoder (exact maximum likelihood); quantum_uf: uf; quantum_match: corr0; quantum_memory_tn: sp<chi> = SparseMl, tn<chi> = TnDecoder (e<K>: cutoff 10^-K; b: weighed by belief propagation); quantum_memory_stream: ufw, mwpmw = ParallelWindows",
                "published_decoders": "each published prediction file scored against the same shots",
                "shots": "shots first_shot to first_shot + shots_per_experiment of each experiment (all of them when not given)",
                "discordant": "per experiment, for each of this crate's decoders against each published decoder: [shots only this crate's decoder got wrong, shots only the published one got wrong]",
            },
            "threads": args.threads,
            "fit_priors_to_every_shot": args.fit_all,
            "source_commit": std::env::var("QEC_SOURCE_COMMIT").ok(),
            "host": host(),
        });
        let shapes: serde_json::Map<String, Value> = flat.iter().filter(|r| !r["decoder_shape"].is_null()).map(|r| (format!("{}|{}", r["decoder"].as_str().unwrap(), r["experiment"].as_str().unwrap()), r["decoder_shape"].clone())).collect();
        let result = json!({
            "experiments": experiments,
            "fits": fits,
            "decoder_shapes": shapes,
            "wall_seconds": wall,
            "energy_figure": energy,
            "peak_rss_bytes": receipt::peak_rss_bytes(),
        });
        receipt::seal(request, result, path);
        eprintln!("sealed {}", path.display());
    }
}

/// One experiment's rows as the receipt holds them: the experiment once,
/// then per decoder (`name|prior`) its errors, and for this crate's decoders
/// the predictions' digest, the mean and 99th-percentile nanoseconds a shot,
/// and against each published decoder the shots only one of the two got
/// wrong (`[this crate's, the published one's]`).
fn compact(rows: &[Value]) -> Value {
    let r0 = &rows[0];
    let mut errors = serde_json::Map::new();
    let mut digests = serde_json::Map::new();
    let mut ns = serde_json::Map::new();
    let mut discordant = serde_json::Map::new();
    for r in rows {
        let key = format!("{}|{}", r["decoder"].as_str().unwrap(), r["prior"].as_str().unwrap());
        errors.insert(key.clone(), r["errors"].clone());
        if !r["predictions_blake3"].is_null() {
            digests.insert(key.clone(), r["predictions_blake3"].clone());
            let cpu = r["cpu_ns_per_shot"].as_f64().map(|c| c.round() as u64);
            // The wall time of the whole decode over every thread, a shot.
            let wall = r["wall_s"].as_f64().map(|w| (w * 1e9 / r["shots"].as_f64().unwrap()).round() as u64);
            ns.insert(key.clone(), json!([r["ns_mean"].as_f64().unwrap().round() as u64, r["ns_p99"], cpu, wall]));
            if r["discordant"].as_object().is_some_and(|d| !d.is_empty()) {
                discordant.insert(key, r["discordant"].clone());
            }
        }
    }
    json!({
        "experiment": r0["experiment"], "distance": r0["distance"], "basis": r0["basis"], "rounds": r0["rounds"], "patch": r0["patch"],
        "detectors": r0["detectors"], "shots": r0["shots"], "event_rate": r0["event_rate"], "shots_blake3": r0["shots_blake3"],
        "errors": errors, "predictions_blake3": digests, "ns_mean_p99": ns, "discordant": discordant,
    })
}

/// A receipt's experiments as flat rows (`decoder`, `prior`, `distance`,
/// `patch`, `basis`, `rounds`, `shots`, `errors`), for [`fit_rows`].
fn flatten(doc: &Value) -> Vec<Value> {
    let mut out = Vec::new();
    for e in doc["result"]["experiments"].as_array().unwrap() {
        for (key, n) in e["errors"].as_object().unwrap() {
            let (dec, prior) = key.split_once('|').unwrap();
            out.push(json!({ "decoder": dec, "prior": prior, "distance": e["distance"], "patch": e["patch"], "basis": e["basis"], "rounds": e["rounds"], "shots": e["shots"], "errors": n }));
        }
    }
    out
}

/// What the receipt says about the machine.
fn host() -> Value {
    json!({ "os": std::env::consts::OS, "arch": std::env::consts::ARCH, "cores": std::thread::available_parallelism().map(|n| n.get()).unwrap_or(0) })
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

/// Every (decoder, prior) in the rows: per (distance, patch, basis) a decay
/// over the rounds the set's publication fits, the mean per distance, and Λ.
/// The fits as JSON, and as text.
fn fit_rows(set: Set, rows: &[Value]) -> (Value, String) {
    use std::fmt::Write;
    let mut text = String::new();
    let mut out = serde_json::Map::new();
    let mut keys: Vec<(String, String)> = rows.iter().map(|r| (r["decoder"].as_str().unwrap().to_string(), r["prior"].as_str().unwrap().to_string())).collect();
    keys.sort();
    keys.dedup();
    for (dec, prior) in keys {
        let mut groups: std::collections::BTreeMap<(u64, String, String), Vec<RoundPoint>> = std::collections::BTreeMap::new();
        for r in rows.iter().filter(|r| r["decoder"] == dec.as_str() && r["prior"] == prior.as_str() && set.fitted(r["rounds"].as_u64().unwrap() as u32)) {
            groups
                .entry((r["distance"].as_u64().unwrap(), r["patch"].as_str().unwrap().to_string(), r["basis"].as_str().unwrap().to_string()))
                .or_default()
                .push(RoundPoint { rounds: r["rounds"].as_u64().unwrap() as u32, shots: r["shots"].as_u64().unwrap(), errors: r["errors"].as_u64().unwrap() });
        }
        let mut per_d: std::collections::BTreeMap<u64, Vec<DecayFit>> = std::collections::BTreeMap::new();
        let mut members = Vec::new();
        for ((d, patch, basis), pts) in &groups {
            if let Some(f) = fit_decay(pts, 0) {
                let _ = writeln!(text, "  {dec:>40} {prior:>9} d{d} {patch:>8} {basis}: eps {:.5}% ± {:.5}% (stat {:.5}%), chi2/dof {:.2}, {} points", 100.0 * f.eps, 100.0 * f.sigma, 100.0 * f.sigma_stat, f.chi2_per_dof, f.points);
                members.push(json!({ "distance": d, "patch": patch, "basis": basis, "eps": f.eps, "sigma": f.sigma, "sigma_stat": f.sigma_stat, "prefactor": f.prefactor, "chi2_per_dof": f.chi2_per_dof, "points": f.points }));
                per_d.entry(*d).or_default().push(f);
            }
        }
        let mut lam_in = Vec::new();
        let mut dist = serde_json::Map::new();
        for (d, fits) in &per_d {
            let m = mean_rate(fits).unwrap();
            let _ = writeln!(text, "{dec:>42} {prior:>9} d{d}: eps {:.4}% (σ mean {:.4}%, indep {:.4}%, rms {:.4}%), {} fits", 100.0 * m.eps, 100.0 * m.sigma_mean, 100.0 * m.sigma_independent, 100.0 * m.sigma_rms, m.members);
            dist.insert(format!("d{d}"), json!({ "eps": m.eps, "sigma_mean": m.sigma_mean, "sigma_independent": m.sigma_independent, "sigma_rms": m.sigma_rms, "fits": m.members }));
            lam_in.push((*d as u32, m.eps, m.sigma_independent));
        }
        let lambda = fit_lambda(&lam_in);
        if let Some((l, s)) = lambda {
            let _ = writeln!(text, "{dec:>42} {prior:>9} Lambda {l:.4} ± {s:.4}");
        }
        out.insert(format!("{dec}|{prior}"), json!({ "decoder": dec, "prior": prior, "members": members, "distances": dist, "lambda": lambda.map(|(l, s)| json!({ "value": l, "sigma": s })) }));
    }
    (Value::Object(out), text)
}

/// Fit the rows of the files given.
fn fit(a: &[String]) {
    let set = Set::parse(&a[0]).unwrap_or_else(|| usage());
    let mut rows: Vec<Value> = Vec::new();
    for f in &a[1..] {
        let t = std::fs::read_to_string(f).unwrap();
        if f.ends_with(".json") {
            let doc: Value = serde_json::from_str(&t).unwrap();
            rows.extend(flatten(&doc));
        } else {
            rows.extend(t.lines().map(|l| serde_json::from_str::<Value>(l).unwrap()));
        }
    }
    print!("{}", fit_rows(set, &rows).1);
}