SpectraRust/src/math/lymlin.rs

//! Lyman 线系不透明度和发射率计算。
//!
//! 重构自 TLUSTY `lymlin.f`
//!
//! 计算前 30 条 Lyman 线的不透明度和发射率贡献。

use crate::state::constants::{MDEPTH, TWO};

// ============================================================================
// 常量参数
// ============================================================================

const SIXTH: f64 = 1.0 / 6.0;
const TTW: f64 = 2.0 / 3.0;
const OS0: f64 = 0.02654;
const CPP: f64 = 4.1412e-16;
const CPJ: f64 = 157803.0;
const C00: f64 = 1.25e-9;
const C18: f64 = 2.997925e18;

/// Lyman 线数量
const MLEVL: usize = 30;

// ============================================================================
// Lymlin 缓存结构
// ============================================================================

/// Lyman 线计算的缓存数据。
///
/// 第一次调用时计算并存储，后续调用直接使用。
#[derive(Debug, Clone)]
pub struct LymlinCache {
    /// 是否已初始化
    pub initialized: bool,
    /// Stark K 系数 (j=2..30)
    pub xkijl: Vec<f64>,
    /// Stark f 值 (j=2..30)
    pub fijl: Vec<f64>,
    /// 谱线波长 (j=2..30)
    pub wl0l: Vec<f64>,
    /// 谱线频率 (j=2..30)
    pub fr0l: Vec<f64>,
    /// 能级占据数 (j, id)
    pub pj: Vec<Vec<f64>>,
    /// 吸收系数 (j, id)
    pub abtr: Vec<Vec<f64>>,
    /// 发射系数 (j, id)
    pub emtr: Vec<Vec<f64>>,
    /// FID 系数 (j, id)
    pub fid: Vec<Vec<f64>>,
    /// ADH 系数 (j, id)
    pub ad0: Vec<Vec<f64>>,
    /// DIVH 系数 (j, id)
    pub div0: Vec<Vec<f64>>,
    /// DBETA 系数 (j, id)
    pub dbet0: Vec<Vec<f64>>,
    /// BETAD 系数 (j, id)
    pub betad0: Vec<Vec<f64>>,
}

impl Default for LymlinCache {
    fn default() -> Self {
        Self {
            initialized: false,
            xkijl: vec![0.0; MLEVL + 1],      // 索引 2..30
            fijl: vec![0.0; MLEVL + 1],
            wl0l: vec![0.0; MLEVL + 1],
            fr0l: vec![0.0; MLEVL + 1],
            pj: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            abtr: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            emtr: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            fid: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            ad0: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            div0: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            dbet0: vec![vec![0.0; MDEPTH]; MLEVL + 1],
            betad0: vec![vec![0.0; MDEPTH]; MLEVL + 1],
        }
    }
}

// ============================================================================
// 参数结构体
// ============================================================================

/// Lymlin 函数的输入参数。
pub struct LymlinParams<'a> {
    /// 深度点数
    pub nd: usize,
    /// 氢元素索引 (0-based, Fortran 中是 1-based)
    pub ielh: i32,
    /// 氢原子能级起始索引 (0-based)
    pub n0h: i32,
    /// 氢原子能级终止索引 (0-based)
    pub n1h: i32,
    /// 氢离子能级索引 (0-based)
    pub nkh: i32,
    /// 氢原子能级数
    pub nlh: i32,
    /// 温度数组 (nd)
    pub temp: &'a [f64],
    /// 电子密度数组 (nd)
    pub elec: &'a [f64],
    /// 占据数数组 (nlevel, nd)
    pub popul: &'a [Vec<f64>],
    /// 湍流速度数组 (nd)
    pub vturb: &'a [f64],
    /// 氢能级占据概率 (nlmx, nd)
    pub wnhint: &'a [Vec<f64>],
    /// 频率数组
    pub freq: &'a [f64],
    /// 当前频率索引 (0-based)
    pub ij: usize,
    /// 不透明度数组 (nd) - 输入/输出
    pub abso1: &'a mut [f64],
    /// 发射率数组 (nd) - 输入/输出
    pub emis1: &'a mut [f64],
}

// ============================================================================
// 主函数
// ============================================================================

/// 计算 Lyman 线系的不透明度和发射率。
///
/// # 参数
///
/// * `params` - 输入参数结构体
/// * `cache` - 缓存数据（第一次调用时初始化）
///
/// # 说明
///
/// 计算前 30 条 Lyman 线 (n=1 到 n=2..30) 对不透明度和发射率的贡献。
/// 使用 Stark 展宽理论处理谱线轮廓。
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE LYMLIN(IJ)
///   ...
///   FR=FREQ(IJ)
///   IF(FR.GT.3.28805E15.OR.FR.LT.1.5E15) RETURN
///   ...
/// ```
pub fn lymlin(params: &mut LymlinParams, cache: &mut LymlinCache) {
    // 如果没有氢，直接返回
    if params.ielh == 0 {
        return;
    }

    let ij = params.ij;
    let fr = params.freq[ij];

    // 频率范围检查：1.5e15 到 3.28805e15 Hz
    if fr > 3.28805e15 || fr < 1.5e15 {
        return;
    }

    let nd = params.nd;
    let n0h = params.n0h;
    let n1h = params.n1h;
    let nkh = params.nkh;
    let nlh = params.nlh;

    // 第一次调用时初始化缓存
    if !cache.initialized {
        initialize_cache(params, cache);
        cache.initialized = true;
    }

    // 主计算循环
    let wl = C18 / fr;
    let f15 = fr * 1.0e-15;

    // 工作数组
    let mut ablym = vec![0.0; MDEPTH];
    let mut emlym = vec![0.0; MDEPTH];

    for id in 0..nd {
        ablym[id] = 0.0;
        emlym[id] = 0.0;

        for j in 2..=30 {
            let beta = cache.dbet0[j][id] * (fr - cache.fr0l[j]).abs();
            let betad = cache.betad0[j][id];
            let adh = cache.ad0[j][id];
            let divh = cache.div0[j][id];

            let sg = super::starka::starka(beta, TWO, adh, betad, divh) * cache.fid[j][id];
            ablym[id] += sg * cache.abtr[j][id];
            emlym[id] += sg * cache.emtr[j][id];
        }

        let xkt = (-4.79928e-11 * fr / params.temp[id]).exp();
        let xkb = xkt * 1.4743e-2 * f15 * f15 * f15;

        ablym[id] -= xkt * emlym[id];
        emlym[id] = xkb * emlym[id];

        // 累加到输出数组
        params.abso1[id] += ablym[id];
        params.emis1[id] += emlym[id];
    }
}

/// 初始化缓存数据。
fn initialize_cache(params: &LymlinParams, cache: &mut LymlinCache) {
    let nd = params.nd;
    let n0h = params.n0h;
    let n1h = params.n1h;
    let nkh = params.nkh;
    let nlh = params.nlh;

    // 计算 j=2..30 的 Stark 参数
    for j in 2..=30 {
        let (xkij0, wl00, fij0) = super::stark0::stark0(1, j, 1);
        cache.xkijl[j] = xkij0;
        cache.fijl[j] = fij0;
        cache.wl0l[j] = wl00;
        cache.fr0l[j] = C18 / wl00;
    }

    // 对每个深度点计算
    for id in 0..nd {
        let t = params.temp[id];
        let t1 = 1.0 / t;
        let sqt = t.sqrt();
        let ane = params.elec[id];
        let anp = params.popul[nkh as usize][id];

        let f00 = C00 * ane.powf(TTW);
        let dop0 = 1.0e8 * (1.65e8 * t + params.vturb[id]).sqrt();

        let p0 = CPP * ane * anp * t1 / sqt;
        let p1 = params.popul[n0h as usize][id];

        for j in 2..=30 {
            let x = (j * j) as f64;
            let jj = n0h + (j as i32) - 1;
            let xjj = 1.0 / (jj as f64);

            // 计算 PJ
            if j as i32 <= nlh {
                cache.pj[j][id] = params.popul[jj as usize][id];
            } else {
                // 使用 LTE 近似
                let wnhint_j = if (j as usize) < params.wnhint.len() {
                    params.wnhint[j][id]
                } else {
                    1.0
                };
                cache.pj[j][id] = p0 * (CPJ / (x * t)).exp() * x * wnhint_j;
            }

            // 计算吸收和发射系数
            cache.abtr[j][id] = p1 * params.wnhint[j][id];
            cache.emtr[j][id] = cache.pj[j][id] * xjj * (CPJ * (1.0 - xjj) * t1).exp();

            // 计算其他参数
            let fxk = f00 * cache.xkijl[j];
            let dbeta = cache.wl0l[j] * cache.wl0l[j] / (C18 * fxk);
            cache.fid[j][id] = OS0 * cache.fijl[j] * dbeta;

            let dop = dop0 / cache.wl0l[j];
            let betad = dop * dbeta;

            // 调用 divstr
            let (adh, divh) = super::divstr::divstr(betad, 1);

            cache.ad0[j][id] = adh;
            cache.div0[j][id] = divh;
            cache.dbet0[j][id] = dbeta;
            cache.betad0[j][id] = betad;
        }
    }
}

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

    fn create_test_params<'a>(
        nd: usize,
        abso1: &'a mut [f64],
        emis1: &'a mut [f64],
        freq: &'a [f64],
        temp: &'a [f64],
        elec: &'a [f64],
        popul: &'a [Vec<f64>],
        vturb: &'a [f64],
        wnhint: &'a [Vec<f64>],
    ) -> LymlinParams<'a> {
        LymlinParams {
            nd,
            ielh: 1,  // 假设有氢
            n0h: 0,   // 氢原子能级起始
            n1h: 29,  // 氢原子能级终止 (30 个能级)
            nkh: 30,  // 氢离子能级
            nlh: 30,  // 氢原子能级数
            temp,
            elec,
            popul,
            vturb,
            wnhint,
            freq,
            ij: 0,
            abso1,
            emis1,
        }
    }

    #[test]
    fn test_lymlin_no_hydrogen() {
        let mut cache = LymlinCache::default();
        let mut abso1 = vec![0.0; 10];
        let mut emis1 = vec![0.0; 10];
        let freq = vec![2.0e15; 100];
        let temp = vec![10000.0; 10];
        let elec = vec![1.0e13; 10];
        let popul = vec![vec![1.0e10; 10]; 50];
        let vturb = vec![0.0; 10];
        let wnhint = vec![vec![1.0; 10]; 80];

        let mut params = LymlinParams {
            nd: 10,
            ielh: 0,  // 无氢
            n0h: 0,
            n1h: 29,
            nkh: 30,
            nlh: 30,
            temp: &temp,
            elec: &elec,
            popul: &popul,
            vturb: &vturb,
            wnhint: &wnhint,
            freq: &freq,
            ij: 0,
            abso1: &mut abso1,
            emis1: &mut emis1,
        };

        lymlin(&mut params, &mut cache);

        // 无氢时应该不修改数组
        for i in 0..10 {
            assert!((params.abso1[i] - 0.0).abs() < 1e-15);
            assert!((params.emis1[i] - 0.0).abs() < 1e-15);
        }
    }

    #[test]
    fn test_lymlin_frequency_out_of_range() {
        let mut cache = LymlinCache::default();
        let mut abso1 = vec![0.0; 10];
        let mut emis1 = vec![0.0; 10];

        // 频率太低
        let freq_low = vec![1.0e15; 100];
        let temp = vec![10000.0; 10];
        let elec = vec![1.0e13; 10];
        let popul = vec![vec![1.0e10; 10]; 50];
        let vturb = vec![0.0; 10];
        let wnhint = vec![vec![1.0; 10]; 80];

        let mut params = create_test_params(
            10, &mut abso1, &mut emis1, &freq_low,
            &temp, &elec, &popul, &vturb, &wnhint,
        );

        lymlin(&mut params, &mut cache);

        // 频率超出范围时应该不修改数组
        for i in 0..10 {
            assert!((params.abso1[i] - 0.0).abs() < 1e-15);
            assert!((params.emis1[i] - 0.0).abs() < 1e-15);
        }
    }

    #[test]
    fn test_lymlin_valid_frequency() {
        let mut cache = LymlinCache::default();
        let mut abso1 = vec![0.0; 5];
        let mut emis1 = vec![0.0; 5];

        // 有效频率范围内的频率
        let freq = vec![2.5e15; 100];
        let temp = vec![10000.0; 5];
        let elec = vec![1.0e13; 5];
        let popul = vec![vec![1.0e10; 5]; 50];
        let vturb = vec![0.0; 5];
        let wnhint = vec![vec![1.0; 5]; 80];

        let mut params = create_test_params(
            5, &mut abso1, &mut emis1, &freq,
            &temp, &elec, &popul, &vturb, &wnhint,
        );

        lymlin(&mut params, &mut cache);

        // 检查缓存已初始化
        assert!(cache.initialized);

        // 检查 Stark 参数已计算
        for j in 2..=30 {
            assert!(cache.xkijl[j] > 0.0);
            assert!(cache.fijl[j] > 0.0);
            assert!(cache.wl0l[j] > 0.0);
            assert!(cache.fr0l[j] > 0.0);
        }
    }

    #[test]
    fn test_cache_initialization() {
        let cache = LymlinCache::default();

        // 检查默认值
        assert!(!cache.initialized);
        assert_eq!(cache.xkijl.len(), MLEVL + 1);
        assert_eq!(cache.pj.len(), MLEVL + 1);
        assert_eq!(cache.pj[0].len(), MDEPTH);
    }
}