SpectraRust/src/math/levsol.rs

//! 能级占据数求解器。
//!
//! 重构自 TLUSTY `levsol.f`。
//!
//! 通过求解速率方程得到新的能级占据数。

use crate::math::lineqs::lineqs_nr;
use crate::state::atomic::AtoPar;
use crate::state::config::{BasNum, InpPar};
use crate::state::constants::MLEVEL;

// ============================================================================
// LEVSOL - 能级求解器
// ============================================================================

/// 求解速率方程得到新的能级占据数。
///
/// 通过两种方式之一：
/// a) 反演全局速率矩阵 (如果 IRSPLT=0)
/// b) 反演各个化学元素的部分速率矩阵
///
/// # 参数
///
/// - `a` - 速率矩阵 A(MLEVEL,MLEVEL)
/// - `b` - 右端向量 B(MLEVEL)
/// - `popp` - 输出占据数 POPP(MLEVEL)
/// - `iical` - 能级索引映射 IICAL(MLEVEL)
/// - `nlvcal` - 实际计算的能级数
/// - `iall` - 标志 (0 表示跳过固定元素)
/// - `basnum` - 基本数值参数 (包含 natom, ioptab)
/// - `inppar` - 输入参数 (包含 irsplt)
/// - `atopar` - 原子参数 (包含 n0a, nka, iifix)
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE LEVSOL(A,B,POPP,IICAL,NLVCAL,IALL)
///   ...
/// END
/// ```
pub fn levsol(
    a: &mut [f64],
    b: &mut [f64],
    popp: &mut [f64],
    iical: &[i32],
    nlvcal: usize,
    iall: i32,
    basnum: &BasNum,
    inppar: &InpPar,
    atopar: &AtoPar,
) {
    // 检查是否跳过
    if basnum.ioptab < 0 {
        return;
    }

    // 方式 a: 反演全局速率矩阵
    if inppar.irsplt == 0 {
        lineqs_nr(a, b, popp, nlvcal, MLEVEL);
        return;
    }

    // 方式 b: 反演各个化学元素的部分速率矩阵
    let natom = basnum.natom as usize;

    // 工作数组
    let mut ap = vec![0.0; MLEVEL * MLEVEL];
    let mut bp = vec![0.0; MLEVEL];
    let mut popp1 = vec![0.0; MLEVEL];

    for iat in 0..natom {
        // 跳过固定元素
        if atopar.iifix[iat] == 1 && iall == 0 {
            continue;
        }

        let mut n1 = atopar.n0a[iat] as usize;
        let nk = atopar.nka[iat] as usize;

        n1 = iical[n1] as usize;
        let mut nk_idx = iical[nk] as usize;

        // 查找有效的起始索引
        if n1 == 0 {
            for i in atopar.n0a[iat] as usize..=atopar.nka[iat] as usize {
                let idx = iical[i] as usize;
                if idx > 0 {
                    n1 = idx;
                    break;
                }
            }
        }

        if n1 == 0 {
            continue;
        }

        // 修正 nk_idx (Fortran 中可能是 0，这里需要调整)
        if nk_idx == 0 {
            nk_idx = n1;
        }

        let nlp = nk_idx - n1 + 1;
        if nlp == 0 || n1 + nlp > MLEVEL {
            continue;
        }

        // 提取部分矩阵
        for i in 0..nlp {
            for j in 0..nlp {
                ap[j * MLEVEL + i] = a[(n1 + i) * MLEVEL + (n1 + j)];
            }
            bp[i] = b[n1 + i];
        }

        // 求解部分矩阵
        lineqs_nr(&mut ap, &mut bp, &mut popp1, nlp, MLEVEL);

        // 存储结果
        for i in 0..nlp {
            popp[n1 + i] = popp1[i];
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::state::atomic::AtoPar;
    use crate::state::config::{BasNum, InpPar};
    use crate::state::constants::MLEVEL;

    fn create_test_atopar() -> AtoPar {
        let mut atopar = AtoPar::default();
        atopar.n0a[0] = 0;
        atopar.n0a[1] = 3;
        atopar.nka[0] = 2;
        atopar.nka[1] = 5;
        atopar.iifix[0] = 0;
        atopar.iifix[1] = 0;
        atopar
    }

    fn create_test_basnum() -> BasNum {
        BasNum {
            natom: 2,
            ioptab: 0,
            ..Default::default()
        }
    }

    fn create_test_inppar() -> InpPar {
        InpPar {
            irsplt: 0,
            ..Default::default()
        }
    }

    #[test]
    fn test_levsol_global_matrix() {
        let mut a = vec![0.0; MLEVEL * MLEVEL];
        let mut b = vec![0.0; MLEVEL];
        let mut popp = vec![0.0; MLEVEL];
        let iical = vec![0i32; MLEVEL];

        // 设置简单的 2x2 系统
        // [2 1] [x]   [3]
        // [1 2] [y] = [3]
        // 解是 x=1, y=1
        // Fortran 列优先: A(j,i) = a[j + i*MLEVEL]
        a[0 + 0 * MLEVEL] = 2.0; // A(1,1)
        a[1 + 0 * MLEVEL] = 1.0; // A(2,1)
        a[0 + 1 * MLEVEL] = 1.0; // A(1,2)
        a[1 + 1 * MLEVEL] = 2.0; // A(2,2)
        b[0] = 3.0;
        b[1] = 3.0;

        let basnum = create_test_basnum();
        let inppar = create_test_inppar();
        let atopar = create_test_atopar();

        levsol(
            &mut a,
            &mut b,
            &mut popp,
            &iical,
            2,
            1,
            &basnum,
            &inppar,
            &atopar,
        );

        // 解应该是 x=1, y=1
        assert!((popp[0] - 1.0).abs() < 1e-10);
        assert!((popp[1] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_levsol_skip_negative_ioptab() {
        let basnum = BasNum {
            ioptab: -2,
            ..Default::default()
        };
        let inppar = create_test_inppar();
        let atopar = create_test_atopar();

        let mut a = vec![0.0; MLEVEL * MLEVEL];
        let mut b = vec![0.0; MLEVEL];
        let mut popp = vec![1.0; MLEVEL]; // 初始值
        let iical = vec![0i32; MLEVEL];

        let popp_before = popp[0];
        levsol(
            &mut a,
            &mut b,
            &mut popp,
            &iical,
            2,
            1,
            &basnum,
            &inppar,
            &atopar,
        );

        // 应该跳过，popp 保持不变
        assert!((popp[0] - popp_before).abs() < 1e-10);
    }
}