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

//! 能级粒子数初始化控制器。
//!
//! 重构自 TLUSTY `change.f`。
//!
//! 功能：
//! - 当显式能级系统与输入模型不一致时，控制初始能级占据数的评估
//! - 支持多种模式来计算新能级的粒子数
//! - 使用 STEQEQ 计算 LTE 粒子数

use crate::tlusty::math::SteqeqConfig;
use crate::tlusty::state::constants::{BOLK, MDEPTH};
// f2r_depends: READBF(I/O layer), STEQEQ

/// CHANGE 配置参数
#[derive(Debug, Clone)]
pub struct ChangeConfig {
    /// ICHANG 参数
    /// < 0: 通用粒子数变化
    /// > 0: 简化变化，ICHANG 是旧模型数据文件的单元号
    pub ichang: i32,
    /// 旧模型的能级数
    pub nlev0: i32,
    /// 是否 LTE
    pub lte: bool,
    /// Saha 常数
    pub saha_const: f64,
}

impl Default for ChangeConfig {
    fn default() -> Self {
        Self {
            ichang: 0,
            nlev0: 0,
            lte: false,
            saha_const: 2.0706e-16, // S = 2.0706D-16
        }
    }
}

/// 能级映射信息（用于 MODE 0-2）
#[derive(Debug, Clone)]
pub struct LevelMapping {
    /// 当前能级索引
    pub ii: usize,
    /// IOLD: 旧能级索引（0 表示无对应）
    pub iold: i32,
    /// MODE: 粒子数计算模式
    pub mode: i32,
    /// NXTOLD: 下一电离态的旧能级索引
    pub nxtold: i32,
    /// ISINEW: 新能级索引（用于 b-因子）
    pub isinew: i32,
    /// ISIOLD: 旧能级索引（用于 b-因子）
    pub isiold: i32,
    /// NXTSIO: 下一电离态的旧索引
    pub nxtsio: i32,
    /// REL: 粒子数乘数
    pub rel: f64,
}

/// STEQEQ 计算闭包包装
pub struct SteqeqFnWrapper(pub Box<dyn Fn(usize) -> Vec<f64>>);

impl std::fmt::Debug for SteqeqFnWrapper {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SteqeqFnWrapper")
    }
}

/// CHANGE 输入参数
pub struct ChangeParams<'a> {
    /// 配置参数
    pub config: ChangeConfig,
    /// 深度点数
    pub nd: usize,
    /// 能级数
    pub nlevel: usize,
    /// 离子数
    pub nion: usize,
    /// 原子数
    pub natom: usize,
    /// 温度数组
    pub temp: &'a [f64],
    /// 电子密度数组
    pub elec: &'a [f64],
    /// 当前粒子数 [能级 × 深度]
    pub popul: &'a [[f64; MDEPTH]],
    /// 能级权重 g
    pub g: &'a [f64],
    /// 电离能 [能级]
    pub enion: &'a [f64],
    /// 能级所属元素索引 [能级]
    pub iel: &'a [i32],
    /// 下一能级索引 [能级]
    pub nnext: &'a [i32],
    /// 离子起始能级 [离子]
    pub nfirst: &'a [i32],
    /// 离子结束能级 [离子]
    pub nlast: &'a [i32],
    /// 能级所属原子 [能级]
    pub iatm: &'a [i32],
    /// 原子编号 [原子]
    pub numat: &'a [i32],
    /// 能级电荷 [能级]
    pub iz: &'a [i32],
    /// 原子能级范围起始 [原子]
    pub nka: &'a [i32],
    /// 能级映射（仅用于 ICHANG < 0）
    pub level_mappings: Vec<LevelMapping>,
    /// STEQEQ 配置
    pub steqeq_config: SteqeqConfig,
    /// 占据概率 [能级 × 深度]
    pub wop: &'a [[f64; MDEPTH]],
    /// STEQEQ 计算函数：给定深度点索引(0-based)，返回 [nlevel] 个 LTE 粒子数
    pub steqeq_fn: Option<SteqeqFnWrapper>,
}

/// CHANGE 输出结果
#[derive(Debug, Clone)]
pub struct ChangeOutput {
    /// 更新后的粒子数 [能级 × 深度]
    pub popul: Vec<Vec<f64>>,
    /// 临时粒子数（LTE）
    pub popull: Vec<Vec<f64>>,
    /// 是否成功
    pub success: bool,
}

/// 处理能级粒子数变化（纯计算函数）。
///
/// # 参数
/// * `params` - 输入参数
///
/// # 返回值
/// 包含更新后的粒子数
pub fn change_pure(params: &ChangeParams) -> ChangeOutput {
    let nd = params.nd;
    let nlevel = params.nlevel;
    let config = &params.config;

    // 初始化输出数组
    let mut popul0 = vec![vec![0.0; nd]; nlevel];
    let mut popull = vec![vec![0.0; nd]; nlevel];
    let mut popul_new = vec![vec![0.0; nd]; nlevel];

    // 复制当前粒子数
    for ii in 0..nlevel {
        for id in 0..nd {
            popul_new[ii][id] = params.popul[ii][id];
        }
    }

    if config.ichang < 0 {
        // 通用粒子数变化
        handle_general_change(params, &mut popul0, &mut popull, &mut popul_new);
    } else if config.ichang > 0 {
        // 简化变化
        handle_simplified_change(params, &mut popul0, &mut popull, &mut popul_new);
    }

    ChangeOutput {
        popul: popul_new,
        popull,
        success: true,
    }
}

/// 处理通用粒子数变化（ICHANG < 0）
fn handle_general_change(
    params: &ChangeParams,
    popul0: &mut [Vec<f64>],
    popull: &mut [Vec<f64>],
    popul_new: &mut [Vec<f64>],
) {
    let nd = params.nd;
    let nlevel = params.nlevel;
    let config = &params.config;
    let mut ifese = 0;

    for mapping in &params.level_mappings {
        let ii = mapping.ii;
        let mode = mapping.mode;
        let iold = mapping.iold;
        let rel = if mapping.rel == 0.0 { 1.0 } else { mapping.rel };

        if mode >= 3 {
            ifese += 1;
        }

        for id in 0..nd {
            if iold != 0 {
                // 直接使用旧粒子数
                popul0[ii][id] = params.popul[(iold - 1) as usize][id];
            } else {
                match mode {
                    0 => {
                        // MODE 0: 粒子数 = 旧粒子数 × REL
                        popul0[ii][id] = params.popul[(mapping.isiold - 1) as usize][id] * rel;
                    }
                    1 | 2 => {
                        // MODE 1/2: 使用 Saha 方程
                        let t = params.temp[id];
                        let ane = params.elec[id];

                        // 计算 Saha 因子
                        let nxtnew = params.nnext[params.iel[ii] as usize] as usize;
                        let sb = config.saha_const / t / t.sqrt()
                            * params.g[ii]
                            / params.g[nxtnew]
                            * (params.enion[ii] / t / BOLK).exp();

                        if mode == 1 {
                            // MODE 1: LTE 相对于下一电离态
                            popul0[ii][id] = sb * ane * params.popul[(mapping.nxtold - 1) as usize][id] * rel;
                        } else {
                            // MODE 2: 使用 b-因子
                            let kk = mapping.isinew as usize;
                            let knext = params.nnext[params.iel[kk] as usize] as usize;
                            let sbk = config.saha_const / t / t.sqrt()
                                * params.g[kk]
                                / params.g[knext]
                                * (params.enion[kk] / t / BOLK).exp();

                            popul0[ii][id] = sb / sbk
                                * params.popul[(mapping.nxtold - 1) as usize][id]
                                / params.popul[(mapping.nxtsio - 1) as usize][id]
                                * params.popul[(mapping.isiold - 1) as usize][id]
                                * rel;
                        }
                    }
                    _ => {
                        // MODE >= 3: LTE
                        if ifese == 1 {
                            // 首次遇到 MODE>=3，调用 STEQEQ 计算 LTE 粒子数
                            if let Some(ref steqeq_fn) = params.steqeq_fn {
                                let popl = (steqeq_fn.0)(id);
                                for iii in 0..nlevel {
                                    popull[iii][id] = popl[iii];
                                }
                            }
                        }
                        popul0[ii][id] = popull[ii][id];
                    }
                }
            }
        }
    }

    // 复制结果到输出
    for ii in 0..nlevel {
        for id in 0..nd {
            popul_new[ii][id] = popul0[ii][id];
        }
    }
}

/// 处理简化粒子数变化（ICHANG > 0）
fn handle_simplified_change(
    params: &ChangeParams,
    popul0: &mut [Vec<f64>],
    _popull: &mut [Vec<f64>],
    popul_new: &mut [Vec<f64>],
) {
    let nd = params.nd;
    let nlevel = params.nlevel;
    let config = &params.config;

    // 设置 LTE 标志并计算所有深度点的 LTE 粒子数
    // Fortran: LTE=.TRUE.; DO ID=1,ND; CALL STEQEQ(ID,POPL,0); POPUL0(II,ID)=POPL(II)
    for id in 0..nd {
        if let Some(ref steqeq_fn) = params.steqeq_fn {
            let popl = (steqeq_fn.0)(id);
            for ii in 0..nlevel {
                popul0[ii][id] = popl[ii];
            }
        }
    }

    // WRITE(6,600) - ' Levels: OLD model -> NEW model'
    eprintln!(" Levels: OLD model -> NEW model");
    eprintln!(" ------------------------------");

    if config.ichang == 1 {
        // 只更新新能级
        for ii in config.nlev0 as usize..nlevel {
            // WRITE(6,601) JL,IL - '**** CHANGE: old, new =',2I5
            eprintln!("          {:8}     {:8}", ii, ii);
            for id in 0..nd {
                popul_new[ii][id] = popul0[ii][id];
            }
        }
    }
    // 其他情况（ichang > 1）需要从文件读取旧模型数据
    // 这部分在 I/O 层处理
}

/// 计算 Saha-Boltzmann 因子
///
/// # 参数
/// * `s` - Saha 常数
/// * `t` - 温度
/// * `g1` - 能级 1 的统计权重
/// * `g2` - 能级 2 的统计权重
/// * `enion` - 电离能
///
/// # 返回值
/// Saha 因子
pub fn saha_factor(s: f64, t: f64, g1: f64, g2: f64, enion: f64) -> f64 {
    s / t / t.sqrt() * g1 / g2 * (enion / t / BOLK).exp()
}

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

    #[test]
    fn test_saha_factor() {
        let s = 2.0706e-16;
        let t = 10000.0;
        let g1 = 2.0;
        let g2 = 1.0;
        let enion = 13.6;

        let result = saha_factor(s, t, g1, g2, enion);
        println!("Saha factor: {}", result);

        // 基本检查：应该是正数
        assert!(result > 0.0);
    }

    #[test]
    fn test_change_basic() {
        // 使用简化的测试参数
        let nd = 5;
        let nlevel = 10;

        // 创建测试数据
        let temp: Vec<f64> = (0..nd).map(|i| 10000.0 - i as f64 * 500.0).collect();
        let elec: Vec<f64> = vec![1e13; nd];
        let mut popul_arr = vec![[0.0; MDEPTH]; MLEVEL];
        let mut wop_arr = vec![[0.0; MDEPTH]; MLEVEL];
        for i in 0..nlevel {
            for j in 0..nd {
                popul_arr[i][j] = 1e10;
                wop_arr[i][j] = 1.0;
            }
        }

        let g = vec![2.0; MLEVEL];
        let enion = vec![13.6; MLEVEL];
        let iel: Vec<i32> = (0..MLEVEL as i32).collect();
        let nnext: Vec<i32> = (1..=MLEVEL as i32).collect();
        let nfirst = vec![0; 200];
        let nlast = vec![0; 200];
        let iatm = vec![1; MLEVEL];
        let numat = vec![0; 100];
        let iz = vec![0; MLEVEL];
        let nka = vec![0; 100];

        // 使用 Box::leak 创建 'static 引用
        let temp: &'static [f64] = Box::leak(temp.into_boxed_slice());
        let elec: &'static [f64] = Box::leak(elec.into_boxed_slice());
        let popul: &'static [[f64; MDEPTH]] = Box::leak(popul_arr.into_boxed_slice().try_into().unwrap());
        let g: &'static [f64] = Box::leak(g.into_boxed_slice().try_into().unwrap());
        let enion: &'static [f64] = Box::leak(enion.into_boxed_slice().try_into().unwrap());
        let iel: &'static [i32] = Box::leak(iel.into_boxed_slice().try_into().unwrap());
        let nnext: &'static [i32] = Box::leak(nnext.into_boxed_slice().try_into().unwrap());
        let nfirst: &'static [i32] = Box::leak(nfirst.into_boxed_slice().try_into().unwrap());
        let nlast: &'static [i32] = Box::leak(nlast.into_boxed_slice().try_into().unwrap());
        let iatm: &'static [i32] = Box::leak(iatm.into_boxed_slice().try_into().unwrap());
        let numat: &'static [i32] = Box::leak(numat.into_boxed_slice().try_into().unwrap());
        let iz: &'static [i32] = Box::leak(iz.into_boxed_slice().try_into().unwrap());
        let nka: &'static [i32] = Box::leak(nka.into_boxed_slice().try_into().unwrap());
        let wop: &'static [[f64; MDEPTH]] = Box::leak(wop_arr.into_boxed_slice().try_into().unwrap());

        let params = ChangeParams {
            config: ChangeConfig::default(),
            nd,
            nlevel,
            nion: 5,
            natom: 1,
            temp,
            elec,
            popul,
            g,
            enion,
            iel,
            nnext,
            nfirst,
            nlast,
            iatm,
            numat,
            iz,
            nka,
            level_mappings: vec![],
            steqeq_config: SteqeqConfig::default(),
            wop,
            steqeq_fn: None,
        };

        let output = change_pure(&params);

        // 检查输出维度
        assert_eq!(output.popul.len(), nlevel);
        assert!(output.success);
    }
}