open-standards/wai/wai-quantum/tests/tn_vs_sv.rs

//! Two independent methods, one answer: an amplitude contracted as a tensor
//! network equals the state vector's, for every bitstring tried and for a
//! batch of open outputs.
#![cfg(all(feature = "quantum_tn", feature = "quantum_sv"))]

use wai_quantum::gates::Gate;
use wai_quantum::quantum_sv::StateVector;
use wai_quantum::quantum_tn::{contract_plan, find_path, Families, Network, PathOptions};

fn circuit(n: u32, depth: u32) -> Vec<Gate> {
    let mut g = Vec::new();
    let mut a = 0.2f64;
    for layer in 0..depth {
        for q in 0..n {
            a += 0.613;
            g.push(match (q * 7 + layer) % 4 {
                0 => Gate::rx(q, a),
                1 => Gate::sqrt_w(q),
                2 => Gate::ry(q, a),
                _ => Gate::sqrt_x(q),
            });
        }
        let mut q = layer % 2;
        while q + 1 < n {
            g.push(Gate::fsim(q, q + 1, 1.2 + 0.1 * q as f64, 0.5));
            q += 2;
        }
        g.push(Gate::cz(0, n - 1));
    }
    g
}

#[test]
fn tensor_network_amplitudes_equal_the_state_vector() {
    let n = 18;
    let gates = circuit(n, 8);
    let mut sv = StateVector::zero(n).unwrap();
    sv.run(&gates, 2).unwrap();
    let opts = PathOptions { trials: 16, seed: 1, max_size: Some(4096.0), reconfigure: 6, families: Families::Both };
    for x in [0usize, 1, 77_777, 131_071, 262_143] {
        let bits: Vec<u8> = (0..n).map(|q| ((x >> (n - 1 - q)) & 1) as u8).collect();
        let mut net = Network::amplitude(n, &gates, &bits, &[]).unwrap();
        net.simplify();
        let plan = find_path(&net.shape(), &net.dims, &net.open, &opts);
        let t = contract_plan(&net, &plan).unwrap().data[0];
        let s = sv.amplitude(x);
        assert!((t.re - s.re).abs() < 1e-12 && (t.im - s.im).abs() < 1e-12, "x {x}: {t:?} vs {s:?}");
    }
    // A batch: qubits 4 and 11 open, the rest fixed to x = 77 777.
    let x = 77_777usize;
    let bits: Vec<u8> = (0..n).map(|q| ((x >> (n - 1 - q)) & 1) as u8).collect();
    let mut net = Network::amplitude(n, &gates, &bits, &[4, 11]).unwrap();
    net.simplify();
    let plan = find_path(&net.shape(), &net.dims, &net.open, &opts);
    let t = contract_plan(&net, &plan).unwrap();
    for j in 0..4usize {
        let (b4, b11) = ((j >> 1) & 1, j & 1);
        let y = (x & !(1 << (n - 1 - 4)) & !(1 << (n - 1 - 11))) | (b4 << (n - 1 - 4)) | (b11 << (n - 1 - 11));
        let s = sv.amplitude(y);
        assert!((t.data[j].re - s.re).abs() < 1e-12 && (t.data[j].im - s.im).abs() < 1e-12, "batch {j}");
    }
}