SpectraRust/src/math/switch.rs

//! 碰撞-辐射切换参数评估。
//!
//! 重构自 TLUSTY `SWITCH` 子程序。
//!
//! # 功能
//!
//! - 计算碰撞-辐射切换参数 lambda(R) = CRSW
//! - 基于 Hummer & Voels (1988) 方法
//! - 支持深度依赖的切换参数

// ============================================================================
// 常量
// ============================================================================

/// 1.0
const UN: f64 = 1.0;
/// h/k (Planck/Boltzmann)
const HK: f64 = 4.7994e0;

// ============================================================================
// 输入/输出结构体
// ============================================================================

/// SWITCH 输入参数（初始化模式）。
pub struct SwitchInitParams<'a> {
    /// 深度点数
    pub nd: usize,
    /// 跃迁数
    pub ntrans: usize,
    /// 碰撞速率 COLRAT(ITR, ID)
    pub colrat: &'a [Vec<f64>],
    /// 向上辐射速率 RRU(ITR, ID)
    pub rru: &'a [Vec<f64>],
    /// 向下辐射速率 RRD(ITR, ID)
    pub rrd: &'a [Vec<f64>],
    /// 参考频率 FR0(ITR)
    pub fr0: &'a [f64],
    /// 温度 TEMP(ID)
    pub temp: &'a [f64],
    /// 是否是谱线 LINE(ITR)
    pub line: &'a [bool],
    /// ICRSW 模式 (0=禁用, 1=全局最小, 2=深度依赖)
    pub icrsw: i32,
    /// SWPFAC 因子
    pub swpfac: f64,
    /// SWPLIM 限制
    pub swplim: f64,
}

/// SWITCH 输入参数（更新模式）。
pub struct SwitchUpdateParams<'a> {
    /// 深度点数
    pub nd: usize,
    /// 当前 CRSW 值
    pub crsw: &'a mut [f64],
    /// SWPINC 增量因子
    pub swpinc: f64,
    /// SWPLIM 限制
    pub swplim: f64,
}

/// SWITCH 输出结果。
#[derive(Debug, Clone)]
pub struct SwitchOutput {
    /// CRSW 数组
    pub crsw: Vec<f64>,
    /// 全局最小值（仅在初始化模式）
    pub swmin: Option<f64>,
}

// ============================================================================
// 核心计算函数
// ============================================================================

/// 执行 SWITCH 初始化计算（INITM = 1）。
///
/// # 参数
/// - `params`: 输入参数
///
/// # 返回
/// CRSW 数组
pub fn switch_init(params: &SwitchInitParams) -> SwitchOutput {
    if params.icrsw == 0 {
        // 碰撞-辐射切换未启用
        return SwitchOutput {
            crsw: vec![UN; params.nd],
            swmin: None,
        };
    }

    let nd = params.nd;
    let ntrans = params.ntrans;

    let mut swtch = vec![UN; nd];
    let mut swmin = UN;

    // 遍历每个深度点
    for id in 0..nd {
        swtch[id] = UN;

        // 遍历每个跃迁
        for itr in 0..ntrans {
            let c = params.colrat.get(itr).and_then(|v| v.get(id)).copied().unwrap_or(0.0);
            let rru_val = params.rru.get(itr).and_then(|v| v.get(id)).copied().unwrap_or(0.0);
            let rrd_val = params.rrd.get(itr).and_then(|v| v.get(id)).copied().unwrap_or(0.0);

            if rru_val.abs() > 1e-30 {
                // 向上速率
                let swu = c / rru_val;

                // 向下速率
                let swd = if params.line.get(itr).copied().unwrap_or(false) {
                    // 谱线
                    let fr0 = params.fr0.get(itr).copied().unwrap_or(0.0);
                    let temp = params.temp.get(id).copied().unwrap_or(10000.0);
                    c * (HK * fr0 / temp).exp() / rrd_val
                } else {
                    // 连续谱
                    c / rrd_val
                };

                // 取最小值
                if swu < swtch[id] {
                    swtch[id] = swu;
                }
                if swd < swtch[id] {
                    swtch[id] = swd;
                }
            }
        }

        // 更新全局最小值
        if swtch[id] < swmin {
            swmin = swtch[id];
        }
    }

    // 计算 CRSW
    let mut crsw = vec![0.0; nd];
    for id in 0..nd {
        crsw[id] = if params.icrsw == 1 {
            // 使用全局最小值
            swmin * params.swpfac
        } else {
            // 使用深度依赖值
            swtch[id] * params.swpfac
        };

        // 限制
        if crsw[id] > params.swplim {
            crsw[id] = UN;
        }
    }

    SwitchOutput {
        crsw,
        swmin: Some(swmin),
    }
}

/// 执行 SWITCH 更新计算（INITM = 0）。
///
/// # 参数
/// - `params`: 输入参数（crsw 会被修改）
///
/// # 返回
/// 更新后的 CRSW 数组
pub fn switch_update(params: &mut SwitchUpdateParams) -> SwitchOutput {
    for id in 0..params.nd {
        params.crsw[id] *= params.swpinc;

        // 限制
        if params.crsw[id] > params.swplim {
            params.crsw[id] = UN;
        }
    }

    SwitchOutput {
        crsw: params.crsw.to_vec(),
        swmin: None,
    }
}

/// 生成 CRSW 输出消息。
pub fn format_crsw_message(crsw: &[f64]) -> String {
    let mut msg = String::new();
    for (i, &val) in crsw.iter().enumerate() {
        msg.push_str(&format!("{:10.3e}", val));
        if (i + 1) % 8 == 0 {
            msg.push('\n');
        }
    }
    if !crsw.is_empty() && crsw.len() % 8 != 0 {
        msg.push('\n');
    }
    msg
}

// ============================================================================
// 测试
// ============================================================================

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

    #[test]
    fn test_switch_disabled() {
        let params = SwitchInitParams {
            nd: 3,
            ntrans: 2,
            colrat: &vec![vec![0.1, 0.1, 0.1], vec![0.1, 0.1, 0.1]],
            rru: &vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]],
            rrd: &vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]],
            fr0: &vec![1.0e15, 2.0e15],
            temp: &vec![10000.0, 9000.0, 8000.0],
            line: &vec![true, false],
            icrsw: 0,  // 禁用
            swpfac: 1.0,
            swplim: 1.0,
        };

        let result = switch_init(&params);

        // 禁用时应该返回全 1
        for &val in &result.crsw {
            assert!((val - 1.0).abs() < 1e-10);
        }
    }

    #[test]
    fn test_switch_global_minimum() {
        let params = SwitchInitParams {
            nd: 3,
            ntrans: 2,
            colrat: &vec![vec![0.1, 0.2, 0.3], vec![0.05, 0.1, 0.15]],
            rru: &vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]],
            rrd: &vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]],
            fr0: &vec![1.0e15, 2.0e15],
            temp: &vec![10000.0, 9000.0, 8000.0],
            line: &vec![false, false],
            icrsw: 1,  // 全局最小
            swpfac: 0.5,
            swplim: 1.0,
        };

        let result = switch_init(&params);

        // 所有深度应该使用相同的全局最小值
        assert!(result.swmin.is_some());
        let swmin = result.swmin.unwrap();
        assert!(swmin < 1.0);

        // CRSW 应该是 swmin * swpfac
        for &val in &result.crsw {
            assert!((val - swmin * 0.5).abs() < 1e-10);
        }
    }

    #[test]
    fn test_switch_depth_dependent() {
        let params = SwitchInitParams {
            nd: 3,
            ntrans: 1,
            colrat: &vec![vec![0.1, 0.2, 0.3]],
            rru: &vec![vec![1.0, 1.0, 1.0]],
            rrd: &vec![vec![1.0, 1.0, 1.0]],
            fr0: &vec![1.0e15],
            temp: &vec![10000.0, 9000.0, 8000.0],
            line: &vec![false],
            icrsw: 2,  // 深度依赖
            swpfac: 1.0,
            swplim: 1.0,
        };

        let result = switch_init(&params);

        // 每个深度应该有不同的值
        assert!((result.crsw[0] - 0.1).abs() < 1e-10);
        assert!((result.crsw[1] - 0.2).abs() < 1e-10);
        assert!((result.crsw[2] - 0.3).abs() < 1e-10);
    }

    #[test]
    fn test_switch_limit() {
        let params = SwitchInitParams {
            nd: 2,
            ntrans: 1,
            colrat: &vec![vec![10.0, 0.5]],  // 第一个深度比例很大
            rru: &vec![vec![1.0, 1.0]],
            rrd: &vec![vec![1.0, 1.0]],
            fr0: &vec![1.0e15],
            temp: &vec![10000.0, 9000.0],
            line: &vec![false],
            icrsw: 2,
            swpfac: 1.0,
            swplim: 0.5,  // 限制为 0.5
        };

        let result = switch_init(&params);

        // 第一个深度应该被限制为 1.0
        assert!((result.crsw[0] - 1.0).abs() < 1e-10);
        // 第二个深度应该正常
        assert!((result.crsw[1] - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_switch_update() {
        let mut crsw = vec![0.1, 0.2, 0.3];

        let mut params = SwitchUpdateParams {
            nd: 3,
            crsw: &mut crsw,
            swpinc: 2.0,
            swplim: 1.0,
        };

        let result = switch_update(&mut params);

        // 应该乘以 swpinc
        assert!((result.crsw[0] - 0.2).abs() < 1e-10);
        assert!((result.crsw[1] - 0.4).abs() < 1e-10);
        assert!((result.crsw[2] - 0.6).abs() < 1e-10);
    }

    #[test]
    fn test_switch_update_limit() {
        let mut crsw = vec![0.8, 0.9];

        let mut params = SwitchUpdateParams {
            nd: 2,
            crsw: &mut crsw,
            swpinc: 2.0,
            swplim: 1.0,
        };

        let result = switch_update(&mut params);

        // 超过限制应该设为 1.0
        assert!((result.crsw[0] - 1.0).abs() < 1e-10);
        assert!((result.crsw[1] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_format_message() {
        let crsw = vec![0.1, 0.2, 0.3];
        let msg = format_crsw_message(&crsw);

        // Rust 的 {:10.3e} 格式产生类似 " 1.000e-1" 的输出
        // 检查包含科学计数法
        assert!(msg.contains("e-1") || msg.contains("E-1"));
        // 检查第一个值
        assert!(msg.contains("1.000e-1") || msg.contains("1.00e-1"));
    }
}