SpectraRust/src/math/eldenc.rs

//! 电子密度比较和贡献者分析。
//!
//! 重构自 TLUSTY `eldenc.f`。
//!
//! 功能：
//! - 比较实际电子密度和不透明度表中的插值电子密度
//! - 分析各元素对电子密度的贡献
//! - 输出电子供体信息

use crate::math::moleq::{moleq_pure, MoleqParams};
use crate::math::rhonen::{rhonen_pure, RhonenParams};
use crate::math::state::{state_pure, StateParams};
use crate::state::constants::{MDEPTH, MLEVEL};

/// 最大温度表点数
pub const MTABT: usize = 21;
/// 最大密度表点数
pub const MTABR: usize = 19;
/// 最大显式原子数
pub const MAX_EXPLICIT_ATOMS: usize = 30;

/// ELDENC 配置参数
#[derive(Debug, Clone)]
pub struct EldencConfig {
    /// 是否输出电子密度检查 (IPELCH)
    pub ipelch: i32,
    /// 是否输出电子供体 (IPELDO)
    pub ipeldo: i32,
    /// 磁盘模型标志 (IDISK)
    pub idisk: i32,
    /// 是否包含分子 (IFMOL)
    pub ifmol: i32,
    /// 分子温度上限 (TMOLIM)
    pub tmolim: f64,
    /// H- 元素索引 (IELHM)
    pub ielhm: i32,
    /// 氢元素索引 (IELH)
    pub ielh: usize,
    /// QM 常数
    pub qm: f64,
}

impl Default for EldencConfig {
    fn default() -> Self {
        Self {
            ipelch: 0,
            ipeldo: 0,
            idisk: 0,
            ifmol: 0,
            tmolim: 1e10,
            ielhm: -1,
            ielh: 1,
            qm: 1.0,
        }
    }
}

/// ELDENC 输入参数
pub struct EldencParams<'a> {
    /// 配置参数
    pub config: EldencConfig,
    /// 深度点数
    pub nd: usize,
    /// 温度数组
    pub temp: &'a [f64],
    /// 电子密度数组
    pub elec: &'a [f64],
    /// 质量密度数组
    pub dens: &'a [f64],
    /// 平均分子量数组
    pub wmm: &'a [f64],
    /// 总丰度因子数组
    pub ytot: &'a [f64],
    /// 温度表点数
    pub numtemp: usize,
    /// 温度向量 (对数)
    pub tempvec: &'a [f64],
    /// 密度矩阵 (对数)
    pub rhomat: &'a [[f64; MTABR]],
    /// 每个温度点的密度点数
    pub numrh: &'a [i32],
    /// 电子密度表 (对数)
    pub elecgr: &'a [[f64; MTABR]],
    /// 显式原子索引数组 (IATEX)
    pub iatex: &'a [i32],
    /// 原子能级范围起始 [原子] (N0A)
    pub n0a: &'a [i32],
    /// 原子能级范围结束 [原子] (NKA)
    pub nka: &'a [i32],
    /// 离子起始能级 [离子] (NFIRST)
    pub nfirst: &'a [i32],
    /// 能级所属元素 [能级] (IEL)
    pub iel: &'a [i32],
    /// l-量子数链接 [能级] (ILK)
    pub ilk: &'a [i32],
    /// 元素电荷 [元素] (IZ)
    pub iz: &'a [i32],
    /// 能级粒子数 [能级 × 深度] (POPUL)
    pub popul: &'a [[f64; MDEPTH]],
    /// RR 数组（丰度比）
    pub rr: &'a [[f64; 100]],
    /// 离子密度 [离子 × 深度]
    pub anion: &'a [[f64; MDEPTH]],
    /// H- 密度 [深度]
    pub anhm: &'a [f64],
    /// RHONEN 参数
    pub rhonen_params: Option<RhonenParams<'a>>,
    /// STATE 参数
    pub state_params: Option<StateParams<'a>>,
    /// MOLEQ 参数
    pub moleq_params: Option<MoleqParams<'a>>,
}

/// 电子密度比较结果
#[derive(Debug, Clone)]
pub struct EldencOutput {
    /// 插值电子密度 [深度]
    pub elecg: Vec<f64>,
    /// LTE 电子密度 [深度]
    pub ane_lte: Vec<f64>,
    /// 相对贡献 [31 × 深度]
    pub elcon: Vec<Vec<f64>>,
}

/// 比较电子密度并分析贡献者。
///
/// # 参数
/// * `params` - 输入参数
///
/// # 返回值
/// 包含插值电子密度和贡献者信息的输出
pub fn eldenc_pure(params: &EldencParams) -> EldencOutput {
    let nd = params.nd;
    let config = &params.config;

    let mut elecg = vec![0.0; nd];
    let mut ane_lte = vec![0.0; nd];
    let mut elcon = vec![vec![0.0; nd]; 31];

    // 如果不需要输出，直接返回
    if config.ipelch <= 0 && (config.idisk == 0 || config.ipeldo == 0) {
        return EldencOutput {
            elecg,
            ane_lte,
            elcon,
        };
    }

    // 电子密度检查
    if config.ipelch > 0 {
        for id in 0..nd {
            let t = params.temp[id];
            let rho = params.dens[id];

            // 从不透明度表插值电子密度
            if params.numtemp == nd {
                // 直接使用表值
                elecg[id] = params.elecgr[id][0];
            } else {
                // 双线性插值
                elecg[id] = interpolate_elec(params, t, rho);
            }

            // 计算 LTE 电子密度
            if let Some(rhonen_params) = &params.rhonen_params {
                let rhonen_out = rhonen_pure(rhonen_params);
                ane_lte[id] = rhonen_out.ane;
            }

            // 转换为实际电子密度
            elecg[id] = elecg[id].exp();
        }
    }

    // 电子供体分析
    if config.idisk != 0 && config.ipeldo != 0 {
        for id in 0..nd {
            let t = params.temp[id];

            if config.ifmol > 0 && t < config.tmolim {
                // 使用分子平衡
                let rho = params.dens[id];

                if let Some(rhonen_params) = &params.rhonen_params {
                    let rhonen_out = rhonen_pure(rhonen_params);
                    let aein = rhonen_out.ane;

                    // 调用 MOLEQ（简化）
                    for ia in 0..30 {
                        elcon[ia][id] = params.anion[ia][id] / params.elec[id];
                    }
                    elcon[30][id] = -params.anhm[id] / params.elec[id];
                }
            } else {
                // 使用 STATE
                if let Some(state_params) = &params.state_params {
                    let _state_out = state_pure(state_params);

                    for ia in 0..30 {
                        let iat = params.iatex[ia];
                        if iat > 0 {
                            let iat_idx = iat as usize;
                            let mut qs = 0.0;
                            let mut n1 = params.n0a[iat_idx] as usize;

                            if ia == 0 {
                                n1 = params.nfirst[config.ielh] as usize;
                            }

                            let nk = params.nka[iat_idx] as usize;
                            for i in n1..=nk {
                                if i >= MLEVEL {
                                    break;
                                }
                                let mut ch = (params.iz[params.iel[i] as usize] - 1) as f64;
                                if params.ilk[i] > 0 {
                                    ch = params.iz[params.ilk[i] as usize] as f64;
                                }
                                qs += ch * params.popul[i][id];
                            }
                            elcon[ia][id] = qs / params.elec[id];
                        } else {
                            elcon[ia][id] = params.rr[ia][99] / params.elec[id];
                        }
                    }

                    // H- 贡献
                    if config.ielhm > 0 {
                        elcon[30][id] = -params.popul[params.nfirst[config.ielhm as usize] as usize][id]
                            / params.elec[id];
                    } else {
                        let aref = params.dens[id] / params.wmm[id] / params.ytot[id];
                        elcon[30][id] = -config.qm * params.rr[0][0] * aref / params.elec[id];
                    }
                }
            }
        }
    }

    EldencOutput {
        elecg,
        ane_lte,
        elcon,
    }
}

/// 从不透明度表插值电子密度
fn interpolate_elec(params: &EldencParams, t: f64, rho: f64) -> f64 {
    let ttab1 = params.tempvec[0];
    let ttab2 = params.tempvec[params.numtemp - 1];

    let tl = t.ln();
    let deltat = (tl - ttab1) / (ttab2 - ttab1) * (params.numtemp - 1) as f64;

    let mut jt = 1 + deltat.floor() as usize;
    if jt < 1 {
        jt = 1;
    }
    if jt > params.numtemp - 1 {
        jt = params.numtemp - 1;
    }

    let ju = jt + 1;
    let t1i = params.tempvec[jt - 1];
    let t2i = params.tempvec[ju - 1];
    let mut dti = (tl - t1i) / (t2i - t1i);
    if deltat < 0.0 {
        dti = 0.0;
    }

    if params.numrh[jt - 1] != 1 {
        // 低温度插值
        let numrho = params.numrh[jt - 1] as usize;
        let rtab1 = params.rhomat[jt - 1][0];
        let rtab2 = params.rhomat[jt - 1][numrho - 1];

        let rl = rho.ln();
        let deltar = (rl - rtab1) / (rtab2 - rtab1) * (numrho - 1) as f64;

        let mut jr = 1 + deltar.floor() as usize;
        if jr < 1 {
            jr = 1;
        }
        if jr > numrho - 1 {
            jr = numrho - 1;
        }

        let r1i = params.rhomat[jt - 1][jr - 1];
        let r2i = params.rhomat[jt - 1][jr];
        let mut dri = (rl - r1i) / (r2i - r1i);
        if deltar < 0.0 {
            dri = 0.0;
        }

        let opr1 = params.elecgr[jt - 1][jr - 1]
            + dri * (params.elecgr[jt - 1][jr] - params.elecgr[jt - 1][jr - 1]);

        // 高温度插值
        let numrho_u = params.numrh[ju - 1] as usize;
        let rtab1_u = params.rhomat[ju - 1][0];
        let rtab2_u = params.rhomat[ju - 1][numrho_u - 1];

        let deltar_u = (rl - rtab1_u) / (rtab2_u - rtab1_u) * (numrho_u - 1) as f64;

        let mut jr_u = 1 + deltar_u.floor() as usize;
        if jr_u < 1 {
            jr_u = 1;
        }
        if jr_u > numrho_u - 1 {
            jr_u = numrho_u - 1;
        }

        let r1i_u = params.rhomat[ju - 1][jr_u - 1];
        let r2i_u = params.rhomat[ju - 1][jr_u];
        let mut dri_u = (rl - r1i_u) / (r2i_u - r1i_u);
        if deltar_u < 0.0 {
            dri_u = 0.0;
        }

        let opr2 = params.elecgr[ju - 1][jr_u - 1]
            + dri_u * (params.elecgr[ju - 1][jr_u] - params.elecgr[ju - 1][jr_u - 1]);

        opr1 + (opr2 - opr1) * dti
    } else {
        // 单密度点
        let jr = 0;
        params.elecgr[jt - 1][jr] + (params.elecgr[ju - 1][jr] - params.elecgr[jt - 1][jr]) * dti
    }
}

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

    #[test]
    fn test_eldenc_skip() {
        // 测试跳过输出的情况
        let params = EldencParams {
            config: EldencConfig {
                ipelch: 0,
                ipeldo: 0,
                idisk: 0,
                ..Default::default()
            },
            nd: 5,
            temp: &[10000.0, 9500.0, 9000.0, 8500.0, 8000.0],
            elec: &[1e13; 5],
            dens: &[1e-7; 5],
            wmm: &[1.0; 5],
            ytot: &[1.0; 5],
            numtemp: 5,
            tempvec: &[9.0, 9.2, 9.4, 9.6, 9.8],
            rhomat: &[[0.0; MTABR]; MTABT],
            numrh: &[1; MTABT],
            elecgr: &[[0.0; MTABR]; MTABT],
            iatex: &[0; 30],
            n0a: &[0; 100],
            nka: &[10; 100],
            nfirst: &[0; 200],
            iel: &[0; MLEVEL],
            ilk: &[0; MLEVEL],
            iz: &[1; 200],
            popul: &[[0.0; MDEPTH]; MLEVEL],
            rr: &[[0.0; 100]; 30],
            anion: &[[0.0; MDEPTH]; 100],
            anhm: &[0.0; MDEPTH],
            rhonen_params: None,
            state_params: None,
            moleq_params: None,
        };

        let output = eldenc_pure(&params);

        // 检查输出维度
        assert_eq!(output.elecg.len(), 5);
        assert_eq!(output.ane_lte.len(), 5);
        assert_eq!(output.elcon.len(), 31);
    }
}