lcdm-core 0.1.0

// lcdm-core/src/neutrino.rs

use crate::constants::{C, EV, G, H0_UNITS_TO_INV_SEC, H_PLANCK, KB};
use crate::spec::{AccuracyPreset, DerivedParams, NeutrinoSpec};

/// Trait for neutrino-background calculations.
///
/// Implementations return the *physical* neutrino density:
///   ω_ν(z) ≡ Ω_ν(z) h²
/// expressed in the same convention as `DerivedParams`.
pub trait NeutrinoBackground: Send + Sync {
    /// Neutrino physical density ω_ν(z) at redshift `z`.
    fn omega_nu_z(&self, z: f64) -> f64;

    /// Present-day (z=0) neutrino physical density ω_ν0.
    fn omega_nu0(&self) -> f64;
}

// -----------------------------------------------------------------------------
// Core integration logic
// -----------------------------------------------------------------------------

/// Composite Simpson's-rule integrator.
///
/// Ensures an even number of sub-intervals and returns an approximation to:
///   ∫[min..max] func(x) dx
fn integrate_simpson<F>(func: F, min: f64, max: f64, steps: usize) -> f64
where
    F: Fn(f64) -> f64,
{
    // Simpson's rule requires an even number of intervals.
    let n = if steps % 2 == 0 { steps } else { steps + 1 };
    let h = (max - min) / n as f64;

    // Endpoints have weight 1.
    let mut s = func(min) + func(max);

    // Interior points alternate weights 4 and 2.
    for i in 1..n {
        let x = min + i as f64 * h;
        let w = if i % 2 == 0 { 2.0_f64 } else { 4.0_f64 };
        s += w * func(x);
    }

    s * h / 3.0_f64
}

/// Compute neutrino physical density ω_ν(z) with explicit numerical settings.
///
/// This evaluates the Fermi–Dirac energy-density integral for each mass in `masses`
/// at temperature `T_nu0` scaled to redshift `z`, and returns the result as a
/// *physical density* normalized to ρ_crit(h=1) (i.e., ω ≡ Ω h²).
///
/// # Parameters
/// - `masses`: Neutrino masses in eV.
/// - `T_nu0`: Present-day neutrino temperature in Kelvin.
/// - `z`: Redshift.
/// - `y_max`: Upper bound for the dimensionless momentum variable in the integral.
/// - `steps`: Integration steps (Simpson composite; rounded to even internally).
/// - `degeneracy_g`: Internal degrees of freedom per species (commonly 2.0).
pub fn calc_omega_nu_z_cfg(
    masses: &[f64],
    T_nu0: f64,
    z: f64,
    y_max: f64,
    steps: usize,
    degeneracy_g: f64,
) -> f64 {
    let mut total_rho = 0.0_f64;

    // Temperature scales as T(z) = T0 (1+z).
    let T_nu_z = T_nu0 * (1.0_f64 + z);
    if T_nu_z <= 0.0_f64 {
        return 0.0_f64;
    }

    // Prefactor for the Fermi–Dirac integral.
    // Standard form:
    //   ρ = g/(2π²) * (kT)^4 / (ħc)^3 * ∫ dy y² √(y² + (mc²/kT)²) / (e^y + 1)
    //
    // Here we use h instead of ħ and fold constants into a convenient prefactor.
    let kBT = KB * T_nu_z;
    let pre_const =
        (4.0_f64 * std::f64::consts::PI * degeneracy_g) / (H_PLANCK.powi(3) * C.powi(3));
    let pre = pre_const * kBT.powi(4);

    // Sum contributions from each mass eigenstate.
    for &m_ev in masses {
        // Convert eV -> kg via E = mc^2 and 1 eV = EV Joules.
        let m_kg = m_ev * EV / C.powi(2);
        let m_c2 = m_kg * C * C;

        // Dimensionless squared mass ratio: (mc^2 / kT)^2.
        let mass_ratio_sq = (m_c2 / kBT).powi(2);

        // Integrand over dimensionless momentum y = p/(kT).
        let integrand = |y: f64| -> f64 {
            let numer = y * y * (y * y + mass_ratio_sq).sqrt();

            // Use a numerically stable approximation for large y.
            if y > 50.0_f64 {
                let exp_neg_y = (-y).exp();
                numer * exp_neg_y / (1.0_f64 + exp_neg_y)
            } else {
                let denom = y.exp() + 1.0_f64;
                numer / denom
            }
        };

        let val = integrate_simpson(integrand, 0.0_f64, y_max, steps);

        // Convert energy density to mass density via division by c^2.
        total_rho += (pre * val) / (C * C);
    }

    // Normalize by critical density for h=1:
    //   ρ_crit(h=1) = 3 (H100)^2 / (8πG), with H100 = 100 (km/s)/Mpc in SI.
    let h100_si = 100.0_f64 * H0_UNITS_TO_INV_SEC;
    let rho_crit_100 = 3.0_f64 * h100_si * h100_si / (8.0_f64 * std::f64::consts::PI * G);

    // Return ω_ν(z) = ρ_ν(z) / ρ_crit(h=1).
    total_rho / rho_crit_100
}

/// Legacy wrapper for `calc_omega_nu_z_cfg` using default numerical settings.
pub fn calc_omega_nu_z(masses: &[f64], T_nu0: f64, z: f64) -> f64 {
    calc_omega_nu_z_cfg(masses, T_nu0, z, 45.0, 1200, 2.0)
}

// -----------------------------------------------------------------------------
// PCHIP slopes (Fritsch–Carlson monotone cubic interpolation)
// -----------------------------------------------------------------------------

/// Compute PCHIP slopes using the Fritsch–Carlson / Fritsch–Butland construction.
///
/// This produces a shape-preserving (monotone) cubic interpolant when `y` is monotone.
fn compute_pchip_slopes(x: &[f64], y: &[f64]) -> Vec<f64> {
    let n = x.len();

    // Interval widths and secant slopes.
    let mut h = vec![0.0_f64; n - 1];
    let mut d = vec![0.0_f64; n - 1];
    for i in 0..n - 1 {
        h[i] = x[i + 1] - x[i];
        d[i] = (y[i + 1] - y[i]) / h[i];
    }

    let mut m = vec![0.0_f64; n];

    // Two-point edge case: constant slope everywhere.
    if n == 2 {
        m[0] = d[0];
        m[1] = d[0];
        return m;
    }

    // 1) Internal points: weighted harmonic mean when slopes agree in sign.
    for i in 1..n - 1 {
        let d0 = d[i - 1];
        let d1 = d[i];
        if d0 * d1 > 0.0_f64 {
            let w1 = 2.0_f64 * h[i] + h[i - 1];
            let w2 = h[i] + 2.0_f64 * h[i - 1];
            m[i] = (w1 + w2) / (w1 / d0 + w2 / d1);
        } else {
            // Enforce monotonicity: zero slope when adjacent secants change sign.
            m[i] = 0.0_f64;
        }
    }

    // 2) Endpoints: one-sided estimates with limiting to preserve shape.
    // Left endpoint.
    let m0 = ((2.0_f64 * h[0] + h[1]) * d[0] - h[0] * d[1]) / (h[0] + h[1]);
    m[0] = if m0.signum() != d[0].signum() {
        0.0_f64
    } else if d[0].signum() != d[1].signum() && m0.abs() > 3.0_f64 * d[0].abs() {
        3.0_f64 * d[0]
    } else {
        m0
    };

    // Right endpoint.
    let mn = ((2.0_f64 * h[n - 2] + h[n - 3]) * d[n - 2] - h[n - 2] * d[n - 3])
        / (h[n - 2] + h[n - 3]);
    m[n - 1] = if mn.signum() != d[n - 2].signum() {
        0.0_f64
    } else if d[n - 2].signum() != d[n - 3].signum() && mn.abs() > 3.0_f64 * d[n - 2].abs() {
        3.0_f64 * d[n - 2]
    } else {
        mn
    };

    m
}

#[derive(Debug, Clone)]
struct Pchip {
    x: Vec<f64>,
    y: Vec<f64>,
    m: Vec<f64>,
}

impl Pchip {
    /// Construct a PCHIP interpolator from x/y samples.
    fn new(x: Vec<f64>, y: Vec<f64>) -> Self {
        let m = compute_pchip_slopes(&x, &y);
        Self { x, y, m }
    }

    /// Evaluate the interpolant at `t`.
    fn eval(&self, t: f64) -> f64 {
        // Locate the interval [x_i, x_{i+1}] containing t.
        let idx = match self.x.binary_search_by(|v| v.partial_cmp(&t).unwrap()) {
            Ok(i) => return self.y[i],
            Err(i) => {
                if i == 0 {
                    0
                } else if i >= self.x.len() {
                    self.x.len() - 1
                } else {
                    i - 1
                }
            }
        };

        let i = if idx >= self.x.len() - 1 {
            self.x.len() - 2
        } else {
            idx
        };

        let h = self.x[i + 1] - self.x[i];
        let d = (t - self.x[i]) / h;

        // Cubic Hermite basis functions.
        let h00 = 2.0 * d.powi(3) - 3.0 * d.powi(2) + 1.0;
        let h10 = d.powi(3) - 2.0 * d.powi(2) + d;
        let h01 = -2.0 * d.powi(3) + 3.0 * d.powi(2);
        let h11 = d.powi(3) - d.powi(2);

        // Hermite interpolation using endpoint values and slopes.
        h00 * self.y[i] + h10 * h * self.m[i] + h01 * self.y[i + 1] + h11 * h * self.m[i + 1]
    }
}

// -----------------------------------------------------------------------------
// Table implementation
// -----------------------------------------------------------------------------

#[derive(Clone)]
pub struct NeutrinoTable {
    interp: Pchip,
    omega_nu0: f64,
}

impl NeutrinoTable {
    /// Build a tabulated neutrino background ω_ν(a) with monotone cubic interpolation.
    ///
    /// The table is generated on a log-spaced scale factor grid a ∈ [a_min, 1].
    /// Massive species are integrated numerically using the accuracy preset controls.
    pub fn new(spec: &NeutrinoSpec, derived: &DerivedParams, accuracy: AccuracyPreset) -> Self {
        let n = accuracy.neutrino_n_points();
        let y_max = accuracy.neutrino_y_max();
        let steps = accuracy.neutrino_integ_steps();

        // Grid in ln(a) for better coverage across early/late times.
        let a_min = 1e-5_f64;
        let a_max = 1.0_f64;
        let ln_a_min = a_min.ln();
        let step = (a_max.ln() - ln_a_min) / (n as f64 - 1.0);

        let mut ln_a_grid = Vec::with_capacity(n);
        let mut omega_grid = Vec::with_capacity(n);

        for i in 0..n {
            let x = ln_a_min + i as f64 * step;
            let a = x.exp();
            let z = 1.0 / a - 1.0;

            // Compute ω_ν(z) based on the chosen neutrino specification.
            let val = match spec {
                NeutrinoSpec::Effective { masses_ev, .. } => {
                    if masses_ev.is_empty() {
                        0.0
                    } else {
                        calc_omega_nu_z_cfg(masses_ev, derived.t_nu0_k, z, y_max, steps, 2.0)
                    }
                }
                NeutrinoSpec::Split {
                    masses_ev,
                    temp_factors,
                    ..
                } => {
                    // In split mode, each massive species can have its own temperature factor.
                    let mut total = 0.0;
                    for (&m, &tf) in masses_ev.iter().zip(temp_factors.iter()) {
                        // Use `t_nu0_base_k` for the split-model base temperature.
                        total += calc_omega_nu_z_cfg(&[m], derived.t_nu0_base_k * tf, z, y_max, steps, 2.0);
                    }
                    total
                }
            };

            ln_a_grid.push(x);
            omega_grid.push(val);
        }

        Self {
            interp: Pchip::new(ln_a_grid, omega_grid),
            omega_nu0: derived.omega_nu0,
        }
    }
}

impl NeutrinoBackground for NeutrinoTable {
    fn omega_nu_z(&self, z: f64) -> f64 {
        debug_assert!(z >= 0.0, "NeutrinoTable does not support z < 0");

        // Use ω_ν0 exactly at (or below) z=0 to avoid log/roundoff issues.
        if z <= 0.0 {
            return self.omega_nu0;
        }

        // Interpolation variable is ln(a).
        let a = 1.0 / (1.0 + z);
        let x = a.ln();

        // If we go earlier than the table, extrapolate with pure radiation scaling (a^-4).
        let x0 = self.interp.x[0];
        if x < x0 {
            return self.interp.y[0] * (-4.0 * (x - x0)).exp();
        }

        // If we go later than the table end, return the cached present-day value.
        if x >= *self.interp.x.last().unwrap() {
            return self.omega_nu0;
        }

        self.interp.eval(x)
    }

    fn omega_nu0(&self) -> f64 {
        self.omega_nu0
    }
}

// -----------------------------------------------------------------------------
// Analytic implementation (massless neutrinos)
// -----------------------------------------------------------------------------

/// Analytic backend for strictly massless neutrinos.
///
/// For a truly massless species, the physical density scales as:
///   ω_ν(z) = ω_ν0 (1+z)^4
#[derive(Clone)]
pub struct NeutrinoAnalytic {
    omega_nu0: f64,
}

impl NeutrinoAnalytic {
    /// Construct the analytic backend given ω_ν0.
    pub fn new(omega_nu0: f64) -> Self {
        Self { omega_nu0 }
    }
}

impl NeutrinoBackground for NeutrinoAnalytic {
    fn omega_nu_z(&self, z: f64) -> f64 {
        debug_assert!(z >= 0.0, "NeutrinoAnalytic: z should be >= 0");

        // Return ω_ν0 exactly at (or below) z=0.
        if z <= 0.0 {
            return self.omega_nu0;
        }

        self.omega_nu0 * (1.0 + z).powi(4)
    }

    fn omega_nu0(&self) -> f64 {
        self.omega_nu0
    }
}

// -----------------------------------------------------------------------------
// Static dispatch wrapper
// -----------------------------------------------------------------------------

/// High-performance neutrino backend selection.
///
/// This avoids dynamic dispatch in hot paths while still allowing multiple
/// backend implementations.
#[derive(Clone)]
pub enum NeutrinoBackend {
    Table(NeutrinoTable),
    Analytic(NeutrinoAnalytic),
}

impl NeutrinoBackground for NeutrinoBackend {
    fn omega_nu_z(&self, z: f64) -> f64 {
        match self {
            Self::Table(t) => t.omega_nu_z(z),
            Self::Analytic(a) => a.omega_nu_z(z),
        }
    }

    fn omega_nu0(&self) -> f64 {
        match self {
            Self::Table(t) => t.omega_nu0(),
            Self::Analytic(a) => a.omega_nu0(),
        }
    }
}