toroidal-noise 0.2.0

#![doc = include_str!("../README.md")]

use num_complex::Complex64;
use std::f64::consts::{FRAC_1_SQRT_2, PI};

/// A 2×2 complex matrix stored in row-major order.
pub type Matrix2x2 = [[Complex64; 2]; 2];

/// Identity matrix.
pub const IDENTITY: Matrix2x2 = [
    [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)],
    [Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)],
];

/// A Kraus operator: a 2×2 complex matrix K such that the channel
/// maps ρ → Σᵢ Kᵢ ρ Kᵢ†.
#[derive(Debug, Clone, Copy)]
pub struct KrausOperator(pub Matrix2x2);

impl KrausOperator {
    /// Conjugate transpose K†.
    #[must_use]
    pub fn adjoint(&self) -> Self {
        let m = &self.0;
        Self([
            [m[0][0].conj(), m[1][0].conj()],
            [m[0][1].conj(), m[1][1].conj()],
        ])
    }

    /// Matrix of the operator.
    #[must_use]
    pub const fn matrix(&self) -> &Matrix2x2 {
        &self.0
    }
}

/// Apply a channel (list of Kraus operators) to a 2×2 density matrix:
/// ρ → Σᵢ Kᵢ ρ Kᵢ†
#[must_use]
pub fn apply_channel(rho: &Matrix2x2, kraus: &[KrausOperator]) -> Matrix2x2 {
    let mut result = [[Complex64::new(0.0, 0.0); 2]; 2];
    for k in kraus {
        let kd = k.adjoint();
        let k_rho = mul2x2(&k.0, rho);
        let k_rho_kd = mul2x2(&k_rho, &kd.0);
        for i in 0..2 {
            for j in 0..2 {
                result[i][j] += k_rho_kd[i][j];
            }
        }
    }
    result
}

/// Apply a unitary gate to a density matrix: ρ → U ρ U†
#[must_use]
pub fn apply_unitary(rho: &Matrix2x2, u: &Matrix2x2) -> Matrix2x2 {
    let ud = adjoint2x2(u);
    let u_rho = mul2x2(u, rho);
    mul2x2(&u_rho, &ud)
}

/// Trace of a 2×2 matrix.
#[must_use]
pub fn trace(m: &Matrix2x2) -> Complex64 {
    m[0][0] + m[1][1]
}

/// Check if Kraus operators satisfy trace preservation: Σᵢ Kᵢ† Kᵢ = I.
#[must_use]
pub fn is_trace_preserving(kraus: &[KrausOperator], tol: f64) -> bool {
    let mut sum = [[Complex64::new(0.0, 0.0); 2]; 2];
    for k in kraus {
        let kd = k.adjoint();
        let kdk = mul2x2(&kd.0, &k.0);
        for i in 0..2 {
            for j in 0..2 {
                sum[i][j] += kdk[i][j];
            }
        }
    }
    (sum[0][0] - Complex64::new(1.0, 0.0)).norm() < tol
        && (sum[0][1]).norm() < tol
        && (sum[1][0]).norm() < tol
        && (sum[1][1] - Complex64::new(1.0, 0.0)).norm() < tol
}

// ─────────────────────────────────────────────────────────────────────────────
// Noise Channels
// ─────────────────────────────────────────────────────────────────────────────

/// Single-qubit phase damping channel.
///
/// Kraus operators:
/// ```text
/// K₀ = [[1, 0], [0, √(1-γ)]]
/// K₁ = [[0, 0], [0, √γ]]
/// ```
///
/// Off-diagonal elements decay as ρ₀₁ → ρ₀₁ · √(1-γ).
///
/// # Panics
///
/// Panics if `gamma` is not in `[0, 1]`.
#[must_use]
pub fn dephasing(gamma: f64) -> Vec<KrausOperator> {
    assert!((0.0..=1.0).contains(&gamma), "gamma must be in [0, 1], got {gamma}");
    let zero = Complex64::new(0.0, 0.0);
    let one = Complex64::new(1.0, 0.0);
    let k0 = KrausOperator([
        [one, zero],
        [zero, Complex64::new((1.0 - gamma).sqrt(), 0.0)],
    ]);
    let k1 = KrausOperator([
        [zero, zero],
        [zero, Complex64::new(gamma.sqrt(), 0.0)],
    ]);
    vec![k0, k1]
}

/// Single-qubit amplitude damping channel.
///
/// Kraus operators:
/// ```text
/// K₀ = [[1, 0], [0, √(1-γ)]]
/// K₁ = [[0, √γ], [0, 0]]
/// ```
///
/// The excited state |1⟩ decays to |0⟩ with probability γ.
///
/// # Panics
///
/// Panics if `gamma` is not in `[0, 1]`.
#[must_use]
pub fn amplitude_damping(gamma: f64) -> Vec<KrausOperator> {
    assert!((0.0..=1.0).contains(&gamma), "gamma must be in [0, 1], got {gamma}");
    let zero = Complex64::new(0.0, 0.0);
    let one = Complex64::new(1.0, 0.0);
    let k0 = KrausOperator([
        [one, zero],
        [zero, Complex64::new((1.0 - gamma).sqrt(), 0.0)],
    ]);
    let k1 = KrausOperator([
        [zero, Complex64::new(gamma.sqrt(), 0.0)],
        [zero, zero],
    ]);
    vec![k0, k1]
}

/// Single-qubit symmetric depolarizing channel.
///
/// Each Pauli error (X, Y, Z) applied with probability p/3.
///
/// Kraus operators:
/// ```text
/// K₀ = √(1-p) · I
/// K₁ = √(p/3) · X
/// K₂ = √(p/3) · Y
/// K₃ = √(p/3) · Z
/// ```
///
/// # Panics
///
/// Panics if `p` is not in `[0, 1]`.
#[must_use]
pub fn depolarizing(p: f64) -> Vec<KrausOperator> {
    assert!((0.0..=1.0).contains(&p), "p must be in [0, 1], got {p}");
    let zero = Complex64::new(0.0, 0.0);
    let s0 = Complex64::new((1.0 - p).sqrt(), 0.0);
    let sp = Complex64::new((p / 3.0).sqrt(), 0.0);
    let im = Complex64::new(0.0, 1.0);

    let k0 = KrausOperator([[s0, zero], [zero, s0]]);                    // √(1-p) I
    let k1 = KrausOperator([[zero, sp], [sp, zero]]);                    // √(p/3) X
    let k2 = KrausOperator([[zero, -sp * im], [sp * im, zero]]);         // √(p/3) Y
    let k3 = KrausOperator([[sp, zero], [zero, -sp]]);                   // √(p/3) Z
    vec![k0, k1, k2, k3]
}

/// Spectral gap of the cycle graph Cₙ Laplacian.
///
/// λ₁ = 2 − 2cos(2π/n)
///
/// This is also the smallest non-zero eigenvalue of the n×n discrete torus
/// Laplacian Cₙ □ Cₙ, since product-graph eigenvalues are pairwise sums.
///
/// # Panics
///
/// Panics if `n < 2`.
#[must_use]
pub fn spectral_gap(n: usize) -> f64 {
    assert!(n >= 2, "n must be >= 2, got {n}");
    2.0 - 2.0 * (2.0 * PI / n as f64).cos()
}

/// Phenomenological effective dephasing rate from a toroidal lattice
/// spectral gap.
///
/// ```text
/// γ_eff = γ · λ₁(n) / (λ₁(n) + α)
/// ```
///
/// where λ₁(n) is the cycle-graph spectral gap. The mapping has the property
/// γ_eff → 0 as λ₁ → 0 and γ_eff → γ as λ₁ → ∞.
///
/// # Status
///
/// This mapping is **phenomenological**. It is offered as a parameterization
/// for studying lattice-geometry-dependent dephasing, not as a derivation
/// of dephasing rates from physical first principles. Use `alpha` as a
/// tunable knob, not as a physically calibrated coupling.
///
/// # Panics
///
/// Panics if `gamma` is not in `[0, 1]`, `grid_n < 2`, or `alpha <= 0.0`.
#[must_use]
pub fn effective_gamma(gamma: f64, grid_n: usize, alpha: f64) -> f64 {
    assert!((0.0..=1.0).contains(&gamma), "gamma must be in [0, 1], got {gamma}");
    assert!(grid_n >= 2, "grid_n must be >= 2, got {grid_n}");
    assert!(alpha > 0.0, "alpha must be positive, got {alpha}");
    let lambda1 = spectral_gap(grid_n);
    gamma * lambda1 / (lambda1 + alpha)
}

/// Single-qubit dephasing parameterized by a toroidal lattice spectral gap.
///
/// Returns the same Kraus operators as [`dephasing`] but with `gamma`
/// replaced by [`effective_gamma`]`(gamma, grid_n, alpha)`. This function
/// is a thin wrapper: it delegates the Kraus construction to [`dephasing`].
///
/// | grid_n | λ₁(n)   | γ_eff/γ (α=1) |
/// |--------|---------|----------------|
/// | 4      | 2.000   | 0.667          |
/// | 6      | 1.000   | 0.500          |
/// | 8      | 0.586   | 0.369          |
/// | 12     | 0.268   | 0.211          |
/// | 32     | 0.0383  | 0.0369         |
///
/// # Status
///
/// Phenomenological — see [`effective_gamma`] for the caveat. The
/// spectral-gap-to-dephasing mapping does not correspond to any particular
/// standard physical mechanism and should not be interpreted as a hardware
/// prediction without independent justification.
///
/// # Panics
///
/// Panics if `gamma` is not in `[0, 1]`, `grid_n < 2`, or `alpha <= 0.0`.
#[must_use]
pub fn toroidal_dephasing(gamma: f64, grid_n: usize, alpha: f64) -> Vec<KrausOperator> {
    dephasing(effective_gamma(gamma, grid_n, alpha))
}

// ─────────────────────────────────────────────────────────────────────────────
// Common gates (for density matrix simulation)
// ─────────────────────────────────────────────────────────────────────────────

/// Hadamard gate.
pub const HADAMARD: Matrix2x2 = {
    let s = FRAC_1_SQRT_2;
    [
        [Complex64::new(s, 0.0), Complex64::new(s, 0.0)],
        [Complex64::new(s, 0.0), Complex64::new(-s, 0.0)],
    ]
};

/// Pauli X gate.
pub const SIGMA_X: Matrix2x2 = [
    [Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)],
    [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)],
];

/// |0⟩⟨0| density matrix.
pub const RHO_ZERO: Matrix2x2 = [
    [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)],
    [Complex64::new(0.0, 0.0), Complex64::new(0.0, 0.0)],
];

// ─────────────────────────────────────────────────────────────────────────────
// Internal 2×2 matrix arithmetic
// ─────────────────────────────────────────────────────────────────────────────

fn mul2x2(a: &Matrix2x2, b: &Matrix2x2) -> Matrix2x2 {
    [
        [
            a[0][0] * b[0][0] + a[0][1] * b[1][0],
            a[0][0] * b[0][1] + a[0][1] * b[1][1],
        ],
        [
            a[1][0] * b[0][0] + a[1][1] * b[1][0],
            a[1][0] * b[0][1] + a[1][1] * b[1][1],
        ],
    ]
}

fn adjoint2x2(m: &Matrix2x2) -> Matrix2x2 {
    [
        [m[0][0].conj(), m[1][0].conj()],
        [m[0][1].conj(), m[1][1].conj()],
    ]
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

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

    const TOL: f64 = 1e-12;

    fn approx_eq(a: Complex64, b: Complex64) -> bool {
        (a - b).norm() < TOL
    }

    fn approx_eq_f64(a: f64, b: f64) -> bool {
        (a - b).abs() < TOL
    }

    // --- Dephasing ---

    #[test]
    fn dephasing_gamma0_is_identity() {
        let k = dephasing(0.0);
        assert!(approx_eq(k[0].0[1][1], Complex64::new(1.0, 0.0)));
        assert!(approx_eq(k[1].0[1][1], Complex64::new(0.0, 0.0)));
    }

    #[test]
    fn dephasing_gamma1_full() {
        let k = dephasing(1.0);
        assert!(approx_eq(k[0].0[1][1], Complex64::new(0.0, 0.0)));
        assert!(approx_eq(k[1].0[1][1], Complex64::new(1.0, 0.0)));
    }

    #[test]
    fn dephasing_trace_preserving() {
        for &g in &[0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0] {
            assert!(is_trace_preserving(&dephasing(g), TOL));
        }
    }

    #[test]
    fn dephasing_off_diagonal_decay() {
        for &g in &[0.1, 0.3, 0.5, 0.9] {
            let rho = apply_unitary(&RHO_ZERO, &HADAMARD); // |+⟩⟨+|
            let rho_noisy = apply_channel(&rho, &dephasing(g));
            let expected = 0.5 * (1.0 - g).sqrt();
            assert!(approx_eq_f64(rho_noisy[0][1].re, expected));
        }
    }

    #[test]
    #[should_panic]
    fn dephasing_invalid_negative() {
        let _ = dephasing(-0.1);
    }

    #[test]
    #[should_panic]
    fn dephasing_invalid_above_one() {
        let _ = dephasing(1.5);
    }

    // --- Amplitude Damping ---

    #[test]
    fn amplitude_damping_gamma0_identity() {
        let k = amplitude_damping(0.0);
        assert!(is_trace_preserving(&k, TOL));
        assert!(approx_eq(k[0].0[1][1], Complex64::new(1.0, 0.0)));
    }

    #[test]
    fn amplitude_damping_full_decay() {
        let rho = apply_unitary(&RHO_ZERO, &SIGMA_X); // |1⟩⟨1|
        let rho_out = apply_channel(&rho, &amplitude_damping(1.0));
        assert!(approx_eq_f64(rho_out[0][0].re, 1.0)); // fully decayed to |0⟩
        assert!(approx_eq_f64(rho_out[1][1].re, 0.0));
    }

    #[test]
    fn amplitude_damping_partial() {
        let rho = apply_unitary(&RHO_ZERO, &SIGMA_X); // |1⟩⟨1|
        let rho_out = apply_channel(&rho, &amplitude_damping(0.3));
        assert!(approx_eq_f64(rho_out[0][0].re, 0.3));
        assert!(approx_eq_f64(rho_out[1][1].re, 0.7));
    }

    #[test]
    fn amplitude_damping_trace_preserving() {
        for &g in &[0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0] {
            assert!(is_trace_preserving(&amplitude_damping(g), TOL));
        }
    }

    // --- Depolarizing ---

    #[test]
    fn depolarizing_p0_identity() {
        let k = depolarizing(0.0);
        assert!(approx_eq(k[0].0[0][0], Complex64::new(1.0, 0.0)));
    }

    #[test]
    fn depolarizing_trace_preserving() {
        for &p in &[0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0] {
            assert!(is_trace_preserving(&depolarizing(p), TOL));
        }
    }

    #[test]
    fn depolarizing_full_on_ground() {
        let rho_out = apply_channel(&RHO_ZERO, &depolarizing(1.0));
        assert!(approx_eq_f64(rho_out[0][0].re, 1.0 / 3.0));
        assert!(approx_eq_f64(rho_out[1][1].re, 2.0 / 3.0));
    }

    #[test]
    fn depolarizing_returns_four_kraus() {
        assert_eq!(depolarizing(0.5).len(), 4);
    }

    // --- Spectral Gap ---

    #[test]
    fn spectral_gap_known_values() {
        assert!(approx_eq_f64(spectral_gap(2), 4.0));
        assert!(approx_eq_f64(spectral_gap(3), 3.0));
        assert!(approx_eq_f64(spectral_gap(4), 2.0));
        assert!(approx_eq_f64(spectral_gap(6), 1.0));
    }

    #[test]
    fn spectral_gap_monotone_decreasing() {
        let mut prev = spectral_gap(2);
        for n in [3, 4, 6, 8, 16, 32, 64] {
            let sg = spectral_gap(n);
            assert!(sg < prev, "spectral_gap({n}) = {sg} >= {prev}");
            prev = sg;
        }
    }

    #[test]
    fn spectral_gap_always_positive() {
        for n in [2, 3, 4, 8, 16, 64, 128] {
            assert!(spectral_gap(n) > 0.0);
        }
    }

    // --- Effective gamma ---

    #[test]
    fn effective_gamma_explicit_formula() {
        let gamma = 0.5_f64;
        let n = 8;
        let alpha = 1.5_f64;
        let lam = spectral_gap(n);
        assert!(approx_eq_f64(
            effective_gamma(gamma, n, alpha),
            gamma * lam / (lam + alpha),
        ));
    }

    #[test]
    fn effective_gamma_zero_in_zero_out() {
        assert_eq!(effective_gamma(0.0, 12, 1.0), 0.0);
    }

    #[test]
    fn effective_gamma_alpha_to_zero_recovers_gamma() {
        let gamma = 0.7_f64;
        let eg = effective_gamma(gamma, 12, 1e-12);
        assert!((eg - gamma).abs() < 1e-6);
    }

    #[test]
    #[should_panic]
    fn effective_gamma_alpha_zero_panics() {
        let _ = effective_gamma(0.5, 12, 0.0);
    }

    // --- Toroidal Dephasing ---

    #[test]
    fn toroidal_gamma0_identity() {
        let k = toroidal_dephasing(0.0, 12, 1.0);
        assert!(approx_eq(k[0].0[1][1], Complex64::new(1.0, 0.0)));
        assert!(approx_eq(k[1].0[1][1], Complex64::new(0.0, 0.0)));
    }

    #[test]
    fn toroidal_matches_dephasing_with_effective_gamma() {
        // Delegation contract: toroidal_dephasing(g, n, a) ≡ dephasing(effective_gamma(g, n, a))
        for &g in &[0.0, 0.25, 0.5, 0.75, 1.0] {
            for n in [4, 8, 12, 32] {
                for &a in &[0.5, 1.0, 2.0] {
                    let k_t = toroidal_dephasing(g, n, a);
                    let k_d = dephasing(effective_gamma(g, n, a));
                    assert_eq!(k_t.len(), k_d.len());
                    for (kt, kd) in k_t.iter().zip(k_d.iter()) {
                        for i in 0..2 {
                            for j in 0..2 {
                                assert!(approx_eq(kt.0[i][j], kd.0[i][j]));
                            }
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn toroidal_smaller_gap_means_smaller_eff_gamma() {
        // With this phenomenological mapping, smaller spectral gap (larger n)
        // produces smaller γ_eff. Note this is a property of the formula, not a
        // physical claim — see effective_gamma docstring.
        let rho = apply_unitary(&RHO_ZERO, &HADAMARD);
        let rho_plain = apply_channel(&rho, &dephasing(0.5));
        let rho_torus = apply_channel(&rho, &toroidal_dephasing(0.5, 12, 1.0));
        assert!(rho_torus[0][1].norm() > rho_plain[0][1].norm());
    }

    #[test]
    fn toroidal_larger_grid_smaller_eff_gamma() {
        let k_small = toroidal_dephasing(0.5, 4, 1.0);
        let k_large = toroidal_dephasing(0.5, 32, 1.0);
        let g_small = k_small[1].0[1][1].norm().powi(2);
        let g_large = k_large[1].0[1][1].norm().powi(2);
        assert!(g_large < g_small);
    }

    #[test]
    fn toroidal_monotonic_in_grid_n() {
        let mut prev_g = 1.0;
        for n in [4, 6, 8, 12, 16, 32, 64] {
            let k = toroidal_dephasing(1.0, n, 1.0);
            let g_eff = k[1].0[1][1].norm().powi(2);
            assert!(g_eff < prev_g);
            prev_g = g_eff;
        }
    }

    #[test]
    fn toroidal_known_value() {
        let k = toroidal_dephasing(1.0, 4, 1.0);
        let g_eff = k[1].0[1][1].norm().powi(2);
        let l1 = spectral_gap(4);
        let expected = l1 / (l1 + 1.0);
        assert!(approx_eq_f64(g_eff, expected));
    }

    #[test]
    fn toroidal_trace_preserving() {
        for &g in &[0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0] {
            for n in [2, 4, 6, 8, 12, 32] {
                assert!(is_trace_preserving(&toroidal_dephasing(g, n, 1.0), TOL));
            }
        }
    }

    #[test]
    fn toroidal_analytical_value() {
        let gamma = 0.5;
        let n = 12;
        let alpha = 1.0;
        let l1 = spectral_gap(n);
        let g_eff = gamma * l1 / (l1 + alpha);
        let rho = apply_unitary(&RHO_ZERO, &HADAMARD);
        let rho_out = apply_channel(&rho, &toroidal_dephasing(gamma, n, alpha));
        let expected_off_diag = 0.5 * (1.0 - g_eff).sqrt();
        assert!(approx_eq_f64(rho_out[0][1].re, expected_off_diag));
    }

    #[test]
    #[should_panic]
    fn toroidal_invalid_grid() {
        let _ = toroidal_dephasing(0.5, 1, 1.0);
    }

    #[test]
    #[should_panic]
    fn toroidal_invalid_alpha() {
        let _ = toroidal_dephasing(0.5, 12, 0.0);
    }

    // --- Density matrix properties ---

    #[test]
    fn density_matrix_hermitian_after_noise() {
        let rho = apply_unitary(&RHO_ZERO, &HADAMARD);
        let rho_out = apply_channel(&rho, &dephasing(0.3));
        assert!(approx_eq(rho_out[0][1], rho_out[1][0].conj()));
    }

    #[test]
    fn density_matrix_unit_trace_after_noise() {
        for &g in &[0.0, 0.1, 0.5, 1.0] {
            let rho = apply_unitary(&RHO_ZERO, &HADAMARD);
            let rho_out = apply_channel(&rho, &dephasing(g));
            assert!(approx_eq_f64(trace(&rho_out).re, 1.0));
        }
    }

    #[test]
    fn composed_noise() {
        let rho = apply_unitary(&RHO_ZERO, &HADAMARD);
        let rho = apply_channel(&rho, &dephasing(0.3));
        let rho = apply_channel(&rho, &amplitude_damping(0.2));
        assert!(approx_eq_f64(trace(&rho).re, 1.0));
    }

    #[test]
    fn full_dephasing_destroys_coherence() {
        let rho = apply_unitary(&RHO_ZERO, &HADAMARD);
        let rho_out = apply_channel(&rho, &dephasing(1.0));
        assert!(rho_out[0][1].norm() < TOL);
        assert!(rho_out[1][0].norm() < TOL);
        assert!(approx_eq_f64(rho_out[0][0].re, 0.5));
        assert!(approx_eq_f64(rho_out[1][1].re, 0.5));
    }
}