SpectraRust/src/synspec/math/inpmod.rs

//! inpmod — 读取初始模型大气。
//!
//! Fortran 原始签名: SUBROUTINE INPMOD
//!
//! 从 unit 8 读取 TLUSTY 模型大气，计算 LTE 能级布居，
//! 并可选地替换为 NLTE 布居。
//!
//! 注意: Fortran 版本直接操作 COMMON 块和文件 I/O。
//! Rust 版本提供纯计算核心函数。

/// 模型大气深度点数据
#[derive(Debug, Clone)]
pub struct ModelDepthPoint {
    /// 质量深度 (g/cm^2)
    pub depth: f64,
    /// 温度 (K)
    pub temp: f64,
    /// 电子密度 (cm^-3)
    pub elec: f64,
    /// 质量密度 (g/cm^3)
    pub dens: f64,
}

/// 计算总粒子数密度
///
/// TOTN = DENS / WMM + ELEC
///
/// Fortran 原始逻辑:
/// ```fortran
/// TOTN(ID)=DENS(ID)/WMM(ID)+ELEC(ID)
/// ```
pub fn total_number_density(dens: f64, wmm: f64, elec: f64) -> f64 {
    dens / wmm + elec
}

/// 计算束缚-自由常数 (BCON)
///
/// BCON = ELEC / TEMP / SQRT(TEMP) * 2.0706E-16
///
/// Fortran 原始逻辑:
/// ```fortran
/// BCON=ELEC(ID)/TEMP(ID)/SQRT(TEMP(ID))*2.0706E-16
/// ```
pub fn bound_free_constant(elec: f64, temp: f64) -> f64 {
    elec / temp / temp.sqrt() * 2.0706e-16
}

/// 计算原子总数密度
///
/// ATTOT(IAT,ID) = DENS / WMM / YTOT * ABUND(IAT,ID)
///
/// Fortran 原始逻辑:
/// ```fortran
/// DO IAT=1,NATOM
///    ATTOT(IAT,ID)=DENS(ID)/WMM(ID)/YTOT(ID)*ABUND(IAT,ID)
/// END DO
/// ```
pub fn compute_atom_densities(
    dens: f64,
    wmm: f64,
    ytot: f64,
    abundances: &[f64],
) -> Vec<f64> {
    let factor = dens / wmm / ytot;
    abundances.iter().map(|&a| factor * a).collect()
}

/// NLTE 布居替换
///
/// POPUL(J,ID) = X(IP+I) * RELAB(IATM(I),ID)
///
/// Fortran 原始逻辑:
/// ```fortran
/// DO I=1,NLEV0
///    j=iltot(i)
///    POPUL(J,ID)=X(IP+I)*RELAB(IATM(I),ID)
/// END DO
/// ```
pub fn replace_nlte_populations(
    populations: &mut [f64],
    nlte_data: &[f64],
    iltot: &[usize],
    relab: &[f64],
    iatm: &[usize],
) {
    for (i, &j) in iltot.iter().enumerate().take(nlte_data.len()) {
        if j > 0 && j <= populations.len() {
            let relab_val = if iatm[i] < relab.len() { relab[iatm[i]] } else { 1.0 };
            populations[j - 1] = nlte_data[i] * relab_val;
        }
    }
}

/// B 因子修正
///
/// POPUL(J,ID) = POPUL(J,ID) * PLTE(J,ID)
///
/// Fortran 原始逻辑:
/// ```fortran
/// if(ibfac.eq.1) then
///    do i=1,nlev0
///       j=iltot(i)
///       popul(j,id)=popul(j,id)*plte(j,id)
///    end do
/// end if
/// ```
pub fn apply_bfactor_correction(
    populations: &mut [f64],
    lte_populations: &[f64],
    iltot: &[usize],
) {
    for &j in iltot {
        if j > 0 && j <= populations.len() {
            populations[j - 1] *= lte_populations[j - 1];
        }
    }
}

/// 模型大气数据
#[derive(Debug, Clone)]
pub struct ModelAtmosphere {
    /// 深度点数据
    pub depths: Vec<ModelDepthPoint>,
    /// 能级布居 [level][depth]
    pub populations: Vec<Vec<f64>>,
    /// LTE 布居 [level][depth]
    pub lte_populations: Vec<Vec<f64>>,
}

/// 模型大气统计信息
#[derive(Debug, Clone)]
pub struct ModelStats {
    /// 最低温度
    pub min_temp: f64,
    /// 最高温度
    pub max_temp: f64,
    /// 最低密度
    pub min_dens: f64,
    /// 最高密度
    pub max_dens: f64,
    /// 深度点数
    pub n_depths: usize,
    /// 能级数
    pub n_levels: usize,
}

/// 计算模型大气统计信息
pub fn compute_model_stats(model: &ModelAtmosphere) -> ModelStats {
    let mut min_temp = f64::MAX;
    let mut max_temp = f64::MIN;
    let mut min_dens = f64::MAX;
    let mut max_dens = f64::MIN;

    for depth in &model.depths {
        if depth.temp < min_temp { min_temp = depth.temp; }
        if depth.temp > max_temp { max_temp = depth.temp; }
        if depth.dens < min_dens { min_dens = depth.dens; }
        if depth.dens > max_dens { max_dens = depth.dens; }
    }

    ModelStats {
        min_temp,
        max_temp,
        min_dens,
        max_dens,
        n_depths: model.depths.len(),
        n_levels: model.populations.len(),
    }
}

/// 模型文件格式
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ModelFormat {
    /// LTE 模型 (3 参数: T, NE, RHO)
    Lte,
    /// NLTE 模型 (>3 参数)
    Nlte,
    /// 磁盘模型 (4 参数)
    Disk,
}

/// 检查模型格式
pub fn detect_model_format(numpar: i32, inmod: i32) -> ModelFormat {
    if inmod == 2 {
        ModelFormat::Disk
    } else if numpar.abs() > 3 {
        ModelFormat::Nlte
    } else {
        ModelFormat::Lte
    }
}

/// 计算 NLTE 修正因子
///
/// PNLT(IAT,ION,ID) = POPUL(NKI,ID) / G(NKI) * BCON
///
/// Fortran 原始逻辑:
/// ```fortran
/// BCON=ELEC(ID)/TEMP(ID)/SQRT(TEMP(ID))*2.0706E-16
/// IF(ION.GT.0) PNLT(IAT,ION,ID)=POPUL(NKI,ID)/G(NKI)*BCON
/// ```
pub fn compute_nlte_correction(
    popul_nki: f64,
    g_nki: f64,
    elec: f64,
    temp: f64,
) -> f64 {
    if g_nki > 0.0 {
        popul_nki / g_nki * bound_free_constant(elec, temp)
    } else {
        0.0
    }
}

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

    #[test]
    fn test_total_number_density() {
        let totn = total_number_density(1e-8, 1.0, 1e-10);
        // 1e-8/1.0 + 1e-10 = 1.01e-8
        assert!((totn - 1.01e-8).abs() < 1e-15);
    }

    #[test]
    fn test_bound_free_constant() {
        let bcon = bound_free_constant(1e10, 10000.0);
        // 1e10 / 10000 / 100 * 2.0706e-16 = 1e10 / 1e6 * 2.0706e-16 = 2.0706e-12
        let expected = 1e10 / 10000.0 / 100.0 * 2.0706e-16;
        assert!((bcon - expected).abs() / expected < 1e-10);
    }

    #[test]
    fn test_compute_atom_densities() {
        let abundances = vec![0.9, 0.1];
        let attot = compute_atom_densities(1e-8, 1.0, 1.5, &abundances);
        assert_eq!(attot.len(), 2);
        let expected0 = 1e-8 / 1.0 / 1.5 * 0.9;
        assert!((attot[0] - expected0).abs() < 1e-20);
    }

    #[test]
    fn test_replace_nlte_populations() {
        let mut popul = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let nlte_data = vec![10.0, 20.0];
        let iltot = vec![1, 3]; // 1-indexed
        let relab = vec![1.0, 1.0, 1.0];
        let iatm = vec![0, 1];
        replace_nlte_populations(&mut popul, &nlte_data, &iltot, &relab, &iatm);
        assert_eq!(popul[0], 10.0); // j=1 → index 0
        assert_eq!(popul[2], 20.0); // j=3 → index 2
        assert_eq!(popul[4], 5.0);  // unchanged
    }

    #[test]
    fn test_apply_bfactor_correction() {
        let mut popul = vec![2.0, 3.0, 4.0];
        let lte = vec![0.5, 0.5, 0.5];
        let iltot = vec![1, 2, 3]; // 1-indexed (Fortran convention)
        apply_bfactor_correction(&mut popul, &lte, &iltot);
        assert_eq!(popul[0], 1.0);  // 2.0 * 0.5
        assert_eq!(popul[1], 1.5);  // 3.0 * 0.5
        assert_eq!(popul[2], 2.0);  // 4.0 * 0.5
    }

    #[test]
    fn test_detect_model_format() {
        assert_eq!(detect_model_format(3, 0), ModelFormat::Lte);
        assert_eq!(detect_model_format(5, 0), ModelFormat::Nlte);
        assert_eq!(detect_model_format(3, 2), ModelFormat::Disk);
    }

    #[test]
    fn test_compute_nlte_correction() {
        let pnlt = compute_nlte_correction(1e12, 2.0, 1e10, 10000.0);
        // 1e12 / 2.0 * 1e10/10000/100 * 2.0706e-16
        let bcon = bound_free_constant(1e10, 10000.0);
        let expected = 1e12 / 2.0 * bcon;
        assert!((pnlt - expected).abs() / expected < 1e-10);
    }

    #[test]
    fn test_compute_model_stats() {
        let model = ModelAtmosphere {
            depths: vec![
                ModelDepthPoint { depth: 1.0, temp: 5000.0, elec: 1e10, dens: 1e-8 },
                ModelDepthPoint { depth: 10.0, temp: 15000.0, elec: 1e12, dens: 1e-6 },
            ],
            populations: vec![vec![1.0, 2.0]],
            lte_populations: vec![vec![1.0, 2.0]],
        };
        let stats = compute_model_stats(&model);
        assert_eq!(stats.min_temp, 5000.0);
        assert_eq!(stats.max_temp, 15000.0);
        assert_eq!(stats.n_depths, 2);
    }
}