stabilizer_qec 2.1.0

//! The joint sampler: Clifford+T circuits with coherent errors on every gate. The non-Clifford
//! *gates* (`I[T]`, `I[R_Z(theta=…)]`, `I[U3(…)]`) are exact terms of the stabilizer-rank
//! engine; the coherent *errors* (`I_ERROR[R_Z(theta=…)]`, `II_ERROR[R_ZZ(…)]`) are drawn as
//! their Pauli twirl and the shot is reweighted, as the coherent sampler reweights it, by the
//! interference the twirl leaves out.
//!
//! Expanding the errors alone, the circuit is U = Σ_S a_S U_S over error sets S, where U_S has
//! the Pauli P_g inserted at each g ∈ S and keeps every gate, T gates included, and
//! a_S = Π_{g∈S}(−i sin θ_g) Π_{g∉S} cos θ_g. For a record m (and the unread outcomes u of
//! its resets) let |Ψ_S⟩ be the term engine's final state along that trajectory with S's
//! Paulis in its frame: P(m, u) = ‖Σ_S a_S |Ψ_S⟩‖² = Σ_{S,S′} a_S a*_{S′} ⟨Ψ_{S′}|Ψ_S⟩. A shot
//! draws S as the twirl (|a_S|²) and (m, u) from the exact terms (‖Ψ_S‖²), and its weight puts
//! back the cross terms: with ρ_G = a_{S⊕G}/a_S,
//!
//! ```text
//!     w = 1 + Σ_{G ≠ ∅} ρ_G ⟨Ψ_{S⊕G}|Ψ_S⟩ / ‖Ψ_S‖²,
//! ```
//!
//! each pair counted once, from its likelier side (so the weight stays real and the estimate
//! unbiased for every pair whose difference G is among the sets summed). Two error sets whose
//! states overlap differ by an element of the spacetime kernel projected onto the error
//! locations — every non-Clifford gate is a location for the kernel's search, so that the
//! interference running through the T gates is found — and the generators are a basis of that
//! projection: the kernel's local generators, the single errors and pairs the exact gates'
//! operators relate, and (up to `max_complete` locations) the directions those leave out, from
//! the null space of the errors' residues modulo the exact gates' (`kernel::null_space`; a
//! logical error at one time with the same logical at another is the commonest, which no local
//! search finds and whose halves add as amplitudes). With at most `SMALL_KERNEL` generators a
//! shot sums every product of them, the whole projected kernel: exact. Otherwise the sets
//! summed are the generators each on its own, products of up to `order` of the generators a
//! fired error touches (per cluster), and products within each group of locations with equal
//! flips (the same fault in different places, a T gate's among them), as the coherent sampler
//! truncates; with the whole kernel, every projection, and the identity is
//! exact (`exact_distribution`, checked against the state vector). Unlike the coherent
//! sampler's class form, the sum is not a product over clusters: through a T gate, members of
//! different clusters overlap. Each member's state is a re-run of the term engine with the
//! member's Paulis inserted and the shot's outcomes forced, brought to the shot's frame so
//! that overlaps are dot products over shared keys.

use std::collections::{BTreeMap, HashMap};

use super::kernel::{CoherentOptions, Kernel, SMALL_KERNEL};
use super::{Element, GaugeSolver, Program as KernelProgram};
use crate::circuit::Circuit;
use crate::stabrank::{Key, Memo, Odometer, Program, Recording, Replay, Run, Source, Stream};
use crate::statevec::C;
use crate::surface_code::Xorshift;

pub struct JointSampler {
    /// The term engine, the coherent errors as twirled locations.
    pub program: Program,
    /// The kernel's circuit: every rotation, exact gate or error, a location.
    pub kernel_program: KernelProgram,
    pub kernel: Kernel,
    solver: GaugeSolver,
    /// Each twirled location's index among the kernel's locations.
    loc_of_twirl: Vec<usize>,
    /// Each kernel location's twirled index, if it is an error.
    twirl_of_loc: Vec<Option<usize>>,
    /// The kernel's generators, and every pair within each group of equal-flip locations.
    generators: Vec<Element>,
    /// How many of the generators are local (found by the kernel's search; the first ones).
    num_local: usize,
    /// How many of the generators complete the local ones to a basis of the projected kernel
    /// (those after the local ones).
    num_nonlocal: usize,
    /// Whether the generators span the whole projected kernel (false past `max_complete`
    /// locations, where interference between distant errors is left out).
    complete: bool,
    /// Each group's pair generators (indices into `generators`).
    group_gens: Vec<Vec<usize>>,
    by_location: Vec<Vec<usize>>,
    options: CoherentOptions,
    /// Whether to sum over the products of every two generators too (quadratic in their
    /// number: a check on small circuits).
    pairs: bool,
    /// With the whole kernel (small circuits): every element.
    whole: Option<Vec<Element>>,
}

impl JointSampler {
    pub fn new(circuit: &Circuit, options: CoherentOptions) -> Result<JointSampler, String> {
        let expanded = crate::stabrank::expand_clifford_rotations(circuit)?;
        let program = Program::with_twirl(&expanded)?;
        let kernel_program = KernelProgram::new(&expanded)?;
        let solver = GaugeSolver::new(&kernel_program);
        // The kernel's local search alone: the completion is of the projection, below.
        let kernel = Kernel::new(&kernel_program, &solver, CoherentOptions { max_complete: 0, ..options });
        let mut loc_of_twirl = Vec::new();
        let mut twirl_of_loc = vec![None; kernel_program.locations.len()];
        for (l, &exact) in kernel_program.exact.iter().enumerate() {
            if !exact {
                twirl_of_loc[l] = Some(loc_of_twirl.len());
                loc_of_twirl.push(l);
            }
        }
        if loc_of_twirl.len() != program.num_twirl() {
            return Err("the kernel's error locations do not match the term engine's".into());
        }
        // The generators: a basis of the kernel's projection onto the error locations, from the
        // kernel's own generators, every single error whose flips the exact gates' operators
        // match (its projection one location: the first-order cross terms) and every pair of
        // errors that agree modulo those operators, smallest projection first and each kept
        // only if independent of those kept. The kernel's basis alone, found over every
        // location, can hold {e₁, T} and {e₁, e₂} and leave {e₂} to a product no shot forms:
        // the sum then misses first-order terms (the review's 1-qubit cases, 60σ off).
        let nl = kernel_program.locations.len();
        let residue = super::kernel::residues(&kernel_program, &solver);
        let mut rows: Vec<(usize, Vec<u64>, Vec<usize>)> = Vec::new();
        let reduce = |v: &mut Vec<u64>, combo: &mut Vec<usize>, rows: &[(usize, Vec<u64>, Vec<usize>)]| {
            for (pivot, rv, rc) in rows {
                if v[pivot / 64] >> (pivot % 64) & 1 == 1 {
                    v.iter_mut().zip(rv).for_each(|(a, b)| *a ^= b);
                    *combo = super::kernel::symmetric_difference(combo, rc);
                }
            }
        };
        for l in (0..nl).filter(|&l| kernel_program.exact[l]) {
            let (mut v, mut combo) = (residue[l].clone(), vec![l]);
            reduce(&mut v, &mut combo, &rows);
            if let Some(w) = v.iter().position(|&w| w != 0) {
                let pivot = w * 64 + v[w].trailing_zeros() as usize;
                for (_, rv, rc) in rows.iter_mut() {
                    if rv[pivot / 64] >> (pivot % 64) & 1 == 1 {
                        rv.iter_mut().zip(&v).for_each(|(a, b)| *a ^= b);
                        *rc = super::kernel::symmetric_difference(rc, &combo);
                    }
                }
                rows.push((pivot, v, combo));
            }
        }
        let mut candidates: std::collections::BTreeSet<Vec<usize>> = kernel.generators.iter().map(|e| e.members.clone()).collect();
        let mut by_reduced: HashMap<Vec<u64>, Vec<usize>> = HashMap::new();
        for &l in &loc_of_twirl {
            let (mut v, mut combo) = (residue[l].clone(), vec![l]);
            reduce(&mut v, &mut combo, &rows);
            if v.iter().all(|&w| w == 0) {
                candidates.insert(combo);
            } else {
                by_reduced.entry(v).or_default().push(l);
            }
        }
        for ls in by_reduced.values() {
            for w in ls.windows(2) {
                let mut v: Vec<u64> = residue[w[0]].iter().zip(&residue[w[1]]).map(|(a, b)| a ^ b).collect();
                let mut combo = vec![w[0], w[1]];
                reduce(&mut v, &mut combo, &rows);
                if v.iter().all(|&w| w == 0) {
                    candidates.insert(combo);
                }
            }
        }
        let projection = |members: &[usize]| -> Vec<usize> { members.iter().filter_map(|&l| twirl_of_loc[l]).collect() };
        let mut candidates: Vec<Vec<usize>> = candidates.into_iter().collect();
        candidates.sort_by_key(|g| (projection(g).len(), g.len(), g.clone()));
        let tw = loc_of_twirl.len().div_ceil(64).max(1);
        let mut basis: Vec<(usize, Vec<u64>)> = Vec::new();
        let mut generators: Vec<Element> = Vec::new();
        for g in candidates {
            let mut v = vec![0u64; tw];
            for t in projection(&g) {
                v[t / 64] ^= 1 << (t % 64);
            }
            for (pivot, row) in &basis {
                if v[pivot / 64] >> (pivot % 64) & 1 == 1 {
                    v.iter_mut().zip(row).for_each(|(a, b)| *a ^= b);
                }
            }
            let Some(w) = v.iter().position(|&w| w != 0) else { continue };
            let Some(e) = kernel_program.element(&g, &solver) else { continue };
            let pivot = w * 64 + v[w].trailing_zeros() as usize;
            for (_, row) in basis.iter_mut() {
                if row[pivot / 64] >> (pivot % 64) & 1 == 1 {
                    row.iter_mut().zip(&v).for_each(|(a, b)| *a ^= b);
                }
            }
            basis.push((pivot, v));
            generators.push(e);
        }
        // The directions the generators and the groups' pairs leave out. A projected kernel
        // element is a set of errors whose residues, reduced by the exact gates', sum to zero;
        // reduced by the span (in echelon form over the errors) it is off that span's pivots, and
        // a nonzero element of the span is on one: so the quotient is the null space of the free
        // errors' reduced residues, each vector carrying the exact gates that complete it.
        let num_local = generators.len();
        let mut complete = nl <= options.max_complete;
        if complete {
            let mut span = basis.clone();
            for g in &kernel.groups {
                for (i, &a) in g.members.iter().enumerate() {
                    for &b in &g.members[i + 1..] {
                        let mut v = vec![0u64; tw];
                        for t in projection(&[a, b]) {
                            v[t / 64] ^= 1 << (t % 64);
                        }
                        for (pivot, row) in &span {
                            if v[pivot / 64] >> (pivot % 64) & 1 == 1 {
                                v.iter_mut().zip(row).for_each(|(a, b)| *a ^= b);
                            }
                        }
                        let Some(w) = v.iter().position(|&w| w != 0) else { continue };
                        let pivot = w * 64 + v[w].trailing_zeros() as usize;
                        for (_, row) in span.iter_mut() {
                            if row[pivot / 64] >> (pivot % 64) & 1 == 1 {
                                row.iter_mut().zip(&v).for_each(|(a, b)| *a ^= b);
                            }
                        }
                        span.push((pivot, v));
                    }
                }
            }
            let mut pinned = vec![false; loc_of_twirl.len()];
            for (pivot, _) in &span {
                pinned[*pivot] = true;
            }
            let free: Vec<(Vec<u64>, Vec<usize>)> = (0..loc_of_twirl.len())
                .filter(|&t| !pinned[t])
                .map(|t| {
                    let l = loc_of_twirl[t];
                    let (mut v, mut combo) = (residue[l].clone(), vec![l]);
                    reduce(&mut v, &mut combo, &rows);
                    (v, combo)
                })
                .collect();
            // Fewest errors first, then fewest exact gates.
            let local: Vec<Vec<usize>> = generators.iter().map(|e| e.members.clone()).collect();
            let cost = |l: usize| -> u64 { if kernel_program.exact[l] { 1 } else { 1 << 16 } };
            for members in super::kernel::shorten(super::kernel::null_space(free), &local, nl, &cost) {
                match kernel_program.element(&members, &solver) {
                    Some(e) => generators.push(e),
                    None => complete = false,
                }
            }
        }
        let num_nonlocal = generators.len() - num_local;
        let mut group_gens = Vec::new();
        for g in &kernel.groups {
            let mut gens = Vec::new();
            for (i, &a) in g.members.iter().enumerate() {
                for &b in &g.members[i + 1..] {
                    if let Some(e) = kernel_program.element(&[a, b], &solver) {
                        gens.push(generators.len());
                        generators.push(e);
                    }
                }
            }
            group_gens.push(gens);
        }
        let mut by_location = vec![Vec::new(); kernel_program.locations.len()];
        for (i, g) in generators.iter().enumerate() {
            for &l in &g.members {
                by_location[l].push(i);
            }
        }
        Ok(JointSampler { program, kernel_program, kernel, solver, loc_of_twirl, twirl_of_loc, generators, num_local, num_nonlocal, complete, group_gens, by_location, options, pairs: false, whole: None })
    }

    /// The same, summing also over the products of every two untouched generators.
    pub fn with_pairs(mut self, pairs: bool) -> JointSampler {
        self.pairs = pairs;
        self
    }

    /// The same, weighing every shot with the whole kernel: for small circuits, the exact
    /// identity.
    pub fn with_whole_kernel(circuit: &Circuit) -> Result<JointSampler, String> {
        let mut s = JointSampler::new(circuit, CoherentOptions::default())?;
        s.whole = Some(s.kernel_program.full_kernel()?);
        Ok(s)
    }

    pub fn num_detectors(&self) -> usize {
        self.program.num_detectors()
    }

    pub fn num_observables(&self) -> usize {
        self.program.num_observables()
    }

    /// The coherent errors.
    pub fn num_locations(&self) -> usize {
        self.loc_of_twirl.len()
    }

    /// The non-Clifford gates, run exactly.
    pub fn num_exact(&self) -> usize {
        self.kernel_program.exact.iter().filter(|&&e| e).count()
    }

    pub fn num_generators(&self) -> usize {
        self.generators.len()
    }

    /// The generators that complete the local ones to a basis of the projected kernel.
    pub fn num_nonlocal_generators(&self) -> usize {
        self.num_nonlocal
    }

    /// Whether the generators span the whole projected kernel.
    pub fn kernel_complete(&self) -> bool {
        self.complete
    }

    /// Whether every shot sums the whole projected kernel (complete, and at most
    /// `SMALL_KERNEL` generators): the weights are then exact.
    pub fn is_whole(&self) -> bool {
        self.whole.is_some() || (self.complete && self.generators.len() <= SMALL_KERNEL)
    }

    /// The error sets G summed for a shot whose errors `fired`: distinct, sorted lists of
    /// twirled locations, never empty. The projections of every untouched generator, of the
    /// products of up to `order` touched generators per cluster, and of the products within
    /// each group; with the whole kernel, of every element.
    fn members(&self, fired: &[bool], cache: &mut HashMap<u64, Option<Element>>) -> Vec<Vec<usize>> {
        let mut out: Vec<Vec<usize>> = Vec::new();
        let add = |elements: Vec<Element>, out: &mut Vec<Vec<usize>>| {
            for e in elements {
                let g: Vec<usize> = e.members.iter().filter_map(|&l| self.twirl_of_loc[l]).collect();
                if !g.is_empty() && !out.contains(&g) {
                    out.push(g);
                }
            }
        };
        if let Some(all) = &self.whole {
            add(all.clone(), &mut out);
            return out;
        }
        // A small kernel: every product of its generators, the same set from every shot (a
        // truncation that sums G from S but not from S ⊕ G is biased at second order; with
        // this many generators the whole span is cheap).
        if self.generators.len() <= SMALL_KERNEL {
            let n = self.generators.len();
            let element_of = |members: &[usize]| self.kernel_program.element(members, &self.solver);
            for mask in 1u32..(1u32 << n) {
                let mut members: Vec<usize> = Vec::new();
                for (g, e) in self.generators.iter().enumerate() {
                    if mask >> g & 1 == 1 {
                        members = super::kernel::symmetric_difference(&members, &e.members);
                    }
                }
                if members.is_empty() {
                    continue;
                }
                let key = members.iter().fold(0u64, |h, &l| h ^ self.kernel.zobrist(l));
                if let Some(e) = cache.entry(key).or_insert_with(|| element_of(&members)).clone() {
                    add(vec![e], &mut out);
                }
            }
            return out;
        }
        let mut fired_locs: Vec<bool> = vec![false; self.kernel_program.locations.len()];
        for (t, &f) in fired.iter().enumerate() {
            if f {
                fired_locs[self.loc_of_twirl[t]] = true;
            }
        }
        // The local generators (and group pairs) the errors touch, by clusters as before 2.1;
        // the non-local ones each alone, and with their neighbours below.
        let nonlocal = |g: usize| g >= self.num_local && g < self.num_local + self.num_nonlocal;
        let touched = self.kernel.touched(&self.generators, &self.by_location, &fired_locs, self.options.hops, &|g| !nonlocal(g));
        let mut untouched = Vec::new();
        let mut t = touched.iter().peekable();
        for (g, e) in self.generators.iter().enumerate() {
            if t.peek() == Some(&&g) {
                t.next();
                continue;
            }
            untouched.push(g);
            add(vec![e.clone()], &mut out);
        }
        let mut element_of = |members: &[usize]| self.kernel_program.element(members, &self.solver);
        if self.pairs {
            // Every pair of generators, touched or not: the sum of G from S and from S ⊕ G must
            // agree on which pairs are in, or the second-order terms stop cancelling.
            let all: Vec<usize> = (0..self.generators.len()).collect();
            for (i, &a) in all.iter().enumerate() {
                for &b in &all[i + 1..] {
                    let members = super::kernel::symmetric_difference(&self.generators[a].members, &self.generators[b].members);
                    if members.is_empty() {
                        continue;
                    }
                    let key = members.iter().fold(0u64, |h, &l| h ^ self.kernel.zobrist(l));
                    if let Some(e) = cache.entry(key).or_insert_with(|| element_of(&members)).clone() {
                        add(vec![e], &mut out);
                    }
                }
            }
        }
        for cluster in self.kernel.clusters(&self.generators, &touched) {
            let elements = self.kernel.products(&self.generators, &cluster, &self.options, cache, &mut element_of, &|_| true, None);
            add(elements.into_iter().map(|(_, e)| e).collect(), &mut out);
        }
        for gens in &self.group_gens {
            if gens.len() > 1 {
                let elements = self.kernel.products(&self.generators, gens, &self.options, cache, &mut element_of, &|_| true, None);
                add(elements.into_iter().map(|(_, e)| e).collect(), &mut out);
            }
        }
        // A non-local generator the errors reach (holding one, or sharing a location with a
        // touched generator: a logical error at one time, its partner the same logical at
        // another): itself, and its products with the products of up to `order` generators
        // around it (those sharing a location with it, in clusters of their own, joined when
        // they share a location), which move the errors on along it. Not merged into the
        // clusters, whose products past `max_cluster` generators stop at pairs.
        let mut reached = self.kernel.touched(&self.generators, &self.by_location, &fired_locs, 0, &nonlocal);
        for &t in &touched {
            for &l in &self.generators[t].members {
                reached.extend(self.by_location[l].iter().copied().filter(|&g| nonlocal(g)));
            }
        }
        reached.sort_unstable();
        reached.dedup();
        for &g in &reached {
            let ge = &self.generators[g];
            let mut around: Vec<usize> = ge.members.iter().flat_map(|&l| self.by_location[l].iter().copied()).filter(|&h| h != g).collect();
            around.sort_unstable();
            around.dedup();
            // And the non-local generators chained to those, up to `order` steps further: the
            // same logical at further times.
            let mut frontier: Vec<usize> = around.iter().copied().filter(|&h| nonlocal(h)).collect();
            for _ in 1..self.options.order {
                let mut next: Vec<usize> = frontier.iter().flat_map(|&h| self.generators[h].members.iter().flat_map(|&l| self.by_location[l].iter().copied())).filter(|&x| nonlocal(x) && x != g && around.binary_search(&x).is_err()).collect();
                next.sort_unstable();
                next.dedup();
                if next.is_empty() {
                    break;
                }
                around.extend_from_slice(&next);
                around.sort_unstable();
                frontier = next;
            }
            for d in self.kernel.clusters(&self.generators, &around) {
                let elements = self.kernel.products(&self.generators, &d, &self.options, cache, &mut element_of, &|_| true, Some(&ge.members));
                add(elements.into_iter().map(|(_, e)| e).collect(), &mut out);
            }
        }
        out
    }

    /// a_{S⊕G} / a_S for the error set G: −i tan θ to add a fault, i / tan θ to take one away.
    fn rho(&self, fired: &[bool], g: &[usize]) -> C {
        g.iter().fold(C::ONE, |acc, &t| {
            let theta = self.kernel_program.locations[self.loc_of_twirl[t]].theta;
            let v = if fired[t] { C::new(0.0, 1.0 / theta.tan()) } else { C::new(0.0, -theta.tan()) };
            acc * v
        })
    }

    /// A member's final state in the shot's frame, as its terms (keys to amplitudes, the
    /// trajectory's probability in their norm), from a re-run with the shot's draws replayed
    /// and the member's errors inserted. Zero (empty) where the shot's outcomes are impossible
    /// for it.
    fn member(&self, base: &Run, fired: &[bool], memo: &mut Memo) -> HashMap<Key, C> {
        let mut src = Replay::new(&base.trace, fired);
        let run = self.program.run_full(&mut src, memo);
        if run.dead {
            return HashMap::new();
        }
        self.program.in_frame(&run, &base.frame)
    }

    /// One shot: its records, and its weight.
    pub(crate) fn shot(&self, src: &mut dyn Source, memo: &mut Memo, cache: &mut HashMap<u64, Option<Element>>) -> (Vec<bool>, f64) {
        let mut rec = Recording::new(src, self.program.num_twirl());
        let base = self.program.run_full(&mut rec, memo);
        let fired = rec.fired.clone();
        let mut base_run = base;
        base_run.trace = rec.into_trace();
        let base_terms = self.program.in_frame(&base_run, &base_run.frame);
        let base_norm: f64 = base_terms.values().map(|c| c.norm2()).sum();
        if base_norm <= 0.0 {
            return (base_run.records, 1.0);
        }
        let mut w = 1.0;
        for g in self.members(&fired, cache) {
            let mut member_fired = fired.clone();
            for &t in &g {
                member_fired[t] = !member_fired[t];
            }
            let rho = self.rho(&fired, &g);
            let terms = self.member(&base_run, &member_fired, memo);
            if terms.is_empty() {
                continue;
            }
            // The pair (S, S ⊕ G) counted from its likelier side, both of its orders at once
            // (2 Re): the twirl's probability of the member relative to the shot's is
            // |ρ|² ‖Ψ_member‖² / ‖Ψ_S‖²; a near tie goes to the smaller firing pattern, which
            // both sides agree on.
            let member_norm: f64 = terms.values().map(|c| c.norm2()).sum();
            let relative = rho.norm2() * member_norm / base_norm;
            let likelier = if (relative - 1.0).abs() < 1e-9 { fired < member_fired } else { relative < 1.0 };
            if !likelier {
                continue;
            }
            let overlap = terms.iter().fold(C::ZERO, |acc, (k, c)| acc + c.conj() * base_terms.get(k).copied().unwrap_or(C::ZERO));
            // a_S a*_{S⊕G} ⟨Ψ_{S⊕G}|Ψ_S⟩ over |a_S|² ‖Ψ_S‖² is ρ* times the overlap.
            w += 2.0 * (rho.conj() * overlap).re / base_norm;
        }
        (base_run.records, w)
    }

    /// Shots `first..first + shots`, each from its own stream: detection events and observable
    /// flips as b8 rows, and each shot's weight.
    pub fn sample_seeded(&self, seed: u64, first: u64, shots: usize, threads: usize) -> super::sampler::CoherentBatch {
        let p = &self.program;
        let (nd, no) = (p.num_detectors(), p.num_observables());
        let (dw, ow) = (nd.div_ceil(8), no.div_ceil(8));
        let rows = crate::parallel::parallel(shots, threads, |range| {
            let mut memo = Memo::default();
            let mut cache = HashMap::new();
            range
                .map(|s| {
                    let mut rng = Xorshift::new(crate::batch_sampler::batch_seed(seed, first + s as u64));
                    let (records, w) = self.shot(&mut Stream(&mut rng), &mut memo, &mut cache);
                    (p.outcome(&records), w)
                })
                .collect()
        });
        let mut batch = super::sampler::CoherentBatch { detectors: vec![0u8; shots * dw], observables: vec![0u8; shots * ow], weights: Vec::with_capacity(shots) };
        for (s, ((dets, obs), w)) in rows.into_iter().enumerate() {
            for (k, &b) in dets.iter().enumerate() {
                if b {
                    batch.detectors[s * dw + k / 8] |= 1 << (k % 8);
                }
            }
            for k in 0..no {
                if obs >> k & 1 == 1 {
                    batch.observables[s * ow + k / 8] |= 1 << (k % 8);
                }
            }
            batch.weights.push(w);
        }
        batch
    }
}

impl JointSampler {
    /// Every outcome's probability by the weighted twirl, summed over every branch (noise
    /// draws, firings, outcomes): with the whole kernel, the exact distribution; with the
    /// truncated one, its bias made visible. At most `max_branches`.
    pub fn exact_distribution(&self, max_branches: usize) -> Result<BTreeMap<(Vec<bool>, u64), f64>, String> {
        let mut out = BTreeMap::new();
        let mut memo = Memo::default();
        let mut cache = HashMap::new();
        let mut odo = Odometer::default();
        let mut branches = 0usize;
        loop {
            odo.restart();
            let (records, w) = self.shot(&mut odo, &mut memo, &mut cache);
            if odo.p > 1e-15 {
                *out.entry(self.program.outcome(&records)).or_insert(0.0) += odo.p * w;
            }
            branches += 1;
            if branches > max_branches {
                return Err(format!("more than {max_branches} branches"));
            }
            if !odo.advance() {
                return Ok(out);
            }
        }
    }
}

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

    fn truth(c: &Circuit) -> BTreeMap<(Vec<bool>, u64), f64> {
        crate::statevec::Program::new(c).unwrap().distribution(1 << 22).unwrap()
    }

    fn close(a: &BTreeMap<(Vec<bool>, u64), f64>, b: &BTreeMap<(Vec<bool>, u64), f64>, text: &str) {
        let keys: std::collections::BTreeSet<_> = a.keys().chain(b.keys()).collect();
        for k in keys {
            let (x, y) = (a.get(k).copied().unwrap_or(0.0), b.get(k).copied().unwrap_or(0.0));
            assert!((x - y).abs() < 1e-9, "{text}\n{k:?}: {x} vs {y}");
        }
    }

    fn check(text: &str) {
        let c = Circuit::parse(text).unwrap();
        let s = JointSampler::with_whole_kernel(&c).unwrap();
        close(&s.exact_distribution(1 << 22).unwrap(), &truth(&c), text);
    }

    #[test]
    fn one_error_before_a_t_gate() {
        check("R 0\nH 0\nI_ERROR[R_X(theta=0.3)] 0\nI[T] 0\nH 0\nM 0\nDETECTOR rec[-1]");
        check("R 0\nH 0\nI_ERROR[R_Z(theta=0.3)] 0\nI[T] 0\nI_ERROR[R_X(theta=0.2)] 0\nH 0\nM 0\nDETECTOR rec[-1]");
    }

    #[test]
    fn errors_either_side_of_t_gates_on_two_qubits() {
        check("R 0 1\nH 0\nI_ERROR[R_Y(theta=0.25)] 0\nCX 0 1\nI[T] 1\nII_ERROR[R_ZZ(theta=0.2)] 0 1\nI[T_DAG] 0\nH 0 1\nM 0 1\nDETECTOR rec[-1]\nDETECTOR rec[-2]\nOBSERVABLE_INCLUDE(0) rec[-1]");
    }

    #[test]
    fn with_noise_resets_and_mid_circuit_measurements() {
        check("R 0 1\nH 0\nI_ERROR[R_X(theta=0.3)] 0\nDEPOLARIZE1(0.1) 1\nCX 0 1\nI[T] 0\nMR 1\nI_ERROR[R_Z(theta=0.2)] 1\nCX 0 1\nI[U3(theta=0.5, phi=0.2, lambda=1.1)] 1\nH 0\nM(0.05) 0 1\nDETECTOR rec[-1] rec[-3]\nDETECTOR rec[-2]\nOBSERVABLE_INCLUDE(0) rec[-1]");
    }

    #[test]
    fn random_circuits() {
        let mut rng = Xorshift::new(99);
        for case in 0..40 {
            let n = 1 + case % 3;
            let text = crate::stabrank::tests::random_circuit(&mut rng, n, 5 + case % 5, true);
            let c = Circuit::parse(&text).unwrap();
            // The kernel's program takes no feedback; the whole kernel is for a few locations.
            let Ok(p) = KernelProgram::new(&c) else { continue };
            if p.locations.len() > 12 {
                continue;
            }
            check(&text);
        }
    }

    /// The sampler as built, its own sums (not `with_whole_kernel`'s), over every branch.
    fn check_built(text: &str) -> JointSampler {
        let c = Circuit::parse(text).unwrap();
        let s = JointSampler::new(&c, CoherentOptions::default()).unwrap();
        close(&s.exact_distribution(1 << 22).unwrap(), &truth(&c), text);
        s
    }

    /// A logical error in one round and the same logical in the next: a kernel element of
    /// weight six that no local search finds, whose halves add as amplitudes (the 2.1 audit's
    /// case: the sampler read the twirl's rate, 44σ off). With a T gate beside the memory, and
    /// with one on a data qubit between the rounds (its Z a kernel element of its own).
    #[test]
    fn logical_errors_at_different_times_interfere() {
        let memory = "R 0 1 2 3 4\nI_ERROR[R_X(theta=0.3)] 0 2 4\nCX 0 1 2 3\nCX 2 1 4 3\nMR 1 3\nDETECTOR rec[-2]\nDETECTOR rec[-1]\nBETWEEN\
                      I_ERROR[R_X(theta=0.3)] 0 2 4\nCX 0 1 2 3\nCX 2 1 4 3\nMR 1 3\nDETECTOR rec[-2] rec[-4]\nDETECTOR rec[-1] rec[-3]\n\
                      M 0 2 4\nDETECTOR rec[-3] rec[-2] rec[-5]\nDETECTOR rec[-2] rec[-1] rec[-4]\nOBSERVABLE_INCLUDE(0) rec[-1]";
        for text in [memory.replace("BETWEEN", "") + "\nR 5\nH 5\nI[T] 5\nH 5\nM 5\nDETECTOR rec[-1]", memory.replace("BETWEEN", "I[T] 0\n")] {
            let s = check_built(&text);
            assert!(s.num_nonlocal_generators() >= 1 && s.kernel_complete() && s.is_whole(), "{text}");
        }
        // Without the completion the same sums miss it.
        let text = memory.replace("BETWEEN", "") + "\nR 5\nH 5\nI[T] 5\nH 5\nM 5\nDETECTOR rec[-1]";
        let c = Circuit::parse(&text).unwrap();
        let s = JointSampler::new(&c, CoherentOptions { max_complete: 0, ..CoherentOptions::default() }).unwrap();
        assert!(!s.kernel_complete() && s.num_nonlocal_generators() == 0);
        let (ours, theirs) = (s.exact_distribution(1 << 22).unwrap(), truth(&c));
        let worst = theirs.iter().map(|(k, v)| (ours.get(k).copied().unwrap_or(0.0) - v).abs()).fold(0.0, f64::max);
        assert!(worst > 1e-4, "{worst}");
    }

    #[test]
    fn without_gates_the_weights_are_the_coherent_samplers() {
        // No non-Clifford gate: every member is Pauli-related and the class weight is 1.9's.
        let text = "R 0 1\nH 0\nI_ERROR[R_Z(theta=0.3)] 0\nCX 0 1\nI_ERROR[R_X(theta=0.2)] 1\nM 0 1\nDETECTOR rec[-1] rec[-2]\nOBSERVABLE_INCLUDE(0) rec[-1]";
        check(text);
        let c = Circuit::parse(text).unwrap();
        let s = JointSampler::new(&c, CoherentOptions::default()).unwrap();
        assert_eq!(s.num_exact(), 0);
        let b = s.sample_seeded(1, 0, 200, 1);
        assert_eq!(b.weights.len(), 200);
        assert!(b.weights.iter().all(|w| w.is_finite() && *w >= 0.0));
    }
}