SpectraRust/src/math/comset.rs

//! Compton 散射参数设置。
//!
//! 重构自 TLUSTY `comset.f`。
//!
//! 设置 Compton 散射所需的频率相关参数：
//! - 导数系数 CDER1M, CDER10, CDER1P, CDER2M, CDER20, CDER2P
//! - 插值系数 DELJ
//! - Klein-Nishina 散射截面 SIGEC

use crate::state::constants::{BN, HALF, HK, MDEPTH, MFREQ, SIGE, UN};

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

/// 频率转换常数 XCON = 8.0935e-21
const XCON: f64 = 8.0935e-21;
/// 温度转换常数 YCON = 1.68638e-10
const YCON: f64 = 1.68638e-10;
/// 频率单位转换 T15 = 1e-15
const T15: f64 = 1e-15;

// ============================================================================
// COMSET 参数结构体
// ============================================================================

/// COMSET 输入参数。
#[derive(Debug, Clone)]
pub struct ComsetParams {
    /// 频率点数
    pub nfreq: usize,
    /// 深度点数
    pub nd: usize,

    /// 频率数组 [MFREQ] (1-indexed 访问)
    pub freq: Vec<f64>,
    /// 频率索引映射 KIJ [MFREQ]
    pub kij: Vec<usize>,
    /// 温度数组 [MDEPTH]
    pub temp: Vec<f64>,

    // 控制参数
    /// Compton 模式 (≤0=禁用)
    pub icompt: i32,
    /// Compton 高阶项标志
    pub ichcoo: i32,
    /// Klein-Nishina 截面标志 (0=一阶近似, 1=完整公式)
    pub knish: i32,
}

impl Default for ComsetParams {
    fn default() -> Self {
        Self {
            nfreq: 1,
            nd: 1,
            freq: vec![0.0; MFREQ],
            kij: vec![1; MFREQ],
            temp: vec![0.0; MDEPTH],
            icompt: 0,
            ichcoo: 0,
            knish: 0,
        }
    }
}

/// COMSET 输出结果。
#[derive(Debug, Clone, Default)]
pub struct ComsetResult {
    /// 频率索引映射 IJORIG [MFREQ]
    pub ijorig: Vec<usize>,
    /// 重排后的频率 FREQI [MFREQ]
    pub freqi: Vec<f64>,
    /// Planck 加权因子 BNUS [MFREQ]
    pub bnus: Vec<f64>,
    /// 对数频率间隔 DLNFR [MFREQ]
    pub dlnfr: Vec<f64>,
    /// 插值系数 DELJ [MFREQ × MDEPTH]
    pub delj: Vec<Vec<f64>>,
    /// Klein-Nishina 散射截面 SIGEC [MFREQ]
    pub sigec: Vec<f64>,

    // 导数系数 [MFREQ]
    /// 前频率导数 CDER1M
    pub cder1m: Vec<f64>,
    /// 当前频率导数 CDER10
    pub cder10: Vec<f64>,
    /// 后频率导数 CDER1P
    pub cder1p: Vec<f64>,
    /// 前二阶导数 CDER2M
    pub cder2m: Vec<f64>,
    /// 当前二阶导数 CDER20
    pub cder20: Vec<f64>,
    /// 后二阶导数 CDER2P
    pub cder2p: Vec<f64>,
}

impl ComsetResult {
    pub fn new() -> Self {
        Self {
            ijorig: vec![0; MFREQ],
            freqi: vec![0.0; MFREQ],
            bnus: vec![0.0; MFREQ],
            dlnfr: vec![0.0; MFREQ],
            delj: vec![vec![0.0; MDEPTH]; MFREQ],
            sigec: vec![0.0; MFREQ],
            cder1m: vec![0.0; MFREQ],
            cder10: vec![0.0; MFREQ],
            cder1p: vec![0.0; MFREQ],
            cder2m: vec![0.0; MFREQ],
            cder20: vec![0.0; MFREQ],
            cder2p: vec![0.0; MFREQ],
        }
    }
}

// ============================================================================
// COMSET 主函数
// ============================================================================

/// 设置 Compton 散射参数。
///
/// # 参数
///
/// * `params` - 输入参数
///
/// # 返回值
///
/// 返回 ComsetResult 包含所有计算结果
///
/// # 说明
///
/// 此函数设置 Compton 散射所需的频率相关参数：
/// 1. 重排频率索引 (IJORIG)
/// 2. 计算 Planck 加权因子 (BNUS)
/// 3. 计算对数频率间隔 (DLNFR)
/// 4. 计算二阶导数系数 (CDER2*)
/// 5. 计算插值系数 (DELJ)
/// 6. 计算 Klein-Nishina 散射截面 (SIGEC)
pub fn comset(params: &ComsetParams) -> ComsetResult {
    let mut result = ComsetResult::new();

    // 如果 ICOMPT ≤ 0，跳过大部分计算
    if params.icompt <= 0 {
        // 只计算 SIGEC
        compute_sigec(params, &mut result);
        return result;
    }

    // ========================================================================
    // 频率相关通用参数
    // ========================================================================
    for ij in 1..=params.nfreq {
        // 初始化导数系数
        result.cder10[ij - 1] = 0.0;
        result.cder1p[ij - 1] = 0.0;
        result.cder1m[ij - 1] = 0.0;
        result.cder20[ij - 1] = 0.0;
        result.cder2p[ij - 1] = 0.0;
        result.cder2m[ij - 1] = 0.0;

        // IJI = NFREQ - KIJ(IJ) + 1
        let iji = params.nfreq - params.kij[ij - 1] + 1;

        // IJORIG(IJI) = IJ
        result.ijorig[iji - 1] = ij;

        // FREQI(IJI) = FREQ(IJ)
        result.freqi[iji - 1] = params.freq[ij - 1];

        // FR = FREQI(IJI)
        let fr = result.freqi[iji - 1];

        // BNUS(IJI) = TWO*XCON*FR/(BN*(FR*T15)**3)
        let fr_t15 = fr * T15;
        result.bnus[iji - 1] = 2.0 * XCON * fr / (BN * fr_t15 * fr_t15 * fr_t15);
    }

    // ========================================================================
    // 计算对数频率间隔和二阶导数系数
    // ========================================================================
    // IJ = 1
    result.dlnfr[0] = (result.freqi[1] / result.freqi[0]).ln();

    for ij in 2..params.nfreq {
        // DLNFR(IJ) = LOG(FREQI(IJ+1)/FREQI(IJ))
        result.dlnfr[ij - 1] = (result.freqi[ij] / result.freqi[ij - 1]).ln();

        let delp = result.dlnfr[ij - 1];
        let delm = result.dlnfr[ij - 2];
        let del0 = delp + delm;
        let cd0 = 2.0 / del0;

        // CDER2M(IJ) = CD0/DELM
        result.cder2m[ij - 1] = cd0 / delm;
        // CDER2P(IJ) = CD0/DELP
        result.cder2p[ij - 1] = cd0 / delp;
        // CDER20(IJ) = -CDER2M(IJ) - CDER2P(IJ)
        result.cder20[ij - 1] = -result.cder2m[ij - 1] - result.cder2p[ij - 1];
    }

    // ========================================================================
    // 计算插值系数 DELJ
    // ========================================================================
    for ij in 1..params.nfreq {
        let frj0 = result.freqi[ij - 1];
        let frjp = result.freqi[ij];
        let frz = (frj0 * frjp).sqrt();

        for id in 1..=params.nd {
            // 避免上溢/下溢
            let fjb0 = if HK * frj0 / params.temp[id - 1] < 200.0 {
                UN / ((HK * frj0 / params.temp[id - 1]).exp() - UN)
            } else {
                0.0
            };

            let fjbp = if HK * frjp / params.temp[id - 1] < 200.0 {
                UN / ((HK * frjp / params.temp[id - 1]).exp() - UN)
            } else {
                0.0
            };

            let fjz0 = fjb0 * (BN * (frj0 * T15).powi(3));
            let fjzp = fjbp * (BN * (frjp * T15).powi(3));

            let (aa, bb, cc) = if params.ichcoo == 0 {
                // 标准模式
                let zj0 = HK * frz / params.temp[id - 1];
                let dfjz = fjz0 - fjzp;
                let dfjb = fjb0 - fjbp;
                let fzz = UN + fjbp - 3.0 / zj0;

                let aa = dfjz * dfjb;
                let bb = dfjz * fzz + fjzp * dfjb;
                let cc = fjzp * fzz - dfjz / result.dlnfr[ij - 1] / zj0;

                (aa, bb, cc)
            } else {
                // 高阶模式
                let e2 = YCON * params.temp[id - 1];
                let zxxp = XCON * frjp * (UN + fjbp) - 3.0 * e2;
                let zxx0 = XCON * frj0 * (UN + fjb0) - 3.0 * e2;
                let dzxx = zxx0 - zxxp;
                let dfjb = fjb0 - fjbp;
                let dfjz = fjz0 - fjzp;

                let aa = dfjz * dzxx;
                let bb = dfjz * zxxp + fjzp * dzxx;
                let cc = fjzp * zxxp - e2 * dfjz / result.dlnfr[ij - 1];

                (aa, bb, cc)
            };

            // 求解二次方程 AA*XX1² + BB*XX1 + CC = 0
            let xx1 = if aa.abs() == 0.0 && bb.abs() == 0.0 {
                0.0
            } else if aa.abs() < 1e-7 * bb.abs() {
                -cc / bb
            } else {
                let dd = bb * bb - 4.0 * aa * cc;
                let dd = if dd < 0.0 { 0.0 } else { dd };
                let dd = dd.sqrt();

                let xx1 = (dd - bb) * HALF / aa;

                if params.ichcoo > 0 {
                    let xx2 = -(dd + bb) * HALF / aa;
                    let dxx1 = (xx1 - HALF).abs();
                    let dxx2 = (xx2 - HALF).abs();

                    let xx1 = if dxx2 < dxx1 { xx2 } else { xx1 };
                    if xx1 > 1.0 || xx1 < 0.0 {
                        HALF
                    } else {
                        xx1
                    }
                } else {
                    xx1
                }
            };

            result.delj[ij - 1][id - 1] = xx1;
        }
    }

    // ========================================================================
    // 计算 Klein-Nishina 散射截面 SIGEC
    // ========================================================================
    compute_sigec(params, &mut result);

    result
}

/// 计算 Klein-Nishina 散射截面。
fn compute_sigec(params: &ComsetParams, result: &mut ComsetResult) {
    for ij in 1..=params.nfreq {
        if params.knish == 0 {
            // 一阶近似
            result.sigec[ij - 1] = SIGE * (UN - 2.0 * params.freq[ij - 1] * XCON);
        } else {
            // 完整 Klein-Nishina 截面 (Rybicki & Lightman 1975)
            let xf = XCON * params.freq[ij - 1];

            if xf < 0.1 {
                // 小 x 展开式
                result.sigec[ij - 1] = SIGE
                    * (1.0
                        - xf
                            * (2.0
                                - xf
                                    * (26.0 / 5.0
                                        - xf
                                            * (13.3
                                                - xf
                                                    * (1144.0 / 35.0
                                                        - xf
                                                            * (544.0 / 7.0
                                                                - xf
                                                                    * (3784.0 / 21.0
                                                                        - xf
                                                                            * (6148.0 / 15.0
                                                                                - xf
                                                                                    * (151552.0
                                                                                        / 165.0
                                                                                        - xf
                                                                                            * 111872.0
                                                                                                / 55.0)))))))));
            } else if xf > 1e3 {
                // 大 x 渐近式
                result.sigec[ij - 1] = SIGE * 3.0 / 8.0 / xf * ((2.0 * xf).ln() + 0.5);
            } else {
                // 完整公式
                let xf1 = xf + 1.0;
                let xf2 = 2.0 * xf + 1.0;
                result.sigec[ij - 1] = SIGE
                    * 0.75
                    * (xf1 / (xf * xf * xf) * (2.0 * xf * xf1 / xf2 - (xf2).ln())
                        + 0.5 * (xf2).ln() / xf
                        - (1.0 + 3.0 * xf) / (xf2 * xf2));
            }
        }
    }
}

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

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

    fn create_test_params() -> ComsetParams {
        let mut params = ComsetParams::default();
        params.nfreq = 10;
        params.nd = 5;
        params.icompt = 1;
        params.ichcoo = 0;
        params.knish = 0;

        // 设置频率 (递增)
        for i in 0..10 {
            params.freq[i] = 1e14 * (i + 1) as f64;
            // KIJ 是反向映射
            params.kij[i] = 10 - i;
        }

        // 设置温度
        for i in 0..5 {
            params.temp[i] = 10000.0 + i as f64 * 1000.0;
        }

        params
    }

    #[test]
    fn test_comset_icompt_zero() {
        let mut params = ComsetParams::default();
        params.nfreq = 5;
        params.icompt = 0; // 禁用
        params.knish = 0;

        for i in 0..5 {
            params.freq[i] = 1e14 * (i + 1) as f64;
            params.kij[i] = 5 - i;
        }

        let result = comset(&params);

        // 当 ICOMPT = 0 时，只计算 SIGEC
        // 检查 SIGEC 有限
        for i in 0..params.nfreq {
            assert!(result.sigec[i].is_finite());
        }
    }

    #[test]
    fn test_comset_basic() {
        let params = create_test_params();
        let result = comset(&params);

        // 检查 IJORIG 映射 (只检查前 nfreq 个元素)
        for i in 0..params.nfreq {
            assert!(result.ijorig[i] >= 1 && result.ijorig[i] <= params.nfreq);
        }

        // 检查 BNUS 有限
        for i in 0..params.nfreq {
            assert!(result.bnus[i].is_finite());
        }

        // 检查 DLNFR 正数
        for i in 0..params.nfreq - 1 {
            assert!(result.dlnfr[i] > 0.0);
        }

        // 检查 DELJ 在 [0, 1] 范围内
        for ij in 0..params.nfreq - 1 {
            for id in 0..params.nd {
                assert!(result.delj[ij][id] >= 0.0 && result.delj[ij][id] <= 1.0);
            }
        }

        // 检查 SIGEC 有限
        for i in 0..params.nfreq {
            assert!(result.sigec[i].is_finite());
        }
    }

    #[test]
    fn test_comset_high_order_mode() {
        let mut params = create_test_params();
        params.ichcoo = 1; // 高阶模式

        let result = comset(&params);

        // 检查结果有限
        for i in 0..params.nfreq {
            assert!(result.bnus[i].is_finite());
            assert!(result.sigec[i].is_finite());
        }
    }

    #[test]
    fn test_comset_klein_nishina() {
        let mut params = create_test_params();
        params.knish = 1; // 完整 Klein-Nishina

        let result = comset(&params);

        // 检查 SIGEC 有限且为正
        for i in 0..params.nfreq {
            assert!(result.sigec[i].is_finite());
            assert!(result.sigec[i] > 0.0);
        }
    }

    #[test]
    fn test_comset_sigec_limits() {
        let mut params = ComsetParams::default();
        params.nfreq = 3;
        params.knish = 1;

        // 低频 (xf < 0.1)
        params.freq[0] = 1e12; // xf ≈ 8e-9
        // 中频 (0.1 < xf < 1000)
        params.freq[1] = 1e16; // xf ≈ 0.08
        // 高频 (xf > 1000)
        params.freq[2] = 1e24; // xf ≈ 8e3

        let result = comset(&params);

        // 所有 SIGEC 应为正且有限
        for i in 0..3 {
            assert!(result.sigec[i] > 0.0);
            assert!(result.sigec[i].is_finite());
        }
    }

    #[test]
    fn test_constants() {
        assert!((XCON - 8.0935e-21).abs() < 1e-30);
        assert!((YCON - 1.68638e-10).abs() < 1e-20);
        assert!((T15 - 1e-15).abs() < 1e-20);
    }
}