toroidal-noise 0.1.0

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

use num_complex::Complex64;
use std::f64::consts::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 nonzero eigenvalue of the n×n 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()
}

/// Single-qubit dephasing with toroidal spectral-gap suppression.
///
/// The effective dephasing probability is reduced by:
///
/// ```text
/// γ_eff = γ · λ₁ / (λ₁ + α)
/// ```
///
/// where λ₁ = 2 - 2cos(2π/n) is the spectral gap of Cₙ and
/// α > 0 controls coupling strength.
///
/// Larger `grid_n` → smaller spectral gap → stronger noise suppression.
///
/// | grid_n | noise reduction (α=1) |
/// |--------|----------------------|
/// | 4      | 1.5×                 |
/// | 8      | 2.7×                 |
/// | 12     | 4.7×                 |
/// | 32     | 27×                  |
///
/// Reference: [Cormier 2026](https://doi.org/10.5281/zenodo.18516477)
///
/// # Panics
///
/// Panics if `gamma` not in `[0, 1]`, `grid_n < 2`, or `alpha <= 0`.
#[must_use]
pub fn toroidal_dephasing(gamma: f64, grid_n: usize, alpha: f64) -> Vec<KrausOperator> {
    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);
    let gamma_eff = gamma * lambda1 / (lambda1 + alpha);

    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_eff).sqrt(), 0.0)],
    ]);
    let k1 = KrausOperator([
        [zero, zero],
        [zero, Complex64::new(gamma_eff.sqrt(), 0.0)],
    ]);
    vec![k0, k1]
}

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

/// Hadamard gate.
pub const HADAMARD: Matrix2x2 = {
    let s = 0.707_106_781_186_547_6; // 1/√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);
        }
    }

    // --- 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_suppresses_noise() {
        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_more_suppression() {
        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_suppression() {
        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_suppression() {
        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));
    }
}