wai-quantum 0.3.30

//! Transport, done honestly: what a noisy channel does to a quantum state, and
//! how you find out what arrived.
//!
//! A quantum channel is the transport layer of this stack. Every physical thing
//! that can go wrong on the way — a bit flipping, a phase smearing, energy leaking
//! into the environment — is one completely-positive trace-preserving map, written
//! as a set of **Kraus operators** `{Kᵢ}` acting as `ρ ↦ Σᵢ Kᵢ ρ Kᵢ†` with
//! `Σᵢ Kᵢ†Kᵢ = I`. That last condition is the statement that probability is
//! conserved: something always arrives, even if it is noise.
//!
//! Two results here are worth more than the machinery:
//!
//! - **Tomography reconstructs what arrived.** Measuring every Pauli expectation
//!   determines the state uniquely — `ρ = (1/d) Σ_P Tr(ρP) · P` — so "what did the
//!   channel actually do" is an answerable question, not a guess.
//! - **Noise cannot create information.** Push an ensemble through any channel and
//!   its Holevo bound ([`crate::quantum_source`]) can only fall. This is the
//!   direction of the arrow for every transport claim: a channel loses, or at best
//!   preserves. Nothing you do downstream recovers what the channel discarded.
//!
//! # Determinism
//!
//! Kraus operators need `√p`, which is irrational, so they are built with the
//! crate's integer [`sqrt_fx`] and everything downstream stays in the same fixed
//! point as the simulator — byte-exact, no floats. The tolerances in the tests are
//! for fixed-point rounding, not for sampling: nothing here is sampled.

use crate::quantum::{fxmul, sqrt_fx, Amp, ONE};
use crate::quantum_info::{pauli_coeff, pauli_flip, DensityMatrix, Pauli};

/// A quantum channel in Kraus form, each operator stored row-major `d × d`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Channel {
    pub d: usize,
    pub kraus: Vec<Vec<Amp>>,
}

fn re(x: i64) -> Amp {
    Amp { re: x, im: 0 }
}

impl Channel {
    /// The channel that does nothing — a perfect wire.
    pub fn identity() -> Channel {
        Channel { d: 2, kraus: vec![vec![Amp::ONE, Amp::ZERO, Amp::ZERO, Amp::ONE]] }
    }

    /// With probability `p`, the bit flips.
    pub fn bit_flip(p: i64) -> Channel {
        let (a, b) = (sqrt_fx(ONE - p), sqrt_fx(p));
        Channel {
            d: 2,
            kraus: vec![
                vec![re(a), Amp::ZERO, Amp::ZERO, re(a)],
                vec![Amp::ZERO, re(b), re(b), Amp::ZERO],
            ],
        }
    }

    /// With probability `p`, the *phase* flips. Populations survive untouched and
    /// only coherence dies — the characteristically quantum failure, invisible to
    /// anyone who only looks at which basis state arrived.
    pub fn dephasing(p: i64) -> Channel {
        let (a, b) = (sqrt_fx(ONE - p), sqrt_fx(p));
        Channel {
            d: 2,
            kraus: vec![
                vec![re(a), Amp::ZERO, Amp::ZERO, re(a)],
                vec![re(b), Amp::ZERO, Amp::ZERO, re(-b)],
            ],
        }
    }

    /// `ρ ↦ (1−p)ρ + p·I/2`: with probability `p` the state is replaced by noise.
    /// At `p = 1` the output is maximally mixed and the message is gone.
    pub fn depolarizing(p: i64) -> Channel {
        let a = sqrt_fx(ONE - 3 * p / 4);
        let b = sqrt_fx(p / 4);
        Channel {
            d: 2,
            kraus: vec![
                vec![re(a), Amp::ZERO, Amp::ZERO, re(a)],
                vec![Amp::ZERO, re(b), re(b), Amp::ZERO],
                vec![Amp::ZERO, Amp { re: 0, im: -b }, Amp { re: 0, im: b }, Amp::ZERO],
                vec![re(b), Amp::ZERO, Amp::ZERO, re(-b)],
            ],
        }
    }

    /// Energy leaking out: `|1⟩` decays toward `|0⟩` with probability `gamma`.
    /// Unlike the others this channel is not unital — it has a preferred
    /// destination, which is what makes it a model of loss rather than of scrambling.
    pub fn amplitude_damping(gamma: i64) -> Channel {
        let keep = sqrt_fx(ONE - gamma);
        let decay = sqrt_fx(gamma);
        Channel {
            d: 2,
            kraus: vec![
                vec![Amp::ONE, Amp::ZERO, Amp::ZERO, re(keep)],
                vec![Amp::ZERO, re(decay), Amp::ZERO, Amp::ZERO],
            ],
        }
    }

    /// `ρ ↦ Σᵢ Kᵢ ρ Kᵢ†`.
    pub fn apply(&self, rho: &DensityMatrix) -> Option<DensityMatrix> {
        if rho.d != self.d {
            return None;
        }
        let d = self.d;
        let mut out = vec![Amp::ZERO; d * d];
        for k in &self.kraus {
            for a in 0..d {
                for b in 0..d {
                    let mut acc = Amp::ZERO;
                    for i in 0..d {
                        if k[a * d + i] == Amp::ZERO {
                            continue;
                        }
                        for j in 0..d {
                            if k[b * d + j] == Amp::ZERO {
                                continue;
                            }
                            acc = acc.add(
                                k[a * d + i].mul(rho.get(i, j)).mul(k[b * d + j].conj()),
                            );
                        }
                    }
                    out[a * d + b] = out[a * d + b].add(acc);
                }
            }
        }
        DensityMatrix::from_entries(out)
    }

    /// How far `Σ Kᵢ†Kᵢ` strays from the identity, as a fixed-point magnitude.
    ///
    /// Zero means probability is exactly conserved. A physical channel should give
    /// something on the order of fixed-point rounding; anything larger means the
    /// operators are not a channel at all.
    pub fn trace_preservation_defect(&self) -> i64 {
        let d = self.d;
        let mut sum = vec![Amp::ZERO; d * d];
        for k in &self.kraus {
            for a in 0..d {
                for b in 0..d {
                    let mut acc = Amp::ZERO;
                    for i in 0..d {
                        acc = acc.add(k[i * d + a].conj().mul(k[i * d + b]));
                    }
                    sum[a * d + b] = sum[a * d + b].add(acc);
                }
            }
        }
        let mut worst = 0i64;
        for a in 0..d {
            for b in 0..d {
                let want = if a == b { ONE } else { 0 };
                let e = sum[a * d + b];
                worst = worst.max((e.re - want).abs()).max(e.im.abs());
            }
        }
        worst
    }
}

/// Every Pauli string on `n` qubits, as measurement settings.
pub fn pauli_basis(n: u8) -> Vec<Vec<(u8, Pauli)>> {
    let mut out = vec![vec![]];
    for q in 0..n {
        let mut next = Vec::with_capacity(out.len() * 4);
        for base in &out {
            for p in [Pauli::I, Pauli::X, Pauli::Y, Pauli::Z] {
                let mut s = base.clone();
                s.push((q, p));
                next.push(s);
            }
        }
        out = next;
    }
    out
}

/// Reconstruct a state from its measurements: `ρ = (1/d) Σ_P Tr(ρP) · P`.
///
/// This is linear-inversion tomography, and it is exact — the Pauli strings are a
/// basis for Hermitian operators, so the expectations determine the state with
/// nothing left over. It is what makes a channel's effect *observable* rather than
/// merely modelled.
pub fn tomography(rho: &DensityMatrix) -> DensityMatrix {
    let d = rho.d;
    let mut out = vec![Amp::ZERO; d * d];
    for ops in pauli_basis(rho.n_qubits) {
        let t = rho.expect_pauli(&ops);
        if t == 0 {
            continue;
        }
        let flip = pauli_flip(&ops);
        for r in 0..d {
            let c = r ^ flip;
            // P[r, c] = coeff(c), and the 1/d normalisation of the basis
            let e = pauli_coeff(c, &ops);
            out[r * d + c] = out[r * d + c].add(Amp {
                re: fxmul(t, e.re) / d as i64,
                im: fxmul(t, e.im) / d as i64,
            });
        }
    }
    DensityMatrix::from_entries(out).expect("square by construction")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantum::Circuit;
    use crate::quantum_source::{holevo_bound, von_neumann_entropy};

    fn pure(c: &mut Circuit) -> DensityMatrix {
        DensityMatrix::from_pure(&c.simulate().unwrap())
    }
    fn ket0() -> DensityMatrix {
        pure(&mut Circuit::new(1))
    }
    fn ket1() -> DensityMatrix {
        let mut c = Circuit::new(1);
        c.x(0);
        pure(&mut c)
    }
    fn plus() -> DensityMatrix {
        let mut c = Circuit::new(1);
        c.h(0);
        pure(&mut c)
    }

    /// Something always arrives. Every channel here conserves probability to
    /// within fixed-point rounding.
    #[test]
    fn every_channel_conserves_probability() {
        for (name, ch) in [
            ("identity", Channel::identity()),
            ("bit_flip", Channel::bit_flip(ONE / 4)),
            ("dephasing", Channel::dephasing(ONE / 3)),
            ("depolarizing", Channel::depolarizing(ONE / 2)),
            ("damping", Channel::amplitude_damping(ONE / 5)),
        ] {
            let defect = ch.trace_preservation_defect();
            assert!(defect < 1 << 12, "{name} defect {defect} is too large to be a channel");
            let out = ch.apply(&plus()).unwrap();
            let tr = out.trace();
            assert!((tr.re - ONE).abs() < 1 << 12, "{name} trace {} != 1", tr.re);
        }
    }

    /// A perfect wire changes nothing.
    #[test]
    fn the_identity_channel_delivers_the_state_untouched() {
        assert_eq!(Channel::identity().apply(&plus()).unwrap(), plus());
    }

    /// Full depolarizing noise destroys the message: whatever went in, a
    /// maximally mixed state comes out, carrying one bit of entropy and no content.
    #[test]
    fn full_depolarizing_noise_destroys_the_message() {
        for input in [ket0(), plus(), ket1()] {
            let out = Channel::depolarizing(ONE).apply(&input).unwrap();
            let s = von_neumann_entropy(&out);
            assert!((s - 1.0).abs() < 1e-3, "entropy {s} should be a full bit of noise");
        }
    }

    /// Damping is loss with a direction: `|1⟩` fully damped arrives as `|0⟩`, and
    /// `|0⟩` — the ground state — is untouched.
    #[test]
    fn amplitude_damping_is_energy_loss_toward_the_ground_state() {
        let out = Channel::amplitude_damping(ONE).apply(&ket1()).unwrap();
        assert!((out.get(0, 0).re - ONE).abs() < 1 << 12, "|1> should decay to |0>");
        assert!(out.get(1, 1).re.abs() < 1 << 12);
        assert_eq!(Channel::amplitude_damping(ONE).apply(&ket0()).unwrap(), ket0());
    }

    /// The quantum-specific failure. Maximal dephasing leaves the populations
    /// exactly as they were — a receiver checking only which basis state arrived
    /// sees nothing wrong — while the coherence that made it a superposition is
    /// gone.
    ///
    /// Maximal is `p = ½`, not `p = 1`: at `p = 1` the `Z` fires on every run,
    /// which is a *unitary*, and a unitary cannot decohere anything. The damage is
    /// done by not knowing whether it fired.
    #[test]
    fn dephasing_destroys_coherence_and_leaves_populations_intact() {
        let out = Channel::dephasing(ONE / 2).apply(&plus()).unwrap();
        assert!((out.get(0, 0).re - ONE / 2).abs() < 1 << 12, "populations must survive");
        assert!((out.get(1, 1).re - ONE / 2).abs() < 1 << 12);
        assert!(out.get(0, 1).re.abs() < 1 << 12, "coherence must be gone");
        assert!(out.get(0, 1).im.abs() < 1 << 12);
        let s = von_neumann_entropy(&out);
        assert!((s - 1.0).abs() < 1e-3, "a dephased |+> is a classical coin: {s}");
    }

    /// The counterpart, and the reason the test above uses `p = ½`: dephasing at
    /// `p = 1` is the unitary `Z`. It sends `|+⟩` to `|−⟩` — a different state, but
    /// still a perfectly pure one, with no information lost at all.
    #[test]
    fn certain_dephasing_is_a_unitary_and_loses_nothing() {
        let out = Channel::dephasing(ONE).apply(&plus()).unwrap();
        let purity = out.purity() as f64 / ONE as f64;
        assert!((purity - 1.0).abs() < 1e-3, "purity {purity} — Z cannot decohere");
        assert!((out.get(0, 1).re + ONE / 2).abs() < 1 << 12, "|+> should have become |->");
        assert!(von_neumann_entropy(&out) < 1e-3);
    }

    /// Measuring every Pauli expectation determines the state uniquely, so what
    /// arrived is knowable rather than assumed.
    #[test]
    fn tomography_reconstructs_what_arrived() {
        let mut bell = Circuit::new(2);
        bell.h(0).cx(0, 1);
        for original in [ket0(), plus(), pure(&mut bell)] {
            let seen = tomography(&original);
            for r in 0..original.d {
                for c in 0..original.d {
                    let (a, b) = (original.get(r, c), seen.get(r, c));
                    assert!(
                        (a.re - b.re).abs() < 1 << 12 && (a.im - b.im).abs() < 1 << 12,
                        "reconstruction differs at ({r},{c}): {a:?} vs {b:?}"
                    );
                }
            }
        }
    }

    /// Tomography of a state that came out of a channel — the actual use. A
    /// half-depolarized `|+⟩` is reconstructed exactly as the channel left it.
    #[test]
    fn tomography_sees_what_the_channel_did() {
        let out = Channel::depolarizing(ONE / 2).apply(&plus()).unwrap();
        let seen = tomography(&out);
        assert!((seen.get(0, 1).re - out.get(0, 1).re).abs() < 1 << 12);
        // half the coherence of a clean |+> survived, and tomography reports it
        assert!(seen.get(0, 1).re > ONE / 8 && seen.get(0, 1).re < ONE / 3);
    }

    /// The direction of the arrow. Push an ensemble through a channel and the
    /// classical information anyone can extract can only fall — noise never
    /// creates information, and nothing downstream recovers what was discarded.
    #[test]
    fn noise_cannot_create_information() {
        let clean = [(ONE / 2, ket0()), (ONE / 2, ket1())];
        let before = holevo_bound(&clean).unwrap();
        let ch = Channel::depolarizing(ONE / 2);
        let noisy = [
            (ONE / 2, ch.apply(&clean[0].1).unwrap()),
            (ONE / 2, ch.apply(&clean[1].1).unwrap()),
        ];
        let after = holevo_bound(&noisy).unwrap();
        assert!(after <= before + 1e-6, "χ rose from {before} to {after}, which is impossible");
        assert!(after < before - 0.1, "this much noise should cost real information");
    }

    /// A non-trivial spectrum, pinned to bits.
    ///
    /// The maximally mixed case exercises the eigensolver barely at all — its
    /// eigenvalues are equal. A half-depolarized `|+⟩` has genuinely distinct
    /// eigenvalues, so this puts the Jacobi sweep and the hand-rolled `log2`
    /// through their real work and fixes the result exactly. Confirmed identical
    /// on `aarch64-apple-darwin` and `wasm32-wasip2` under wasmtime.
    #[test]
    fn a_non_trivial_spectrum_is_bit_identical_across_architectures() {
        let out = Channel::depolarizing(ONE / 2).apply(&plus()).unwrap();
        let s = von_neumann_entropy(&out);
        assert_eq!(
            s.to_bits(),
            0.8112781263732944f64.to_bits(),
            "entropy drifted from the pinned bits: {s:.17}"
        );
        let chi = holevo_bound(&[
            (ONE / 2, Channel::depolarizing(ONE / 2).apply(&ket0()).unwrap()),
            (ONE / 2, Channel::depolarizing(ONE / 2).apply(&ket1()).unwrap()),
        ])
        .unwrap();
        assert_eq!(
            chi.to_bits(),
            0.18872187458378642f64.to_bits(),
            "Holevo drifted from the pinned bits: {chi:.17}"
        );
    }

    /// A channel that does nothing costs nothing — the boundary case of the same law.
    #[test]
    fn a_perfect_wire_loses_nothing() {
        let ch = Channel::identity();
        let before = holevo_bound(&[(ONE / 2, ket0()), (ONE / 2, ket1())]).unwrap();
        let after = holevo_bound(&[
            (ONE / 2, ch.apply(&ket0()).unwrap()),
            (ONE / 2, ch.apply(&ket1()).unwrap()),
        ])
        .unwrap();
        assert!((after - before).abs() < 1e-9, "{before} -> {after}");
    }
}