open-standards/wai/wai-quantum/examples/tn_path.rs

//! Contraction trees for a circuit's amplitude, from the command line.
//!
//! Reads a circuit as JSON lines, one gate per line:
//! `{"q":[0,1],"m":[re,im, re,im, …]}` (row-major, rows = outputs, the first
//! qubit most significant). Writes the simplified network's shape (for any
//! other tool to search the same network) and the best plan found.
//!
//! ```text
//! cargo run --release --example tn_path --features quantum_tn -- CIRCUIT N TRIALS SEED [LOG2_CAP] [SHAPE_OUT] [--contract]
//! ```
use wai_quantum::quantum_tn::{contract_plan, find_path, Families, Gate, Network, PathOptions};

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let text = std::fs::read_to_string(&a[1]).expect("read circuit");
    let n: u32 = a[2].parse().unwrap();
    let trials: usize = a[3].parse().unwrap();
    let seed: u64 = a[4].parse().unwrap();
    let cap: Option<f64> = a.get(5).filter(|s| *s != "-").map(|s| 2f64.powf(s.parse().unwrap()));
    let shape_out = a.get(6).filter(|s| *s != "-");
    let gates: Vec<Gate> = text
        .lines()
        .filter(|l| !l.trim().is_empty())
        .map(|l| {
            let v: serde_json::Value = serde_json::from_str(l).expect("json gate");
            let qubits: Vec<u32> = v["q"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect();
            let m: Vec<f64> = v["m"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect();
            let matrix = m.chunks(2).map(|c| wai_quantum::quantum_tn::complex(c[0], c[1])).collect();
            Gate { qubits, matrix }
        })
        .collect();
    let mut net = Network::amplitude(n, &gates, &vec![0; n as usize], &[]).expect("network");
    let raw = net.tensors.len();
    net.simplify();
    let shape = net.shape();
    if let Some(path) = shape_out {
        let inputs: Vec<String> = shape.iter().map(|s| format!("{s:?}")).collect();
        let dims: Vec<String> = net.dims.iter().map(|d| d.to_string()).collect();
        std::fs::write(path, format!("{{\"inputs\":[{}],\"output\":{:?},\"dims\":[{}]}}\n", inputs.join(","), net.open, dims.join(","))).unwrap();
    }
    let t0 = std::time::Instant::now();
    let families = match std::env::var("TN_FAMILIES").as_deref() {
        Ok("greedy") => Families::Greedy,
        Ok("partition") => Families::Partition,
        _ => Families::Both,
    };
    let reconfigure = std::env::var("TN_RECONF").ok().and_then(|v| v.parse().ok()).unwrap_or(8);
    let plan = find_path(&shape, &net.dims, &net.open, &PathOptions { trials, seed, max_size: cap, reconfigure, families });
    let dt = t0.elapsed().as_secs_f64();
    println!(
        "tensors {raw} -> {}  log10(flops) {:.3}  log2(max size) {:.1}  slices {}  search {:.1}s",
        shape.len(),
        plan.total_flops().log10(),
        plan.per_slice.max_size.log2(),
        plan.slices,
        dt
    );
    if a.iter().any(|x| x == "--contract") {
        let t0 = std::time::Instant::now();
        let t = contract_plan(&net, &plan).expect("contract");
        println!("amplitude {:.12e} {:+.12e}i  contract {:.1}s", t.data[0].re, t.data[0].im, t0.elapsed().as_secs_f64());
    }
}