wai-quantum 0.6.0

//! Sampling on the fabric: the kernel a quantum processor offers, and how a
//! member that returns measurement counts rather than a state is certified,
//! placed and receipted.
//!
//! - **The kernel** ([`sample_kernel`]): run a circuit from `|0…0⟩`, measure
//!   every qubit in the computational basis, `shots` times, and return the
//!   count of each basis state. Its reference is the exact distribution of the
//!   CPU's integer simulation ([`reference_distribution`]); the CPU reference
//!   sampler ([`CpuSampler`]) draws from it with the core's seeded integer
//!   sampler, so its counts are the same on every machine.
//! - **A member that samples** implements [`Sampler`]. Its certificate is
//!   statistical ([`Parity::Statistical`](super::Parity::Statistical)): over a corpus of circuits shaped
//!   for quantum processors ([`QpuCorpus`]), the worst classical fidelity
//!   between its counts and the reference distribution, and the shots each
//!   case took ([`certify_sampler`]).
//! - **Placement** is [`place`](super::place) with a statistical requirement:
//!   a bit-exact reference meets every requirement, so the CPU reference is
//!   placed unless the caller offers only other members; a simulated
//!   processor's certificate ([`Provenance::Simulator`]) ranks after every
//!   silicon one.
//! - **A receipt** binds [`counts_hash`]: the circuit's hash, the shots and the
//!   counts.

use super::{Certificate, Member, MemberError, Presence, Provenance, REFERENCE, RunMeter, Window, certify_sampling, detect_local, lanes};
use crate::quantum::{BaseGate, Circuit, Gate, Gateset, QuantumError, circuit_hash};
use crate::quantum_energy::EnergyClass;
use std::future::Future;
use std::pin::Pin;

/// The sampling kernel of the frozen gateset.
pub const KERNEL_CIRCUIT_SAMPLE: &str = "wai.quantum.circuit/sample";

/// The sampling kernel of circuit2.
pub const KERNEL_CIRCUIT2_SAMPLE: &str = "wai.quantum.circuit2/sample";

/// The sampling kernel of `gs`.
pub const fn sample_kernel(gs: Gateset) -> &'static str {
    match gs {
        Gateset::V1 => KERNEL_CIRCUIT_SAMPLE,
        _ => KERNEL_CIRCUIT2_SAMPLE,
    }
}

/// Counts over the `2^n` basis states, by index: qubit `q` is bit `q` of the
/// index, as in a state vector.
pub type Counts = Vec<u64>;

/// What a [`Sampler`] returns: the counts, and the microjoules and class its
/// meter read (zero and `None` when unmetered).
pub type SampleFuture<'r> = Pin<Box<dyn Future<Output = Result<(Counts, u64, Option<EnergyClass>), MemberError>> + 'r>>;

/// A member that measures: a quantum processor, a simulator of one, or the
/// CPU reference sampler.
pub trait Sampler {
    /// The id of the member this sampler is.
    fn member_id(&self) -> String;

    /// `shots` measurements of every qubit of `c`, run from `|0…0⟩` with
    /// `implementation`, inside a window of `meter`: one count per basis state,
    /// summing to `shots`.
    fn sample<'r>(&'r self, implementation: &'r str, c: &'r Circuit, shots: u64, meter: &'r mut dyn RunMeter) -> SampleFuture<'r>;
}

impl<S: Sampler + ?Sized> Sampler for &S {
    fn member_id(&self) -> String {
        (**self).member_id()
    }

    fn sample<'r>(&'r self, implementation: &'r str, c: &'r Circuit, shots: u64, meter: &'r mut dyn RunMeter) -> SampleFuture<'r> {
        (**self).sample(implementation, c, shots, meter)
    }
}

/// The probability of each basis state of `c` run from `|0…0⟩`: the CPU's
/// integer state, its weights `|amp|²` divided by their sum. Refuses what the
/// simulator refuses.
pub fn reference_distribution(c: &Circuit) -> Result<Vec<f64>, QuantumError> {
    let weights = c.simulate()?.prob_weights();
    let total: i128 = weights.iter().sum();
    if total <= 0 {
        return Ok(vec![0.0; weights.len()]);
    }
    let t = total as f64;
    Ok(weights.iter().map(|&w| w as f64 / t).collect())
}

/// The CPU reference as a [`Sampler`]: the exact integer state, sampled with
/// the core's seeded sampler (`StateVector::sample_shots`), so the same
/// circuit, seed and shots give the same counts on every machine.
#[derive(Clone, Copy, Debug, Default)]
pub struct CpuSampler {
    pub seed: u64,
}

impl Sampler for CpuSampler {
    fn member_id(&self) -> String {
        detect_local().id
    }

    fn sample<'r>(&'r self, implementation: &'r str, c: &'r Circuit, shots: u64, meter: &'r mut dyn RunMeter) -> SampleFuture<'r> {
        Box::pin(async move {
            if implementation != REFERENCE {
                return Err(MemberError::Other(format!("the CPU reference sampler does not run implementation {implementation:?}")));
            }
            let mut window = Window::new(meter);
            window.open();
            let counts = c.simulate().map_err(MemberError::Circuit)?.sample_shots(self.seed, shots);
            let (joules_micro, class) = window.close();
            Ok((counts, joules_micro, class))
        })
    }
}

/// The CPU reference's sampling entries: one per gateset, bit-exact by
/// definition ([`Certificate::reference`]).
pub fn local_sampling_member() -> Member {
    let local = detect_local();
    let kernels = Gateset::ALL
        .into_iter()
        .map(|gs| super::KernelEntry {
            kernel: sample_kernel(gs).into(),
            implementation: REFERENCE.into(),
            certificate: Certificate::reference(sample_kernel(gs), &local.id),
            energy: None,
        })
        .collect();
    Member { kernels, ..local }
}

/// A certification corpus for samplers, shaped for quantum processors: few
/// qubits, shallow, at most one control per gate and phases no finer than
/// `P(4)`, so a processor's compiler takes it as written in either gateset.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct QpuCorpus {
    pub seed: u64,
    /// Circuits in all: a Bell pair, GHZ on three qubits, GHZ on
    /// `max_qubits`, then random circuits.
    pub cases: usize,
    /// The widest circuit, at least 3.
    pub max_qubits: u8,
    /// Ops per random circuit.
    pub depth: usize,
}

/// The default sampler corpus: qpu-v1, seed 7, 16 circuits of up to 5 qubits,
/// 12 ops each.
pub const QPU_DEFAULT: QpuCorpus = QpuCorpus { seed: 7, cases: 16, max_qubits: 5, depth: 12 };

fn next(s: &mut u64) -> u64 {
    *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *s;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

impl QpuCorpus {
    /// `"qpu-v1 seed=<seed> cases=<cases> qubits=<max_qubits> depth=<depth>"`,
    /// enough to draw the corpus again.
    pub fn generator(&self) -> String {
        format!("qpu-v1 seed={} cases={} qubits={} depth={}", self.seed, self.cases, self.max_qubits, self.depth)
    }

    /// The circuits, in order, in `gs`. Both gatesets draw the same ops.
    pub fn circuits(&self, gs: Gateset) -> Vec<Circuit> {
        let one = |base, target| Gate { base, controls: Vec::new(), target, param: 0 };
        let cx = |control, target| Gate { base: BaseGate::X, controls: vec![control], target, param: 0 };
        let ghz = |n: u8| {
            let mut c = Circuit::with_gateset(n, gs);
            c.ops.push(one(BaseGate::H, 0));
            (1..n).for_each(|q| c.ops.push(cx(q - 1, q)));
            c
        };
        let mut out = vec![ghz(2), ghz(3), ghz(self.max_qubits.max(3))];
        let bases = [BaseGate::H, BaseGate::X, BaseGate::Y, BaseGate::Z, BaseGate::S, BaseGate::Sdg, BaseGate::T, BaseGate::Tdg, BaseGate::P];
        let mut s = self.seed;
        while out.len() < self.cases {
            let n = 2 + (next(&mut s) % u64::from(self.max_qubits.max(3) - 1)) as u8;
            let mut c = Circuit::with_gateset(n, gs);
            for _ in 0..self.depth {
                let target = (next(&mut s) % u64::from(n)) as u8;
                if next(&mut s).is_multiple_of(3) {
                    let control = (target + 1 + (next(&mut s) % u64::from(n - 1)) as u8) % n;
                    let base = if next(&mut s).is_multiple_of(2) { BaseGate::X } else { BaseGate::Z };
                    c.ops.push(Gate { base, controls: vec![control], target, param: 0 });
                } else {
                    let base = bases[(next(&mut s) % bases.len() as u64) as usize];
                    let param = if base == BaseGate::P { 1 + (next(&mut s) % 4) as u16 } else { 0 };
                    c.ops.push(Gate { base, controls: Vec::new(), target, param });
                }
            }
            out.push(c);
        }
        out.truncate(self.cases);
        out
    }
}

/// A case's input bytes for a sampling certificate: the circuit's bytes
/// ([`lanes::circuit_bytes`]) and the shots, little-endian.
pub fn case_input(c: &Circuit, shots: u64) -> Vec<u8> {
    let mut o = lanes::circuit_bytes(c);
    o.extend_from_slice(&shots.to_le_bytes());
    o
}

/// Checks `counts` for a run of `c` with `shots`: one count per basis state,
/// summing to `shots`.
pub fn check_counts(c: &Circuit, shots: u64, counts: &[u64]) -> Result<(), MemberError> {
    let want = 1usize.checked_shl(u32::from(c.n_qubits)).unwrap_or(0);
    if counts.len() != want {
        return Err(MemberError::Other(format!("{} counts for a circuit of {} qubits; {want} expected", counts.len(), c.n_qubits)));
    }
    let total = counts.iter().try_fold(0u64, |a, &b| a.checked_add(b));
    if total != Some(shots) {
        return Err(MemberError::Other(format!("the counts sum to {total:?}, not {shots}")));
    }
    Ok(())
}

/// The statistical certificate of `sampler` running `implementation` for
/// [`sample_kernel`]`(gs)`, over `corpus` with `shots` per circuit: each case
/// folds in the circuit and shots, the reference distribution and the counts
/// ([`certify_sampling`]); the parity is the worst case's classical fidelity.
/// The certificate is earned under `provenance`. Counts of the wrong shape
/// stop the certification.
pub async fn certify_sampler<S: Sampler + ?Sized>(sampler: &S, implementation: &str, gs: Gateset, corpus: QpuCorpus, shots: u64, provenance: Provenance) -> Result<Certificate, MemberError> {
    let mut cases = Vec::new();
    for c in corpus.circuits(gs) {
        let reference = reference_distribution(&c).map_err(MemberError::Circuit)?;
        let (counts, _, _) = sampler.sample(implementation, &c, shots, &mut super::Unmetered).await?;
        check_counts(&c, shots, &counts)?;
        cases.push((case_input(&c, shots), reference, counts));
    }
    Ok(certify_sampling(sample_kernel(gs), implementation, &sampler.member_id(), &corpus.generator(), cases).with_provenance(provenance))
}

const DOMAIN_COUNTS: &[u8] = b"wai:quantum-fabric-counts\x01";

/// What a receipt over a sampled result binds: BLAKE3 of the domain, the
/// circuit's hash, the shots and every count, little-endian. `None` for a
/// circuit with no hash (one its gateset does not allow).
pub fn counts_hash(c: &Circuit, shots: u64, counts: &[u64]) -> Option<[u8; 32]> {
    let mut h = blake3::Hasher::new();
    h.update(DOMAIN_COUNTS);
    h.update(&circuit_hash(c).ok()?);
    h.update(&shots.to_le_bytes());
    h.update(&(counts.len() as u64).to_le_bytes());
    for k in counts {
        h.update(&k.to_le_bytes());
    }
    Some(*h.finalize().as_bytes())
}

/// A member's sampling entries are present: its id and presence, for
/// [`place`](super::place).
pub fn sampling_member(id: String, class: super::FabricClass, detail: String, entries: Vec<super::KernelEntry>) -> Member {
    Member { id, class, presence: Presence::Detected, kernels: entries, detail }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantum_fabric::{FabricClass, KernelEntry, Parity, place, ready};

    #[test]
    fn the_sample_kernels_name_their_gateset() {
        assert_eq!(sample_kernel(Gateset::V1), KERNEL_CIRCUIT_SAMPLE);
        assert_eq!(sample_kernel(Gateset::V2), KERNEL_CIRCUIT2_SAMPLE);
        assert_ne!(KERNEL_CIRCUIT_SAMPLE, crate::quantum_fabric::KERNEL_CIRCUIT_APPLY);
    }

    #[test]
    fn the_reference_distribution_is_the_states_weights_normalised() {
        let mut bell = Circuit::new(2);
        bell.h(0).cx(0, 1);
        let p = reference_distribution(&bell).unwrap();
        assert_eq!(p.len(), 4);
        assert!((p[0] - 0.5).abs() < 1e-9 && (p[3] - 0.5).abs() < 1e-9 && p[1] == 0.0 && p[2] == 0.0, "{p:?}");
        let qft = Circuit::qft(4);
        let p = reference_distribution(&qft).unwrap();
        assert!((p.iter().sum::<f64>() - 1.0).abs() < 1e-12);
        assert!(p.iter().all(|&x| (x - 1.0 / 16.0).abs() < 1e-6), "the QFT of |0⟩ is uniform");
    }

    #[test]
    fn the_qpu_corpus_is_small_shallow_singly_controlled_and_the_same_in_both_gatesets() {
        let v1 = QPU_DEFAULT.circuits(Gateset::V1);
        let v2 = QPU_DEFAULT.circuits(Gateset::V2);
        assert_eq!(v1.len(), QPU_DEFAULT.cases);
        assert_eq!(v1.iter().map(|c| (&c.ops, c.n_qubits)).collect::<Vec<_>>(), v2.iter().map(|c| (&c.ops, c.n_qubits)).collect::<Vec<_>>());
        for c in v1.iter().chain(&v2) {
            assert!(c.validate().is_ok(), "{c:?}");
            assert!((2..=QPU_DEFAULT.max_qubits).contains(&c.n_qubits));
            assert!(c.ops.iter().all(|g| g.controls.len() <= 1 && g.param <= 4));
        }
        assert_eq!(QPU_DEFAULT.circuits(Gateset::V2), v2, "deterministic");
        assert_eq!(QPU_DEFAULT.generator(), "qpu-v1 seed=7 cases=16 qubits=5 depth=12");
    }

    #[test]
    fn the_cpu_sampler_is_deterministic_and_certifies_near_one() {
        let s = CpuSampler { seed: 0x5a };
        let qft = Circuit::qft(3);
        let a = ready(s.sample(REFERENCE, &qft, 1000, &mut crate::quantum_fabric::Unmetered)).unwrap().unwrap().0;
        let b = ready(s.sample(REFERENCE, &qft, 1000, &mut crate::quantum_fabric::Unmetered)).unwrap().unwrap().0;
        assert_eq!(a, b);
        assert_eq!(a.iter().sum::<u64>(), 1000);
        assert!(ready(s.sample("other", &qft, 10, &mut crate::quantum_fabric::Unmetered)).unwrap().is_err());
        for gs in Gateset::ALL {
            let cert = ready(certify_sampler(&s, REFERENCE, gs, QPU_DEFAULT, 4000, Provenance::Silicon)).unwrap().unwrap();
            let Parity::Statistical { fidelity, shots } = cert.parity else { panic!("{:?}", cert.parity) };
            assert_eq!(shots, 4000);
            assert!(fidelity > 0.99, "{gs:?}: {fidelity}");
            assert_eq!((cert.cases, cert.kernel.as_str()), (QPU_DEFAULT.cases as u64, sample_kernel(gs)));
        }
    }

    /// A sampler that returns what it is told to.
    struct Fixed(Counts);

    impl Sampler for Fixed {
        fn member_id(&self) -> String {
            "local/test".into()
        }

        fn sample<'r>(&'r self, _: &'r str, _: &'r Circuit, _: u64, _: &'r mut dyn RunMeter) -> SampleFuture<'r> {
            let counts = self.0.clone();
            Box::pin(async move { Ok((counts, 0, None)) })
        }
    }

    #[test]
    fn counts_of_the_wrong_shape_stop_a_certification() {
        let short = ready(certify_sampler(&Fixed(vec![1, 1]), "x", Gateset::V2, QPU_DEFAULT, 2, Provenance::Silicon)).unwrap();
        assert!(matches!(short, Err(MemberError::Other(_))), "{short:?}");
        let mut bell = Circuit::new(2);
        bell.h(0).cx(0, 1);
        assert!(check_counts(&bell, 10, &[5, 0, 0, 5]).is_ok());
        assert!(check_counts(&bell, 10, &[5, 0, 0, 4]).is_err());
        assert!(check_counts(&bell, 10, &[10, 0, 0]).is_err());
    }

    #[test]
    fn the_counts_hash_binds_circuit_shots_and_every_count() {
        let mut bell = Circuit::new(2);
        bell.h(0).cx(0, 1);
        let h = counts_hash(&bell, 10, &[5, 0, 0, 5]).unwrap();
        assert_ne!(counts_hash(&bell, 10, &[4, 0, 0, 6]), Some(h));
        assert_ne!(counts_hash(&bell, 11, &[5, 0, 0, 5]), Some(h));
        let mut other = bell.clone();
        other.gateset = Gateset::V1;
        assert_ne!(counts_hash(&other, 10, &[5, 0, 0, 5]), Some(h));
    }

    #[test]
    fn placement_keeps_the_reference_and_puts_silicon_before_a_simulator() {
        let local = local_sampling_member();
        let simulated = |id: &str, fidelity| {
            let mut entries = Vec::new();
            for gs in Gateset::ALL {
                let cert = Certificate {
                    kernel: sample_kernel(gs).into(),
                    implementation: "q".into(),
                    member: id.into(),
                    generator: QPU_DEFAULT.generator(),
                    parity: Parity::Statistical { fidelity, shots: 1000 },
                    cases: 16,
                    corpus: [3; 32],
                    provenance: Provenance::Silicon,
                };
                entries.push(KernelEntry { kernel: sample_kernel(gs).into(), implementation: "q".into(), certificate: cert, energy: None });
            }
            sampling_member(id.into(), FabricClass::Qpu, String::new(), entries)
        };
        let need = Parity::Statistical { fidelity: 0.9, shots: 1000 };
        // A silicon QPU whose id sorts first, and a simulator.
        let qpu = simulated("a/qpu", 0.95);
        let mut sim = simulated("a/aaa-sim", 0.999);
        sim.kernels.iter_mut().for_each(|k| k.certificate.provenance = Provenance::Simulator("test 1".into()));
        for gs in Gateset::ALL {
            let all = [local.clone(), qpu.clone(), sim.clone()];
            let (placed, _) = place(&all, sample_kernel(gs), &need).unwrap();
            assert_eq!(placed.id, qpu.id, "silicon first, whatever the ids");
            let pair = [sim.clone(), local.clone()];
            let (placed, _) = place(&pair, sample_kernel(gs), &need).unwrap();
            assert_eq!(placed.id, local.id, "the reference before a simulator");
            let (placed, _) = place(std::slice::from_ref(&sim), sample_kernel(gs), &need).unwrap();
            assert_eq!(placed.id, sim.id, "a simulator alone is placed");
            // A requirement above a member's fidelity leaves it out.
            let strict = Parity::Statistical { fidelity: 0.97, shots: 1000 };
            assert!(place(std::slice::from_ref(&qpu), sample_kernel(gs), &strict).is_none());
        }
    }
}