SpectraRust/src/tlusty/math/atomic/dietot.rs

//! 计算双电子复合对光电离截面的修正。
//!
//! 重构自 TLUSTY `dietot.f`
//! 遍历所有离子和深度点，计算双电子复合速率和伪截面。

use crate::tlusty::math::dielrc;
use crate::tlusty::state::atomic::{AtoPar, IonPar, LevPar};
use crate::tlusty::state::config::{BasNum, InpPar};
use crate::tlusty::state::model::{LevAdd, ModPar};

/// DIETOT 参数结构体
pub struct DietotParams<'a> {
    /// 基本数值
    pub basnum: &'a BasNum,
    /// 输入参数（包含 wmm, ytot）
    pub inppar: &'a InpPar,
    /// 原子参数
    pub atopar: &'a AtoPar,
    /// 离子参数
    pub ionpar: &'a IonPar,
    /// 能级参数
    pub levpar: &'a LevPar,
    /// 模型参数（包含 temp, dens）
    pub modpar: &'a ModPar,
    /// 能级附加数据（包含 diesig 输出）
    pub levadd: &'a mut LevAdd,
}

/// 计算双电子复合截面修正。
///
/// 遍历所有离子和深度点，调用 `dielrc` 计算双电子复合速率和伪截面。
/// 结果存储在 `levadd.diesig[ion][id]` 中。
///
/// # 参数
/// * `params` - 参数结构体
pub fn dietot(params: &mut DietotParams) {
    let nion = params.basnum.nion as usize;
    let nd = params.basnum.nd as usize;

    for ion in 0..nion {
        // 获取离子的第一个能级索引
        let i = (params.ionpar.nfirst[ion] - 1) as usize; // 转换为 0-indexed

        // 获取原子序数和离子电荷
        let iatm_idx = (params.levpar.iatm[i] - 1) as usize;
        let ia = params.atopar.numat[iatm_idx] as usize;
        let io = params.ionpar.iz[ion] as usize;

        for id in 0..nd {
            let t = params.modpar.temp[id];
            let xpx = params.modpar.dens[id] / params.inppar.wmm[id] / params.inppar.ytot[id];

            // 调用 dielrc 计算双电子复合速率和伪截面
            let (dirt, sig0) = dielrc(ia, io, t, xpx);

            // 存储结果
            params.levadd.diesig[ion][id] = sig0;

            if id == 0 || id == 34 || id == nd - 1 {
                eprintln!("{:5}{:5}{:5}{:5}{:5}{:5}{:12.4}{:12.4}",
                    ion + 1, ia, io, id + 1, i + 1,
                    params.ionpar.nnext[ion], dirt, sig0);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tlusty::state::atomic::{AtoPar, IonPar, LevPar};
    use crate::tlusty::state::config::{BasNum, InpPar};
    use crate::tlusty::state::constants::{MDEPTH, MION, MLEVEL, MATOM};
    use crate::tlusty::state::model::{LevAdd, ModPar};

    fn create_test_state() -> (
        BasNum,
        InpPar,
        AtoPar,
        IonPar,
        LevPar,
        ModPar,
        LevAdd,
    ) {
        let mut basnum = BasNum::default();
        basnum.nion = 3;
        basnum.nd = 5;

        let mut inppar = InpPar::new();
        for i in 0..5 {
            inppar.wmm[i] = 1.0 + i as f64 * 0.1;
            inppar.ytot[i] = 1.0;
        }

        let mut atopar = AtoPar::default();
        // 设置原子序数：H=1, He=2
        atopar.numat[0] = 1; // H
        atopar.numat[1] = 2; // He

        let mut ionpar = IonPar::default();
        // 设置离子 0: H I (nfirst=1, iz=1)
        ionpar.nfirst[0] = 1;
        ionpar.iz[0] = 1;
        // 设置离子 1: He I (nfirst=2, iz=1)
        ionpar.nfirst[1] = 2;
        ionpar.iz[1] = 1;
        // 设置离子 2: He II (nfirst=3, iz=2)
        ionpar.nfirst[2] = 3;
        ionpar.iz[2] = 2;

        let mut levpar = LevPar::default();
        // 设置能级的原子归属
        levpar.iatm[0] = 1; // 能级 1 属于原子 1 (H)
        levpar.iatm[1] = 2; // 能级 2 属于原子 2 (He)
        levpar.iatm[2] = 2; // 能级 3 属于原子 2 (He)

        let mut modpar = ModPar::default();
        for i in 0..5 {
            modpar.temp[i] = 10000.0 + i as f64 * 1000.0;
            modpar.dens[i] = 1e-10 + i as f64 * 1e-11;
        }

        let levadd = LevAdd::new();

        (basnum, inppar, atopar, ionpar, levpar, modpar, levadd)
    }

    #[test]
    fn test_dietot_basic() {
        let (basnum, inppar, atopar, ionpar, levpar, modpar, mut levadd) = create_test_state();

        let mut params = DietotParams {
            basnum: &basnum,
            inppar: &inppar,
            atopar: &atopar,
            ionpar: &ionpar,
            levpar: &levpar,
            modpar: &modpar,
            levadd: &mut levadd,
        };

        dietot(&mut params);

        // 验证结果已存储
        for ion in 0..3 {
            for id in 0..5 {
                // diesig 应该有一些值（可能是 0.0 如果没有数据）
                assert!(
                    params.levadd.diesig[ion][id].is_finite(),
                    "diesig[{}][{}] should be finite",
                    ion,
                    id
                );
            }
        }
    }

    #[test]
    fn test_dietot_values() {
        let (basnum, inppar, atopar, ionpar, levpar, modpar, mut levadd) = create_test_state();

        let mut params = DietotParams {
            basnum: &basnum,
            inppar: &inppar,
            atopar: &atopar,
            ionpar: &ionpar,
            levpar: &levpar,
            modpar: &modpar,
            levadd: &mut levadd,
        };

        dietot(&mut params);

        // 验证某些已知离子的值
        // 对于 H I (离子 0，ia=1, io=1)，dielrc 可能返回 0（没有双电子复合数据）
        // 对于 He I (离子 1，ia=2, io=1)，可能有数据
        // 对于 He II (离子 2，ia=2, io=2)，可能有数据

        // 验证 H I 的第一个深度点
        let diesig_h1 = params.levadd.diesig[0][0];
        // H I 的双电子复合通常为 0（氢没有双电子复合）
        assert!(
            diesig_h1 >= 0.0,
            "H I diesig should be non-negative, got {}",
            diesig_h1
        );
    }
}