nufast 0.3.0

//! # NuFast
//!
//! Fast and accurate three-flavor neutrino oscillation probabilities in vacuum and matter.
//!
//! This is a Rust port of [NuFast](https://github.com/PeterDenton/NuFast) by Peter Denton.
//!
//! ## Features
//!
//! - **Vacuum oscillations**: Full 3-flavor PMNS oscillation probabilities
//! - **Matter effects (MSW)**: Constant-density matter with arbitrary electron fraction
//! - **CP violation**: Full δ_CP phase support
//! - **Newton refinement**: Optional iterative improvement for matter eigenvalues
//! - **Zero dependencies**: Pure Rust, no external crates required
//! - **Fast**: ~61 ns vacuum, ~95 ns matter (27% faster than C++ for matter!)
//!
//! ## Quick Start
//!
//! ```rust
//! use nufast::{VacuumParameters, probability_vacuum_lbl};
//! use std::f64::consts::PI;
//!
//! let params = VacuumParameters {
//!     s12sq: 0.307,
//!     s13sq: 0.0218,
//!     s23sq: 0.545,
//!     delta: 1.36 * PI,
//!     Dmsq21: 7.42e-5,
//!     Dmsq31: 2.517e-3,
//!     L: 295.0,
//!     E: 0.6,
//! };
//!
//! let probs = probability_vacuum_lbl(&params);
//! assert!(probs[1][0] >= 0.0 && probs[1][0] <= 1.0); // P(νμ → νe)
//! ```
//!
//! ## Performance
//!
//! Benchmarks on AMD Ryzen (10M iterations):
//!
//! | Language | Vacuum | Matter N=0 |
//! |----------|--------|------------|
//! | **Rust** | 61 ns  | **95 ns**  |
//! | C++      | 49 ns  | 130 ns     |
//! | Fortran  | 51 ns  | 107 ns     |
//! | Python   | 14,700 ns | 21,900 ns |
//!
//! Rust is **27% faster than C++** for matter calculations.

#![cfg_attr(feature = "no_std", no_std)]
// Allow physics naming conventions (standard in neutrino oscillation literature)
#![allow(non_snake_case)]

use core::f64::consts::PI;

/// Conversion factor: eV² × km → GeV (divided by 4)
/// 
/// This combines ħc and unit conversions for the oscillation phase:
/// Δ = Δm² × L / (4E) in natural units
const EV_SQ_KM_TO_GEV_OVER4: f64 = 1e-9 / 1.97327e-7 * 1e3 / 4.0;

/// Matter potential conversion factor: Y_e × ρ × E → A
/// 
/// A = √2 G_F N_e E where N_e = Y_e × ρ × N_A / m_nucleon
/// This constant is approximately 1.52 × 10⁻⁴ eV² / (g/cm³ × GeV)
const YE_RHO_E_TO_A: f64 = 1.52e-4;

/// Parameters for vacuum oscillation calculations.
///
/// All angles are specified as sin²θ, CP phase in radians.
/// Mass splittings in eV², baseline in km, energy in GeV.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct VacuumParameters {
    /// sin²θ₁₂ (solar mixing angle)
    pub s12sq: f64,
    /// sin²θ₁₃ (reactor mixing angle)
    pub s13sq: f64,
    /// sin²θ₂₃ (atmospheric mixing angle)
    pub s23sq: f64,
    /// CP-violating phase δ in radians
    pub delta: f64,
    /// Δm²₂₁ in eV² (solar mass splitting, always positive)
    pub Dmsq21: f64,
    /// Δm²₃₁ in eV² (atmospheric mass splitting, positive for NO, negative for IO)
    pub Dmsq31: f64,
    /// Baseline L in km
    pub L: f64,
    /// Neutrino energy E in GeV
    pub E: f64,
}

impl VacuumParameters {
    /// Create parameters with NuFit 5.2 best-fit values (Normal Ordering).
    ///
    /// Only L and E need to be specified.
    pub fn nufit52_no(l: f64, e: f64) -> Self {
        Self {
            s12sq: 0.307,
            s13sq: 0.02203,
            s23sq: 0.546,
            delta: 1.36 * PI,
            Dmsq21: 7.42e-5,
            Dmsq31: 2.517e-3,
            L: l,
            E: e,
        }
    }

    /// Create parameters with NuFit 5.2 best-fit values (Inverted Ordering).
    pub fn nufit52_io(l: f64, e: f64) -> Self {
        Self {
            s12sq: 0.307,
            s13sq: 0.02219,
            s23sq: 0.539,
            delta: 1.56 * PI,
            Dmsq21: 7.42e-5,
            Dmsq31: -2.498e-3,
            L: l,
            E: e,
        }
    }
}

/// Parameters for matter oscillation calculations.
///
/// Extends [VacuumParameters] with matter density and electron fraction.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MatterParameters {
    /// sin²θ₁₂ (solar mixing angle)
    pub s12sq: f64,
    /// sin²θ₁₃ (reactor mixing angle)
    pub s13sq: f64,
    /// sin²θ₂₃ (atmospheric mixing angle)
    pub s23sq: f64,
    /// CP-violating phase δ in radians
    pub delta: f64,
    /// Δm²₂₁ in eV² (solar mass splitting)
    pub Dmsq21: f64,
    /// Δm²₃₁ in eV² (atmospheric mass splitting)
    pub Dmsq31: f64,
    /// Baseline L in km
    pub L: f64,
    /// Neutrino energy E in GeV
    pub E: f64,
    /// Matter density ρ in g/cm³
    pub rho: f64,
    /// Electron fraction Y_e (typically 0.5 for Earth)
    pub Ye: f64,
    /// Number of Newton iterations for eigenvalue refinement (0 = DMP approximation)
    pub N_Newton: u8,
}

impl MatterParameters {
    /// Create parameters with NuFit 5.2 best-fit values (Normal Ordering).
    ///
    /// Uses Earth crust density (2.6 g/cm³) and Y_e = 0.5.
    pub fn nufit52_no(l: f64, e: f64) -> Self {
        Self {
            s12sq: 0.307,
            s13sq: 0.02203,
            s23sq: 0.546,
            delta: 1.36 * PI,
            Dmsq21: 7.42e-5,
            Dmsq31: 2.517e-3,
            L: l,
            E: e,
            rho: 2.6,
            Ye: 0.5,
            N_Newton: 0,
        }
    }

    /// Create from vacuum parameters, adding matter properties.
    pub fn from_vacuum(vac: &VacuumParameters, rho: f64, ye: f64, n_newton: u8) -> Self {
        Self {
            s12sq: vac.s12sq,
            s13sq: vac.s13sq,
            s23sq: vac.s23sq,
            delta: vac.delta,
            Dmsq21: vac.Dmsq21,
            Dmsq31: vac.Dmsq31,
            L: vac.L,
            E: vac.E,
            rho,
            Ye: ye,
            N_Newton: n_newton,
        }
    }
}

/// 3×3 probability matrix type.
///
/// Indexed as `probs[α][β]` = P(ν_α → ν_β) where:
/// - 0 = electron (e)
/// - 1 = muon (μ)
/// - 2 = tau (τ)
pub type ProbabilityMatrix = [[f64; 3]; 3];

/// Calculate vacuum oscillation probabilities for long-baseline experiments.
///
/// Returns a 3×3 matrix where `probs[α][β]` = P(ν_α → ν_β).
///
/// # Example
///
/// ```rust
/// use nufast::{VacuumParameters, probability_vacuum_lbl};
///
/// let params = VacuumParameters::nufit52_no(295.0, 0.6);
/// let probs = probability_vacuum_lbl(&params);
///
/// // Check unitarity
/// let row_sum: f64 = probs[1].iter().sum();
/// assert!((row_sum - 1.0).abs() < 1e-10);
/// ```
pub fn probability_vacuum_lbl(parameters: &VacuumParameters) -> ProbabilityMatrix {
    let VacuumParameters {
        s12sq,
        s13sq,
        s23sq,
        delta,
        Dmsq21,
        Dmsq31,
        L,
        E,
    } = *parameters;

    // Pre-calculate trig functions
    let c13sq = 1.0 - s13sq;

    // Ueisq's
    let Ue3sq = s13sq;
    let Ue2sq = c13sq * s12sq;

    // Umisq's, Utisq's and Jvac
    let Um3sq = c13sq * s23sq;
    // Um2sq and Ut2sq are used here as temporary variables
    let Ut2sq = s13sq * s12sq * s23sq;
    let Um2sq = (1.0 - s12sq) * (1.0 - s23sq);

    let Jrr = (Um2sq * Ut2sq).sqrt();
    let sind = delta.sin();
    let cosd = delta.cos();
    let Um2sq = Um2sq + Ut2sq - 2.0 * Jrr * cosd;
    let Jvac = 8.0 * Jrr * c13sq * sind;

    // Get all elements of Usq using unitarity
    let Ue1sq = 1.0 - Ue3sq - Ue2sq;
    let Um1sq = 1.0 - Um3sq - Um2sq;

    let Ut3sq = 1.0 - Um3sq - Ue3sq;
    let Ut2sq = 1.0 - Um2sq - Ue2sq;
    let Ut1sq = 1.0 - Um1sq - Ue1sq;

    // Get the kinematic terms
    let Lover4E = EV_SQ_KM_TO_GEV_OVER4 * L / E;

    let D21 = Dmsq21 * Lover4E;
    let D31 = Dmsq31 * Lover4E;

    let sinD21 = D21.sin();
    let sinD31 = D31.sin();
    let sinD32 = (D31 - D21).sin();

    let triple_sin = sinD21 * sinD31 * sinD32;

    let sinsqD21_2 = 2.0 * sinD21 * sinD21;
    let sinsqD31_2 = 2.0 * sinD31 * sinD31;
    let sinsqD32_2 = 2.0 * sinD32 * sinD32;

    // Calculate the three probabilities, separating CPC and CPV
    let Pme_CPC = (Ut3sq - Um2sq * Ue1sq - Um1sq * Ue2sq) * sinsqD21_2
        + (Ut2sq - Um3sq * Ue1sq - Um1sq * Ue3sq) * sinsqD31_2
        + (Ut1sq - Um3sq * Ue2sq - Um2sq * Ue3sq) * sinsqD32_2;

    let Pme_CPV = -Jvac * triple_sin;

    let Pmm = 1.0
        - 2.0
            * (Um2sq * Um1sq * sinsqD21_2
                + Um3sq * Um1sq * sinsqD31_2
                + Um3sq * Um2sq * sinsqD32_2);

    let Pee = 1.0
        - 2.0
            * (Ue2sq * Ue1sq * sinsqD21_2
                + Ue3sq * Ue1sq * sinsqD31_2
                + Ue3sq * Ue2sq * sinsqD32_2);

    // Assign all the probabilities
    let mut probs = [[0.0; 3]; 3];

    probs[0][0] = Pee; // Pee
    probs[0][1] = Pme_CPC - Pme_CPV; // Pem
    probs[0][2] = 1.0 - Pee - probs[0][1]; // Pet

    probs[1][0] = Pme_CPC + Pme_CPV; // Pme
    probs[1][1] = Pmm; // Pmm
    probs[1][2] = 1.0 - probs[1][0] - Pmm; // Pmt

    probs[2][0] = 1.0 - Pee - probs[1][0]; // Pte
    probs[2][1] = 1.0 - probs[0][1] - Pmm; // Ptm
    probs[2][2] = 1.0 - probs[0][2] - probs[1][2]; // Ptt

    probs
}

/// Calculate matter oscillation probabilities for long-baseline experiments.
///
/// Uses the DMP approximation with optional Newton refinement for the
/// matter-modified eigenvalues.
///
/// # Example
///
/// ```rust
/// use nufast::{MatterParameters, probability_matter_lbl};
///
/// let params = MatterParameters::nufit52_no(1300.0, 2.5);
/// let probs = probability_matter_lbl(&params);
///
/// // Matter enhances νμ → νe for normal ordering
/// assert!(probs[1][0] > 0.0);
/// ```
pub fn probability_matter_lbl(parameters: &MatterParameters) -> ProbabilityMatrix {
    let MatterParameters {
        s12sq,
        s13sq,
        s23sq,
        delta,
        Dmsq21,
        Dmsq31,
        L,
        E,
        rho,
        Ye,
        N_Newton,
    } = *parameters;

    // Pre-calculate trig functions
    let c13sq = 1.0 - s13sq;

    // Ueisq's
    let Ue2sq = c13sq * s12sq;
    let Ue3sq = s13sq;

    // Umisq's, Utisq's and Jmatter
    let Um3sq = c13sq * s23sq;
    let Ut2sq = s13sq * s12sq * s23sq;
    let Um2sq = (1.0 - s12sq) * (1.0 - s23sq);

    let Jrr = (Um2sq * Ut2sq).sqrt();
    let sind = delta.sin();
    let cosd = delta.cos();

    let Um2sq = Um2sq + Ut2sq - 2.0 * Jrr * cosd;
    let Jmatter = 8.0 * Jrr * c13sq * sind;
    let Amatter = Ye * rho * E * YE_RHO_E_TO_A;
    let Dmsqee = Dmsq31 - s12sq * Dmsq21;

    // Calculate A, B, C, See, Tee, and part of Tmm
    let A_sum = Dmsq21 + Dmsq31; // temporary variable
    let See = A_sum - Dmsq21 * Ue2sq - Dmsq31 * Ue3sq;
    let Tmm_base = Dmsq21 * Dmsq31;
    let Tee = Tmm_base * (1.0 - Ue3sq - Ue2sq);
    let C = Amatter * Tee;
    let A = A_sum + Amatter;

    // Get lambda3 from lambda+ of MP/DMP
    let xmat = Amatter / Dmsqee;
    let tmp = 1.0 - xmat;
    let mut lambda3 = Dmsq31 + 0.5 * Dmsqee * (xmat - 1.0 + (tmp * tmp + 4.0 * s13sq * xmat).sqrt());

    // Newton iterations to improve lambda3 arbitrarily
    let B = Tmm_base + Amatter * See; // B is only needed for N_Newton >= 1
    for _ in 0..N_Newton {
        lambda3 = (lambda3 * lambda3 * (lambda3 - A) + C)
            / (lambda3 * (2.0 * lambda3 - A) + B);
    }

    // Get Delta lambdas
    let tmp = A - lambda3;
    let Dlambda21 = (tmp * tmp - 4.0 * C / lambda3).sqrt();
    let lambda2 = 0.5 * (A - lambda3 + Dlambda21);
    let Dlambda32 = lambda3 - lambda2;
    let Dlambda31 = Dlambda32 + Dlambda21;

    // Use Rosetta for Veisq's (matter-modified PMNS elements)
    let PiDlambdaInv = 1.0 / (Dlambda31 * Dlambda32 * Dlambda21);
    let Xp3 = PiDlambdaInv * Dlambda21;
    let Xp2 = -PiDlambdaInv * Dlambda31;

    let Ue3sq = (lambda3 * (lambda3 - See) + Tee) * Xp3;
    let Ue2sq = (lambda2 * (lambda2 - See) + Tee) * Xp2;

    let Smm = A - Dmsq21 * Um2sq - Dmsq31 * Um3sq;
    let Tmm = Tmm_base * (1.0 - Um3sq - Um2sq) + Amatter * (See + Smm - A_sum);

    let Um3sq = (lambda3 * (lambda3 - Smm) + Tmm) * Xp3;
    let Um2sq = (lambda2 * (lambda2 - Smm) + Tmm) * Xp2;

    // Use NHS for J (matter-modified Jarlskog)
    let Jmatter = Jmatter * Dmsq21 * Dmsq31 * (Dmsq31 - Dmsq21) * PiDlambdaInv;

    // Get all elements of Usq
    let Ue1sq = 1.0 - Ue3sq - Ue2sq;
    let Um1sq = 1.0 - Um3sq - Um2sq;

    let Ut3sq = 1.0 - Um3sq - Ue3sq;
    let Ut2sq = 1.0 - Um2sq - Ue2sq;
    let Ut1sq = 1.0 - Um1sq - Ue1sq;

    // Get the kinematic terms (using matter eigenvalues)
    let Lover4E = EV_SQ_KM_TO_GEV_OVER4 * L / E;

    let D21 = Dlambda21 * Lover4E;
    let D32 = Dlambda32 * Lover4E;

    let sinD21 = D21.sin();
    let sinD31 = (D32 + D21).sin();
    let sinD32 = D32.sin();

    let triple_sin = sinD21 * sinD31 * sinD32;

    let sinsqD21_2 = 2.0 * sinD21 * sinD21;
    let sinsqD31_2 = 2.0 * sinD31 * sinD31;
    let sinsqD32_2 = 2.0 * sinD32 * sinD32;

    // Calculate the three necessary probabilities
    let Pme_CPC = (Ut3sq - Um2sq * Ue1sq - Um1sq * Ue2sq) * sinsqD21_2
        + (Ut2sq - Um3sq * Ue1sq - Um1sq * Ue3sq) * sinsqD31_2
        + (Ut1sq - Um3sq * Ue2sq - Um2sq * Ue3sq) * sinsqD32_2;
    let Pme_CPV = -Jmatter * triple_sin;

    let Pmm = 1.0
        - 2.0
            * (Um2sq * Um1sq * sinsqD21_2
                + Um3sq * Um1sq * sinsqD31_2
                + Um3sq * Um2sq * sinsqD32_2);

    let Pee = 1.0
        - 2.0
            * (Ue2sq * Ue1sq * sinsqD21_2
                + Ue3sq * Ue1sq * sinsqD31_2
                + Ue3sq * Ue2sq * sinsqD32_2);

    // Assign all the probabilities
    let mut probs = [[0.0; 3]; 3];

    probs[0][0] = Pee;
    probs[0][1] = Pme_CPC - Pme_CPV;
    probs[0][2] = 1.0 - Pee - probs[0][1];

    probs[1][0] = Pme_CPC + Pme_CPV;
    probs[1][1] = Pmm;
    probs[1][2] = 1.0 - probs[1][0] - Pmm;

    probs[2][0] = 1.0 - Pee - probs[1][0];
    probs[2][1] = 1.0 - probs[0][1] - Pmm;
    probs[2][2] = 1.0 - probs[0][2] - probs[1][2];

    probs
}

/// Clamp all probabilities to [0, 1] and ensure row unitarity.
/// 
/// Useful when numerical precision issues cause slight violations.
pub fn normalize_probabilities(probs: &mut ProbabilityMatrix) {
    for row in probs.iter_mut() {
        for p in row.iter_mut() {
            *p = p.clamp(0.0, 1.0);
        }
        let sum: f64 = row.iter().sum();
        if sum > 0.0 && (sum - 1.0).abs() > 1e-10 {
            for p in row.iter_mut() {
                *p /= sum;
            }
        }
    }
}

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

    const EPSILON: f64 = 1e-10;

    #[test]
    fn test_vacuum_unitarity() {
        let params = VacuumParameters::nufit52_no(1000.0, 2.0);
        let probs = probability_vacuum_lbl(&params);

        // Check all rows sum to 1
        for (i, row) in probs.iter().enumerate() {
            let sum: f64 = row.iter().sum();
            assert!(
                (sum - 1.0).abs() < EPSILON,
                "Row {} sum = {}, expected 1.0",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_vacuum_zero_distance() {
        let params = VacuumParameters::nufit52_no(0.0, 2.0);
        let probs = probability_vacuum_lbl(&params);

        // At L=0, no oscillation: diagonal = 1, off-diagonal = 0
        for i in 0..3 {
            for j in 0..3 {
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!(
                    (probs[i][j] - expected).abs() < EPSILON,
                    "probs[{}][{}] = {}, expected {}",
                    i,
                    j,
                    probs[i][j],
                    expected
                );
            }
        }
    }

    #[test]
    fn test_matter_unitarity() {
        let params = MatterParameters::nufit52_no(1300.0, 2.5);
        let mut probs = probability_matter_lbl(&params);
        normalize_probabilities(&mut probs);

        for (i, row) in probs.iter().enumerate() {
            let sum: f64 = row.iter().sum();
            assert!(
                (sum - 1.0).abs() < 1e-6,
                "Row {} sum = {}, expected 1.0",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_matter_enhances_appearance() {
        // For normal ordering, matter enhances νμ → νe at DUNE energies
        let vac = VacuumParameters::nufit52_no(1300.0, 2.5);
        let mat = MatterParameters::nufit52_no(1300.0, 2.5);

        let probs_vac = probability_vacuum_lbl(&vac);
        let probs_mat = probability_matter_lbl(&mat);

        // Matter should modify the appearance probability
        assert!(
            (probs_mat[1][0] - probs_vac[1][0]).abs() > 0.001,
            "Matter effect should modify P(νμ → νe)"
        );
    }

    #[test]
    fn test_cp_conjugate() {
        // P(ν, δ) should differ from P(ν, -δ) due to CP violation
        let mut params_pos = VacuumParameters::nufit52_no(1000.0, 2.0);
        params_pos.delta = 1.36 * PI;

        let mut params_neg = params_pos;
        params_neg.delta = -1.36 * PI;

        let probs_pos = probability_vacuum_lbl(&params_pos);
        let probs_neg = probability_vacuum_lbl(&params_neg);

        // P(μ→e, δ) ≠ P(μ→e, -δ) unless δ = 0 or π
        assert!(
            (probs_pos[1][0] - probs_neg[1][0]).abs() > 1e-6,
            "CP conjugate probabilities should differ"
        );
    }
}