SpectraRust/src/synspec/math/inkur.rs

//! Input of a Kurucz model atmosphere.
//!
//! Translated from SYNSPEC54.FOR subroutine INKUR at line 11048.
//!
//! Reads a Kurucz model atmosphere from file (unit 8) and initializes
//! the model state arrays (DM, TEMP, ELEC, DENS, POPUL, etc.).

/// Parameters for INKUR initialization.
pub struct InkurParams<'a> {
    /// Boltzmann constant (BOLK)
    pub bolk: f64,
    /// Mean molecular weight at each depth
    pub wmm: &'a [f64],
    /// Total number of particles per H atom
    pub ytot: &'a [f64],
    /// Number of atoms
    pub natom: usize,
    /// Number of levels
    pub nlevel: usize,
    /// Molecular equilibrium flag
    pub ifmol: i32,
    /// Molecular temperature limit
    pub tmolim: f64,
    /// Maximum number of depth points
    pub nd_max: usize,
}

/// Result of INKUR initialization.
pub struct InkurResult {
    /// Number of depth points
    pub nd: usize,
    /// Effective temperature (from file header)
    pub tef: f64,
    /// Surface gravity log g (from file header)
    pub grav: f64,
    /// Mass depth coordinate
    pub dm: Vec<f64>,
    /// Temperature at each depth
    pub temp: Vec<f64>,
    /// Electron density at each depth
    pub elec: Vec<f64>,
    /// Mass density at each depth
    pub dens: Vec<f64>,
    /// Population of each level at each depth (nlevel x nd)
    pub popul: Vec<Vec<f64>>,
}

/// Input of a Kurucz model atmosphere.
///
/// Reads model atmosphere data and initializes the depth-dependent arrays.
/// For each depth point, computes density from pressure and temperature,
/// optionally solves molecular equilibrium, and computes LTE populations.
///
/// # Arguments
/// * `params` - Initialization parameters
/// * `tef` - Effective temperature from file header
/// * `grav` - Surface gravity from file header
/// * `depth_data` - Slice of (dm, temp, pressure, elec) for each depth
/// * `moleq_fn` - Optional molecular equilibrium callback: (id, t, an, aein) -> ane
/// * `attot_fn` - Callback to compute ATTOT: (iat, id, dens, wmm, ytot, abund) -> f64
/// * `post_depth_fn` - Callback after each depth: (id) for WNSTOR, SABOLF, RATMAT, LEVSOL
///
/// # Returns
/// Initialized model arrays.
pub fn inkur(
    params: &InkurParams,
    tef: f64,
    grav: f64,
    depth_data: &[(f64, f64, f64, f64)],
    moleq_fn: Option<&dyn Fn(usize, f64, f64, f64) -> f64>,
    attot_fn: &dyn Fn(usize, usize, f64, f64, f64, f64) -> f64,
    post_depth_fn: &dyn Fn(usize, &mut [Vec<f64>]),
) -> InkurResult {
    let bolk = params.bolk;
    let nd = depth_data.len().min(params.nd_max);
    let mut dm = Vec::with_capacity(nd);
    let mut temp = Vec::with_capacity(nd);
    let mut elec = Vec::with_capacity(nd);
    let mut dens = Vec::with_capacity(nd);
    let mut popul = vec![vec![0.0; nd]; params.nlevel];

    for (id, &(dm_i, temp_i, p, elec_i)) in depth_data.iter().enumerate().take(nd) {
        dm.push(dm_i);
        temp.push(temp_i);
        elec.push(elec_i);

        // Compute density: DENS = WMM * (P/(T*BOLK) - ELEC)
        let an = p / temp_i / bolk;
        let dens_i = params.wmm[id] * (an - elec_i);
        dens.push(dens_i);

        let t = temp_i;

        // Molecular equilibrium or simple abundance
        if params.ifmol > 0 && t < params.tmolim {
            if let Some(ref moleq) = moleq_fn {
                let aein = elec_i;
                let _ane = moleq(id, t, an, aein);
            }
        } else {
            // Compute total atom abundance for each atom
            // Fortran: ATTOT(IAT,ID)=DENS(ID)/WMM(ID)/YTOT(ID)*ABUND(IAT,ID)
            for iat in 0..params.natom {
                let _ = attot_fn(iat, id, dens_i, params.wmm[id], params.ytot[id], 0.0);
            }
        }

        // Post-depth processing: WNSTOR, SABOLF, RATMAT, LEVSOL
        post_depth_fn(id, &mut popul);
    }

    InkurResult {
        nd,
        tef,
        grav,
        dm,
        temp,
        elec,
        dens,
        popul,
    }
}

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

    #[test]
    fn test_inkur_basic() {
        let params = InkurParams {
            bolk: 1.380649e-16,
            wmm: &[1.0, 1.0],
            ytot: &[1.0, 1.0],
            natom: 1,
            nlevel: 2,
            ifmol: 0,
            tmolim: 0.0,
            nd_max: 10,
        };

        let depth_data = [
            (1e-4, 5000.0, 1e5, 1e11),
            (1e-3, 4000.0, 1e4, 1e10),
        ];

        let attot_fn = |_iat: usize, _id: usize, dens: f64, wmm: f64, ytot: f64, _abund: f64| {
            dens / wmm / ytot
        };

        let post_depth_fn = |_id: usize, _popul: &mut [Vec<f64>]| {
            // No-op for test
        };

        let result = inkur(
            &params,
            5777.0,
            4.44,
            &depth_data,
            None,
            &attot_fn,
            &post_depth_fn,
        );

        assert_eq!(result.nd, 2);
        assert_eq!(result.dm.len(), 2);
        assert_eq!(result.temp.len(), 2);
        assert_eq!(result.elec.len(), 2);
        assert_eq!(result.dens.len(), 2);
        assert_eq!(result.tef, 5777.0);
        assert_eq!(result.grav, 4.44);
        assert!(result.dens[0].is_finite());
        assert!(result.dens[1].is_finite());
        // DENS = WMM * (P/(T*BOLK) - ELEC)
        let an0 = 1e5 / 5000.0 / 1.380649e-16;
        let expected_dens0 = 1.0 * (an0 - 1e11);
        assert!((result.dens[0] - expected_dens0).abs() < 1.0);
    }
}