open-standards/wai/wai-quantum-cli/src/qec.rs

//! `wai-quantum qec`: error correction, as the job service's MCP serves it
//! ([`wai_quantum::engine::qec`]): a code shown with its distance certified,
//! a syndrome or a circuit's detectors decoded, and a threshold measured,
//! each answer signed with this user's key (`--json` prints it sealed, and
//! `wai-quantum verify` checks it).
//!
//! ```text
//! wai-quantum qec code color:7                      # its checks, k, the decoders it takes, its certified distance
//! wai-quantum qec code gross --json > cert.json     # sealed
//! wai-quantum qec decode --code surface:5 --error 3,12
//! wai-quantum qec decode --memory 5,5,0.003 --shots 20000 --decoder correlated
//! wai-quantum qec decode --model run.dem --detectors 4,9
//! wai-quantum qec threshold --distances 3,5,7 --p 0.002,0.004,0.006 --shots 20000
//! ```
//!
//! The research forms stay where they were: `wai-quantum lab qec`.

use serde_json::{Value, json};
use std::path::PathBuf;
use wai_quantum::engine::{job, qec};

#[derive(clap::Subcommand)]
pub enum Cmd {
    /// A code: its checks and stabilizers, k, the decoders it takes, and its distance certificate
    Code {
        /// surface:D, toric:L, steane, color:D, bb:72, bb:90, bb:108, bb:144 (gross) or bb:288
        name: String,
        /// The heaviest support the distance search tries (default the code's nominal distance)
        #[arg(long)]
        bound: Option<u64>,
        /// Show the code without certifying its distance
        #[arg(long)]
        no_certify: bool,
        /// Check this certificate (a sealed answer or the certificate alone) against the code
        #[arg(long, value_name = "FILE")]
        check: Option<PathBuf>,
        #[command(flatten)]
        out: Out,
    },
    /// Decode: a code's syndrome or error, or a circuit's or a model's detectors
    Decode {
        /// The code, for decoding at code capacity
        #[arg(long)]
        code: Option<String>,
        /// The violated checks, as 0,4,9
        #[arg(long, value_delimiter = ',')]
        syndrome: Vec<u64>,
        /// The flipped qubits, as 3,12 (its syndrome is decoded and its residual judged)
        #[arg(long, value_delimiter = ',')]
        error: Vec<u64>,
        /// The rotated surface-code memory experiment, as DISTANCE,ROUNDS,P
        #[arg(long, value_delimiter = ',')]
        memory: Vec<f64>,
        /// A circuit in the circuit text
        #[arg(long, value_name = "FILE")]
        circuit: Option<PathBuf>,
        /// A detector error model in its text
        #[arg(long, value_name = "FILE")]
        model: Option<PathBuf>,
        /// The detectors that fired, as 4,9
        #[arg(long, value_delimiter = ',')]
        detectors: Vec<u64>,
        /// Sample this many shots and decode each
        #[arg(long)]
        shots: Option<u64>,
        /// matching, union_find, color, relay_bp or search at code capacity; matching, correlated, union_find, bp_osd or window at circuit level
        #[arg(long)]
        decoder: Option<String>,
        /// window: COMMIT,BUFFER rounds
        #[arg(long, value_delimiter = ',')]
        window: Vec<u64>,
        /// The sampling seed
        #[arg(long)]
        seed: Option<u64>,
        #[command(flatten)]
        out: Out,
    },
    /// Logical failure rates over distances and physical rates, and the code model they fit
    Threshold {
        /// surface (circuit level), or color, toric or surface_capacity (code capacity)
        #[arg(long, default_value = "surface")]
        family: String,
        /// Odd distances, as 3,5,7
        #[arg(long, value_delimiter = ',')]
        distances: Vec<u64>,
        /// Physical rates, as 0.002,0.004
        #[arg(long = "p", value_delimiter = ',')]
        rates: Vec<f64>,
        /// Shots at each point
        #[arg(long)]
        shots: Option<u64>,
        /// matching, correlated or union_find (surface)
        #[arg(long)]
        decoder: Option<String>,
        /// Rounds (surface; default the distance)
        #[arg(long)]
        rounds: Option<u64>,
        #[arg(long)]
        seed: Option<u64>,
        #[command(flatten)]
        out: Out,
    },
}

#[derive(clap::Args)]
pub struct Out {
    /// Print the answer sealed with this user's key, as JSON (`wai-quantum verify` checks it)
    #[arg(long)]
    pub json: bool,
    #[command(flatten)]
    pub key: super::KeyArgs,
}

fn read(p: &PathBuf) -> Result<String, String> {
    std::fs::read_to_string(p).map_err(|e| format!("{}: {e}", p.display()))
}

/// The request a command makes of the engine.
pub fn request(cmd: &Cmd) -> Result<(&'static str, Value), String> {
    Ok(match cmd {
        Cmd::Code { name, bound, no_certify, check, .. } => {
            let mut r = json!({ "code": name });
            if let Some(b) = bound {
                r["bound"] = json!(b);
            }
            if *no_certify {
                r["certify"] = json!(false);
            }
            if let Some(p) = check {
                let v: Value = serde_json::from_str(&read(p)?).map_err(|e| format!("{}: {e}", p.display()))?;
                r["certificate"] = v.pointer("/result/certificate").cloned().unwrap_or(v);
            }
            ("code", r)
        }
        Cmd::Decode { code, syndrome, error, memory, circuit, model, detectors, shots, decoder, window, seed, .. } => {
            let mut r = json!({});
            if let Some(c) = code {
                r["code"] = json!(c);
                if !syndrome.is_empty() {
                    r["syndrome"] = json!(syndrome);
                }
                if !error.is_empty() {
                    r["error"] = json!(error);
                }
            } else {
                match (memory.as_slice(), circuit, model) {
                    ([d, rounds, p], None, None) => r["memory"] = json!({ "distance": *d as u64, "rounds": *rounds as u64, "p": p }),
                    ([], Some(c), None) => r["circuit"] = json!(read(c)?),
                    ([], None, Some(m)) => r["model"] = json!(read(m)?),
                    _ => return Err("give --code (with --syndrome or --error), or one of --memory DISTANCE,ROUNDS,P, --circuit FILE and --model FILE (with --detectors or --shots)".into()),
                }
                if !detectors.is_empty() {
                    r["detectors"] = json!(detectors);
                }
                if let Some(s) = shots {
                    r["shots"] = json!(s);
                }
                if let [c, b] = window.as_slice() {
                    r["window"] = json!({ "commit": c, "buffer": b });
                }
            }
            if let Some(d) = decoder {
                r["decoder"] = json!(d);
            }
            if let Some(s) = seed {
                r["seed"] = json!(s);
            }
            ("decode", r)
        }
        Cmd::Threshold { family, distances, rates, shots, decoder, rounds, seed, .. } => {
            let mut r = json!({ "family": family });
            if !distances.is_empty() {
                r["distances"] = json!(distances);
            }
            if !rates.is_empty() {
                r["p"] = json!(rates);
            }
            for (k, v) in [("shots", shots), ("rounds", rounds), ("seed", seed)] {
                if let Some(x) = v {
                    r[k] = json!(x);
                }
            }
            if let Some(d) = decoder {
                r["decoder"] = json!(d);
            }
            ("threshold", r)
        }
    })
}

/// An answer as a reader takes it in.
pub fn report(op: &str, r: &Value) -> String {
    let mut s = String::new();
    let num = |v: &Value| v.as_f64().map_or_else(|| v.to_string(), |x| format!("{x}"));
    match op {
        "code" => {
            let c = &r["code"];
            s += &format!("code        {} ({}): n = {}, k = {}, {} checks and {} stabilizers\n", c["name"].as_str().unwrap_or(""), c["family"].as_str().unwrap_or(""), c["n"], c["k"], c["checks"].as_array().map_or(0, Vec::len), c["stabs"].as_array().map_or(0, Vec::len));
            s += &format!("decoders    {}\n", c["decoders"].as_array().into_iter().flatten().filter_map(Value::as_str).collect::<Vec<_>>().join(", "));
            if let Some(cert) = r.get("certificate") {
                for side in ["x", "z"] {
                    let x = &cert[side];
                    s += &format!("{side} sector    distance {}, no logical below {}, witness {}, {} search nodes\n", num(&x["distance"]), num(&x["no_logical_below"]), x["witness"], x["nodes"]);
                }
                s += &format!("certificate {}\n", cert["reading"].as_str().unwrap_or(""));
            }
            if let Some(ch) = r.get("check") {
                s += &format!("check       {}: {}\n", if ch["verified"] == true { "verified" } else { "NOT VERIFIED" }, ch["why"].as_str().unwrap_or(""));
            }
        }
        "decode" => {
            if r.get("correction").is_some() {
                s += &format!("decoder     {} on {}\n", r["decoder"].as_str().unwrap_or(""), r["code"].as_str().unwrap_or(""));
                s += &format!("syndrome    {}\n", r["syndrome"]);
                s += &format!("correction  {} (weight {}), {}\n", r["correction"], r["weight"], if r["explains"] == true { "explains the syndrome" } else { "DOES NOT explain the syndrome" });
                if let Some(res) = r.get("residual") {
                    s += &format!("residual    {}: {}\n", res["qubits"], res["reading"].as_str().unwrap_or(""));
                }
            } else {
                s += &format!("model       {} detectors, {} observables, {} mechanisms; decoder {}\n", r["detectors"], r["observables"], r["mechanisms"], r["decoder"].as_str().unwrap_or(""));
                if let Some(p) = r.get("predicted") {
                    s += &format!("fired       {}\npredicted   observables flipped: {p}\n", r["fired"]);
                }
                if let Some(f) = r.get("failures") {
                    s += &format!("failures    {f} of {} shots: {:.3e} ± {:.1e}\n", r["shots"], r["rate"].as_f64().unwrap_or(0.0), r["stderr"].as_f64().unwrap_or(0.0));
                }
            }
        }
        _ => {
            s += &format!("{} ({}), decoder {}\n", r["family"].as_str().unwrap_or(""), r["noise"].as_str().unwrap_or(""), r["decoder"].as_str().unwrap_or(""));
            s += &format!("{:>3} {:>9} {:>9} {:>9} {:>12}\n", "d", "p", "shots", "failures", "rate");
            for p in r["points"].as_array().into_iter().flatten() {
                s += &format!("{:>3} {:>9} {:>9} {:>9} {:>12}\n", p["distance"], num(&p["p"]), p["shots"], p["failures"], format!("{:.3e}", p["rate"].as_f64().unwrap_or(0.0)));
            }
            if let Some(f) = r.get("fit").filter(|f| !f.is_null()) {
                s += &format!("fit         {}: a = {:.4}, p* = {:.5}\n", f["model"].as_str().unwrap_or(""), f["prefactor"].as_f64().unwrap_or(0.0), f["threshold"].as_f64().unwrap_or(0.0));
            }
            for b in r["below_threshold"].as_array().into_iter().flatten() {
                s += &format!("p = {}: {}\n", num(&b["p"]), if b["distance_helps"] == true { "every larger distance fails less" } else { "a larger distance does not fail less" });
            }
        }
    }
    s
}

pub fn main(cmd: Cmd) {
    let r = (|| -> Result<(), String> {
        let (op, req) = request(&cmd)?;
        let result = qec::call(op, &req)?;
        let out = match &cmd {
            Cmd::Code { out, .. } | Cmd::Decode { out, .. } | Cmd::Threshold { out, .. } => out,
        };
        let k = super::load_key(&out.key);
        let body = qec::seal(&req, result, &k.key);
        if out.json {
            println!("{body}");
            return Ok(());
        }
        print!("{}", report(op, &body["result"]));
        let ok = job::verify(&body).map(|(ok, _)| ok).unwrap_or(false);
        let signer = body["result"]["receipt"]["key"].as_str().and_then(job::did_of).unwrap_or_default();
        println!("receipt     {}, signed by {signer}", if ok { "verifies" } else { "DOES NOT VERIFY" });
        Ok(())
    })();
    r.unwrap_or_else(|e| super::die(&e));
}

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

    fn cmd(args: &[&str]) -> Cmd {
        match super::super::Cli::try_parse_from(args).unwrap().command {
            super::super::Command::Qec { cmd } => cmd,
            _ => panic!("not qec"),
        }
    }

    #[test]
    fn each_verb_reads_its_arguments_and_answers() {
        let key = wai_quantum::ed25519_dalek::SigningKey::from_bytes(&[5; 32]);
        for (args, op, want) in [
            (vec!["wai-quantum", "qec", "code", "color:5"], "code", "distance 5"),
            (vec!["wai-quantum", "qec", "decode", "--code", "surface:3", "--error", "4"], "decode", "the logical state is unharmed"),
            (vec!["wai-quantum", "qec", "decode", "--memory", "3,3,0.004", "--shots", "500", "--decoder", "correlated"], "decode", "of 500 shots"),
            (vec!["wai-quantum", "qec", "threshold", "--family", "color", "--distances", "3,5", "--p", "0.03", "--shots", "300"], "threshold", "p = 0.03"),
        ] {
            let c = cmd(&args);
            let (o, req) = request(&c).unwrap();
            assert_eq!(o, op);
            let body = qec::seal(&req, qec::call(o, &req).unwrap(), &key);
            assert!(job::verify(&body).unwrap().0, "{args:?}");
            let text = report(o, &body["result"]);
            assert!(text.contains(want), "{args:?}:\n{text}");
        }
        assert!(request(&cmd(&["wai-quantum", "qec", "decode", "--detectors", "1"])).is_err());
    }
}