SpectraRust/src/tlusty/math/utils/wnstor.rs

//! 存储氢能级占据概率。
//!
//! 重构自 TLUSTY `wnstor.f`
//!
//! 将氢能级的占据概率存储到 WNCOM 公共块中，以便后续使用。

use crate::tlusty::math::wn;
use crate::tlusty::state::constants::{MION, MLEVEL, NLMX, UN};

/// 存储氢能级占据概率。
///
/// 计算氢能级的占据概率，并将结果存储到相应数组中。
///
/// # 参数
///
/// * `id` - 深度索引 (0-based)
/// * `temp` - 温度数组
/// * `elec` - 电子密度数组
/// * `xi2` - XI2 系数数组
/// * `wnhint` - 输出：氢能级占据概率积分 [NLMX][深度]
/// * `wop` - 输出：能级权重 [MLEVEL][深度]
/// * `ifwop` - 能级权重标志
/// * `nlevel` - 能级数
/// * `nquant` - 主量子数数组
/// * `iz` - 原子序数数组
/// * `io_ptab` - 占据概率表标志
/// * `lte` - LTE 标志
///
/// # 示例
///
/// ```
/// use tlusty_rust::math::wnstor::wnstor;
///
/// let temp = vec![10000.0; 100];
/// let elec = vec![1e12; 100];
/// let xi2 = vec![1.0; 30];
/// let mut wnhint = vec![vec![0.0; 100]; 30];
/// let mut wop = vec![vec![0.0; 100]; 1134];
/// let ifwop = vec![1; 1134];
/// let nquant = vec![1; 1134];
/// let iz = vec![1; 1134];
///
/// wnstor(0, &temp, &elec, &xi2, &mut wnhint, &mut wop,
///        &ifwop, 10, &nquant, &iz, 0, false);
/// ```
pub fn wnstor(
    id: usize,
    temp: &[f64],
    elec: &[f64],
    xi2: &[f64],
    wnhint: &mut [Vec<f64>],
    wop: &mut [Vec<f64>],
    ifwop: &[i32],
    nlevel: usize,
    nquant: &[i32],
    iel: &[i32],
    iz: &[i32],
    io_ptab: i32,
    lte: bool,
) {
    // 常数参数
    const P1: f64 = 0.1402;
    const P2: f64 = 0.1285;
    const P3: f64 = 1.0;
    const P4: f64 = 3.15;
    const P5: f64 = 4.0;
    const TKN: f64 = 3.01;
    const CKN: f64 = 5.33333333;
    const CB0: f64 = 8.59e14;
    const F23: f64 = -2.0 / 3.0;

    if io_ptab < 0 {
        return;
    }

    // 获取深度点数据
    let ane = elec[id];
    let t = temp[id];

    // 计算 A 参数
    let a = CKN * ane * (ane.ln() / 6.0).exp() / t.sqrt();

    let z = UN;
    let x = (UN + P3 * a).powf(P4);
    let c1 = P1 * (x + P5 * (z - UN) * a * a * a);
    let c2 = P2 * x;
    let beta0 = CB0 * z * z * z * ane.powf(F23);

    // 计算 WNHINT
    for i in 1..=NLMX {
        let xn = i as f64;
        let xkn = if xn <= TKN {
            UN
        } else {
            let xn1 = UN / (xn + UN);
            CKN * xn * xn1 * xn1
        };

        // XI2 系数
        let xi2_val = *xi2.get(i - 1).unwrap_or(&UN);
        let beta = beta0 * xkn * xi2_val * xi2_val;
        let f = (c1 * beta * beta * beta) / (UN + c2 * beta * beta.sqrt());
        wnhint[i - 1][id] = f / (UN + f);
    }

    // 计算 WOP - 显式能级的占据概率
    for ii in 0..nlevel {
        let ifwop_val = ifwop[ii];

        if ifwop_val <= 0 {
            continue;
        }

        let nq = nquant[ii] as usize;
        let iel_idx = iel[ii] as usize;
        let iz_val = if iel_idx > 0 && iel_idx <= iz.len() { iz[iel_idx - 1] as f64 } else { 0.0 };

        if iz_val == 1.0 {
            // 氢：使用 WNHINT
            if nq > 0 && nq <= NLMX {
                wop[ii][id] = wnhint[nq - 1][id];
            }
        } else {
            // 其他元素：使用 WN 函数
            let xn = nq as f64;
            // bergfc = 0.0 表示使用标准 Hummer-Mihalas 公式
            wop[ii][id] = wn(xn, a, ane, iz_val, 0.0);
        }

        if ifwop_val > 1 && lte {
            wop[ii][id] = UN;
        }
    }
}

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

    const ND: usize = 10;

    fn create_test_data() -> (
        Vec<f64>,
        Vec<f64>,
        Vec<f64>,
        Vec<Vec<f64>>,
        Vec<Vec<f64>>,
        Vec<i32>,
        Vec<i32>,
        Vec<i32>,
        Vec<i32>,
    ) {
        let temp = vec![10000.0; ND];
        let elec = vec![1e12; ND];
        let xi2 = vec![1.0; NLMX];
        let wnhint = vec![vec![0.0; ND]; NLMX];
        let wop = vec![vec![0.0; ND]; MLEVEL];
        let ifwop = vec![1; MLEVEL];
        let nquant = vec![1; MLEVEL];
        let iel = vec![1; MLEVEL];
        let iz = vec![1; MION];

        (temp, elec, xi2, wnhint, wop, ifwop, nquant, iel, iz)
    }

    #[test]
    fn test_wnstor_io_ptab_negative() {
        let (temp, elec, xi2, mut wnhint, mut wop, ifwop, nquant, iel, iz) = create_test_data();

        // io_ptab < 0 应该直接返回
        wnstor(
            0, &temp, &elec, &xi2, &mut wnhint, &mut wop,
            &ifwop, 10, &nquant, &iel, &iz, -1, false,
        );

        // wnhint 应该保持为 0
        assert!((wnhint[0][0] - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_wnstor_basic() {
        let (temp, elec, xi2, mut wnhint, mut wop, ifwop, nquant, iel, iz) = create_test_data();

        wnstor(
            0, &temp, &elec, &xi2, &mut wnhint, &mut wop,
            &ifwop, 10, &nquant, &iel, &iz, 0, false,
        );

        // 检查 WNHINT 已被计算，值应该在 [0, 1] 范围内
        for i in 0..NLMX {
            let val = wnhint[i][0];
            assert!(val >= 0.0 && val <= 1.0, "WNHINT[{}][0] = {} should be in [0,1]", i, val);
        }
    }

    #[test]
    fn test_wnstor_lte_flag() {
        let (temp, elec, xi2, mut wnhint, mut wop, ifwop, nquant, iel, iz) = create_test_data();

        wnstor(
            0, &temp, &elec, &xi2, &mut wnhint, &mut wop,
            &ifwop, 10, &nquant, &iel, &iz, 0, true, // lte = true
        );

        // 当 ifwop > 1 且 lte = true 时，wop 应该是 1.0
        for ii in 0..10 {
            if ifwop[ii] > 1 {
                assert!((wop[ii][0] - UN).abs() < 1e-10);
            }
        }
    }
}