SpectraRust/src/tlusty/math/eos/rhonen.rs

//! 迭代求解粒子密度和电子密度。
//!
//! 重构自 TLUSTY `rhonen.f`。
//!
//! 功能：
//! - 从给定的温度和质量密度迭代求解总粒子密度和电子密度
//! - 使用 eldens 计算电子密度

use crate::tlusty::math::{eldens_pure, EldensConfig, EldensOutput, EldensParams};
use crate::tlusty::state::constants::{HMASS, UN};

/// RHONEN 输入参数
pub struct RhonenParams<'a> {
    /// 深度点索引 (1-indexed)
    pub id: usize,
    /// 温度 [K]
    pub t: f64,
    /// 质量密度 [g/cm³]
    pub rho: f64,
    /// 平均分子量数组
    pub wmm: &'a [f64],
    /// 初始电子相对密度 (0 表示自动估计)
    pub anerel: f64,
    /// eldens 配置
    pub eldens_config: EldensConfig,
    /// eldens 需要的额外参数
    pub eldens_ytot: f64,
    pub eldens_qref: f64,
    pub eldens_dqnr: f64,
    pub eldens_wmy: f64,
}

/// RHONEN 输出结果
#[derive(Debug, Clone)]
pub struct RhonenOutput {
    /// 总粒子密度 [cm⁻³]
    pub an: f64,
    /// 电子密度 [cm⁻³]
    pub ane: f64,
    /// 更新后的电子相对密度
    pub anerel: f64,
    /// 更新后的平均分子量
    pub wm: f64,
    /// 内能 [erg]
    pub enrgi: f64,
    /// 熵 [erg/K]
    pub entt: f64,
    /// 迭代次数
    pub iterations: i32,
    /// 是否收敛
    pub converged: bool,
}

/// 迭代求解 N 和 Ne。
///
/// # 参数
/// * `params` - 输入参数
///
/// # 返回值
/// 包含粒子密度、电子密度等的输出结构体
pub fn rhonen_pure(params: &RhonenParams) -> RhonenOutput {
    let id = params.id;
    let t = params.t;
    let rho = params.rho;

    // 初始化电子相对密度
    let mut anerel = if id == 1 && params.anerel == 0.0 {
        // 根据温度估计初始值
        if t < 4000.0 {
            1e-6
        } else if t < 5000.0 {
            1e-5
        } else if t < 5500.0 {
            1e-4
        } else if t < 6000.0 {
            1e-3
        } else if t < 7000.0 {
            0.01
        } else if t < 8000.0 {
            0.1
        } else if t < 9000.0 {
            0.4
        } else {
            0.5
        }
    } else {
        params.anerel
    };

    // 计算初始平均分子量
    // Fortran: wm = wmm(id) * (un - anerel) / hmass
    // 注意：wmm 数组是 0-indexed
    let wmm_id = params.wmm[id - 1];
    let mut wm = wmm_id * (UN - anerel) / HMASS;
    let mut wm0 = wm;

    let mut an = 0.0;
    let mut ane = 0.0;
    let mut enrgi = 0.0;
    let mut entt = 0.0;
    let mut converged = false;
    let max_iterations = 30;

    for it in 1..=max_iterations {
        // 计算粒子密度
        // Fortran: an = rho / wm / hmass
        an = rho / wm / HMASS;

        // 保存旧值用于收敛检查
        let ane0 = anerel * an;
        wm0 = wm;

        // 调用 eldens 计算电子密度
        let eldens_params = EldensParams {
            id,
            t,
            an,
            ytot: params.eldens_ytot,
            qref: params.eldens_qref,
            dqnr: params.eldens_dqnr,
            wmy: params.eldens_wmy,
            config: params.eldens_config.clone(),
            state_params: None,
            molecule_data: None,
            anato_data: None,
        };

        let eldens_output = eldens_pure(&eldens_params, 0);

        ane = eldens_output.ane;
        enrgi = eldens_output.energ;
        entt = eldens_output.entt;
        wm = eldens_output.wm;

        // 更新电子相对密度
        if an > 0.0 {
            anerel = ane / an;
        }

        // 收敛检查
        // Fortran: abs((ane-ane0)/ane0).lt.1.e-5 .and. abs((wm-wm0)/wm0).lt.1.e-5
        let ane_converged = if ane0.abs() > 1e-30 {
            ((ane - ane0) / ane0).abs() < 1e-5
        } else {
            ane.abs() < 1e-30
        };

        let wm_converged = if wm0.abs() > 1e-30 {
            ((wm - wm0) / wm0).abs() < 1e-5
        } else {
            wm.abs() < 1e-30
        };

        if ane_converged && wm_converged {
            converged = true;
            return RhonenOutput {
                an,
                ane,
                anerel,
                wm,
                enrgi,
                entt,
                iterations: it,
                converged,
            };
        }
    }

    // 未收敛，返回最后一次迭代结果
    RhonenOutput {
        an,
        ane,
        anerel,
        wm,
        enrgi,
        entt,
        iterations: max_iterations,
        converged,
    }
}

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

    #[test]
    fn test_rhonen_hot_star() {
        // 高温恒星情况 (T > 9000 K)
        let wmm = vec![1.4e-24; 10];
        let config = EldensConfig::default();

        let params = RhonenParams {
            id: 1,
            t: 15000.0,  // 高温
            rho: 1e-7,   // 典型光球层密度
            wmm: &wmm,
            anerel: 0.0, // 自动估计
            eldens_config: config,
            eldens_ytot: 1.0,
            eldens_qref: 1.0,
            eldens_dqnr: 0.0,
            eldens_wmy: 1.0,
        };

        let result = rhonen_pure(&params);

        // 验证基本物理约束
        assert!(result.an > 0.0, "粒子密度应为正");
        assert!(result.ane > 0.0, "电子密度应为正");
        assert!(result.anerel > 0.0 && result.anerel < 1.0, "电子相对密度应在 (0, 1) 范围内");
        assert!(result.iterations > 0, "应至少迭代一次");
    }

    #[test]
    fn test_rhonen_cool_star() {
        // 低温恒星情况 (T < 6000 K)
        let wmm = vec![1.4e-24; 10];
        let config = EldensConfig::default();

        let params = RhonenParams {
            id: 1,
            t: 5000.0,   // 低温
            rho: 1e-6,   // 较高密度
            wmm: &wmm,
            anerel: 0.0, // 自动估计
            eldens_config: config,
            eldens_ytot: 1.0,
            eldens_qref: 1.0,
            eldens_dqnr: 0.0,
            eldens_wmy: 1.0,
        };

        let result = rhonen_pure(&params);

        // 验证基本物理约束
        assert!(result.an > 0.0, "粒子密度应为正");
        assert!(result.ane > 0.0, "电子密度应为正");
        // 低温时电子相对密度应该较低
        assert!(result.anerel < 0.1, "低温时电子相对密度应较低");
    }

    #[test]
    fn test_rhonen_initial_anerel() {
        // 测试提供初始 anerel 的情况
        let wmm = vec![1.4e-24; 10];
        let config = EldensConfig::default();

        let params = RhonenParams {
            id: 2,  // 非 1，所以不会重新估计
            t: 8000.0,
            rho: 1e-7,
            wmm: &wmm,
            anerel: 0.3,  // 提供初始值
            eldens_config: config,
            eldens_ytot: 1.0,
            eldens_qref: 1.0,
            eldens_dqnr: 0.0,
            eldens_wmy: 1.0,
        };

        let result = rhonen_pure(&params);

        assert!(result.an > 0.0);
        assert!(result.ane > 0.0);
    }

    #[test]
    fn test_anerel_temperature_estimate() {
        // 测试温度对初始 anerel 估计的影响
        let wmm = vec![1.4e-24; 10];
        let config = EldensConfig::default();

        // 测试不同温度
        let temps = [3500.0, 4500.0, 5500.0, 6500.0, 7500.0, 8500.0, 10000.0];

        for &t in &temps {
            let params = RhonenParams {
                id: 1,
                t,
                rho: 1e-7,
                wmm: &wmm,
                anerel: 0.0,
                eldens_config: config.clone(),
                eldens_ytot: 1.0,
                eldens_qref: 1.0,
                eldens_dqnr: 0.0,
                eldens_wmy: 1.0,
            };

            let result = rhonen_pure(&params);
            assert!(result.ane >= 0.0, "温度 {} K 时电子密度应为非负", t);
        }
    }
}