SpectraRust/src/math/compt0.rs

//! Compton 散射源函数辅助量计算。
//!
//! 重构自 TLUSTY `compt0.f`
//!
//! 计算 Compton 散射源函数的辅助量 COMPA, COMPB, COMPC, COMPD, COMPE, COMPS。

use crate::state::constants::{MDEPTH, MFREQ, TWO, UN};

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

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

// ============================================================================
// 辅助数组 (对应 COMMON /auxcbc/)
// ============================================================================

/// Compton 散射导数辅助数组。
/// 对应 COMMON /auxcbc/
#[derive(Debug, Clone, Default)]
pub struct AuxCbc {
    /// CDER1M(MDEPTH) - 频率导数 (前)
    pub cden1m: Vec<f64>,
    /// CDER10(MDEPTH) - 频率导数 (当前)
    pub cden10: Vec<f64>,
    /// CDER2M(MDEPTH) - 二阶频率导数 (前)
    pub cden2m: Vec<f64>,
    /// CDER20(MDEPTH) - 二阶频率导数 (当前)
    pub cden20: Vec<f64>,
}

impl AuxCbc {
    pub fn new() -> Self {
        Self {
            cden1m: vec![0.0; MDEPTH],
            cden10: vec![0.0; MDEPTH],
            cden2m: vec![0.0; MDEPTH],
            cden20: vec![0.0; MDEPTH],
        }
    }
}

// ============================================================================
// COMPT0 参数结构体
// ============================================================================

/// COMPT0 输入参数。
#[derive(Debug, Clone)]
pub struct Compt0Params {
    /// 频率索引 IJ (1-indexed)
    pub ij: usize,
    /// 深度索引 ID (1-indexed)
    pub id: usize,
    /// 吸收系数
    pub ab: f64,

    /// 总频率点数
    pub nfreq: usize,

    // 频率相关数组 [MFREQ]
    pub freq: Vec<f64>,
    pub ijorig: Vec<usize>,
    pub kij: Vec<usize>,
    pub dlnfr: Vec<f64>,
    pub delj: Vec<Vec<f64>>,
    pub bnus: Vec<f64>,
    pub sigec: Vec<f64>,

    // 深度相关数组 [MDEPTH]
    pub temp: Vec<f64>,
    pub elec: Vec<f64>,

    // 辐射场 [MFREQ × MDEPTH]
    pub rad: Vec<Vec<f64>>,

    // Compton 控制参数
    /// Compton 密度导数标志 (0=禁用)
    pub icomde: i32,
    /// Compton 源函数标志 (0=禁用)
    pub icomst: i32,
    /// Compton 高阶项标志
    pub ichcoo: i32,
    /// Compton 模式 (1=标准, 2=无非对角项, 3=全禁用)
    pub icompt: i32,

    /// 汤姆逊散射截面
    pub sige: f64,

    // 可变导数数组 (CDER1M, CDER10, CDER1P, CDER2M, CDER20, CDER2P)
    /// CDER1M(IJI) - 前频率导数
    pub cder1m: Vec<f64>,
    /// CDER10(IJI) - 当前频率导数
    pub cder10: Vec<f64>,
    /// CDER1P(IJI) - 后频率导数
    pub cder1p: Vec<f64>,
    /// CDER2M(IJI) - 前二阶导数
    pub cder2m: Vec<f64>,
    /// CDER20(IJI) - 当前二阶导数
    pub cder20: Vec<f64>,
    /// CDER2P(IJI) - 后二阶导数
    pub cder2p: Vec<f64>,
}

impl Default for Compt0Params {
    fn default() -> Self {
        Self {
            ij: 1,
            id: 1,
            ab: 1.0,
            nfreq: 1,
            freq: vec![0.0; MFREQ],
            ijorig: vec![0; MFREQ],
            kij: vec![0; MFREQ],
            dlnfr: vec![0.0; MFREQ],
            delj: vec![vec![0.0; MDEPTH]; MFREQ],
            bnus: vec![0.0; MFREQ],
            sigec: vec![0.0; MFREQ],
            temp: vec![0.0; MDEPTH],
            elec: vec![0.0; MDEPTH],
            rad: vec![vec![0.0; MDEPTH]; MFREQ],
            icomde: 0,
            icomst: 0,
            ichcoo: 0,
            icompt: 1,
            sige: 6.6516e-25, // 汤姆逊散射截面
            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],
        }
    }
}

/// COMPT0 输出结果。
#[derive(Debug, Clone, Default)]
pub struct Compt0Result {
    /// COMPA - 前频率导数系数
    pub compa: f64,
    /// COMPB - 当前频率导数系数
    pub compb: f64,
    /// COMPC - 后频率导数系数
    pub compc: f64,
    /// COMPD - 密度导数系数
    pub compd: f64,
    /// COMPE - 发射系数
    pub compe: f64,
    /// COMPS - 源函数系数
    pub comps: f64,
}

// ============================================================================
// COMPT0 主函数
// ============================================================================

/// 计算 Compton 散射源函数的辅助量。
///
/// # 参数
///
/// * `params` - 输入参数
///
/// # 返回值
///
/// 返回 Compt0Result 包含 COMPA, COMPB, COMPC, COMPD, COMPE, COMPS
///
/// # Fortran 索引说明
///
/// - IJ, ID, IJI 都是 1-indexed
/// - Rust 中使用 0-indexed，需要转换
pub fn compt0(params: &mut Compt0Params) -> Compt0Result {
    let mut result = Compt0Result::default();

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

    // 如果 IJI = 1，所有输出为 0
    if iji == 1 {
        return result;
    }

    let ij_idx = params.ij - 1;  // 0-indexed
    let id_idx = params.id - 1;  // 0-indexed
    let iji_idx = iji - 1;       // 0-indexed

    // FR = FREQ(IJ)
    let fr = params.freq[ij_idx];
    // FRP = FREQ(IJORIG(IJI+1))
    let frp = params.freq[params.ijorig[iji_idx + 1] - 1];
    // FRM = FREQ(IJORIG(IJI-1))
    let frm = params.freq[params.ijorig[iji_idx - 1] - 1];

    let xcomp = fr * XCON;
    let e2 = YCON * params.temp[id_idx];
    let e1 = xcomp - 3.0 * e2;

    // DEL0 = TWO / (DLNFR(IJI) + DLNFR(IJI-1))
    let del0 = TWO / (params.dlnfr[iji_idx] + params.dlnfr[iji_idx - 1]);

    // 计算导数系数
    // Fortran: CDER1P(IJI) = (UN - DELJ(IJI, ID)) * DEL0
    // Rust: cder1p[iji_idx] = (1.0 - delj[iji_idx][id_idx]) * del0
    params.cder1p[iji_idx] = (UN - params.delj[iji_idx][id_idx]) * del0;
    // Fortran: CDER1M(IJI) = -DELJ(IJI-1, ID) * DEL0
    params.cder1m[iji_idx] = -params.delj[iji_idx - 1][id_idx] * del0;
    // Fortran: CDER10(IJI) = -DEL0 * (UN - DELJ(IJI-1, ID) - DELJ(IJI, ID))
    params.cder10[iji_idx] = -del0 * (UN - params.delj[iji_idx - 1][id_idx] - params.delj[iji_idx][id_idx]);

    // SS0 = ELEC(ID) * SIGE / AB
    let ss0 = params.elec[id_idx] * params.sige / params.ab;

    // 临时变量
    let mut compu = 0.0;
    let mut compv = 0.0;
    let mut cbs = 0.0;

    if params.ichcoo == 0 {
        // 标准模式
        params.cder10[iji_idx] = -params.cder1m[iji_idx] - params.cder1p[iji_idx];

        result.compa = ss0 * (e1 * params.cder1m[iji_idx] + e2 * params.cder2m[iji_idx]);
        result.compb = ss0
            * (UN - xcomp - params.sigec[ij_idx] / params.sige + e1 * params.cder10[iji_idx]
                + e2 * params.cder20[iji_idx]);
        result.compc = ss0 * (e1 * params.cder1p[iji_idx] + e2 * params.cder2p[iji_idx]);
    } else {
        // 高阶模式
        // EPSNU = (AB - ELEC(ID) * SIGEC(IJ)) / AB
        let epsnu = (params.ab - params.elec[id_idx] * params.sigec[ij_idx]) / params.ab;

        // ZXXP, ZXX0, ZXXM
        let zxxp = XCON * frp + 0.5 * params.bnus[iji_idx + 1] * params.rad[iji_idx + 1][id_idx]
            - 3.0 * e2;
        let zxx0 = xcomp + 0.5 * params.bnus[iji_idx] * params.rad[iji_idx][id_idx] - 3.0 * e2;
        let zxxm = XCON * frm + 0.5 * params.bnus[iji_idx - 1] * params.rad[iji_idx - 1][id_idx]
            - 3.0 * e2;

        let zxxp12 = ((UN - params.delj[iji_idx][id_idx]) * zxxp
            + params.delj[iji_idx][id_idx] * zxx0)
            * del0;
        let zxxm12 = ((UN - params.delj[iji_idx - 1][id_idx]) * zxx0
            + params.delj[iji_idx - 1][id_idx] * zxxm)
            * del0;

        result.compa = ss0 * (-params.delj[iji_idx - 1][id_idx] * zxxm12 + e2 * params.cder2m[iji_idx]);
        result.compc = ss0
            * ((UN - params.delj[iji_idx][id_idx]) * zxxp12 + e2 * params.cder2p[iji_idx]);
        result.compb = ss0
            * (params.delj[iji_idx][id_idx] * zxxp12
                - (UN - params.delj[iji_idx - 1][id_idx]) * zxxm12
                + e2 * params.cder20[iji_idx]
                - params.sigec[ij_idx] / params.sige)
            - epsnu
            + 1.0;
        result.compe = 0.0;
    }

    // COMPD = (-3 * CDER10(IJI) + CDER20(IJI)) * RAD(IJI, ID)
    result.compd =
        (-3.0 * params.cder10[iji_idx] + params.cder20[iji_idx]) * params.rad[iji_idx][id_idx];

    // 如果 ICOMDE = 0，禁用密度导数项
    if params.icomde == 0 {
        result.compa = 0.0;
        result.compc = 0.0;
        result.compb = 0.0;
    }

    // X0 = SS0 * BNUS(IJI)
    let x0 = ss0 * params.bnus[iji_idx];

    // 如果 ICOMST = 0，禁用源函数项
    let x0 = if params.icomst == 0 { 0.0 } else { x0 };

    if params.ichcoo == 0 {
        result.compe = x0 * (params.cder10[iji_idx] - UN);
        compu = x0 * params.cder1m[iji_idx];
        compv = x0 * params.cder1p[iji_idx];
        cbs = result.compe * params.rad[iji_idx][id_idx];
        result.compe = cbs;
    }

    result.comps = result.compb * params.rad[iji_idx][id_idx];

    // IJI > 1 时的额外项
    if iji > 1 {
        if params.ichcoo == 0 {
            cbs += compu * params.rad[iji_idx - 1][id_idx];
        }
        result.comps += result.compa * params.rad[iji_idx - 1][id_idx];
        result.compd +=
            (-3.0 * params.cder1m[iji_idx] + params.cder2m[iji_idx]) * params.rad[iji_idx - 1][id_idx];
    }

    // IJI < NFREQ 时的额外项
    if iji < params.nfreq {
        if params.ichcoo == 0 {
            cbs += compv * params.rad[iji_idx + 1][id_idx];
        }
        result.comps += result.compc * params.rad[iji_idx + 1][id_idx];
        result.compd +=
            (-3.0 * params.cder1p[iji_idx] + params.cder2p[iji_idx]) * params.rad[iji_idx + 1][id_idx];
    }

    if params.ichcoo == 0 {
        result.compb += cbs;
        result.compa += compu * params.rad[iji_idx][id_idx];
        result.compc += compv * params.rad[iji_idx][id_idx];
        result.comps += cbs * params.rad[iji_idx][id_idx];
    }

    result.compd *= ss0 * YCON;

    // 如果 ICOMDE = 0，禁用密度导数
    if params.icomde == 0 {
        result.compd = 0.0;
    }

    // ICOMPT = 2 模式：无非对角强度项
    if params.icompt == 2 {
        if iji > 1 {
            result.compb += result.compa * params.rad[iji_idx - 1][id_idx];
        }
        if iji < params.nfreq {
            result.compb += result.compc * params.rad[iji_idx + 1][id_idx];
        }
        result.compa = 0.0;
        result.compc = 0.0;
    } else if params.icompt == 3 {
        // ICOMPT = 3 模式：全禁用
        result.compa = 0.0;
        result.compb = 0.0;
        result.compc = 0.0;
    }

    result
}

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

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

    fn create_test_params() -> Compt0Params {
        let mut params = Compt0Params::default();
        params.ij = 5;
        params.id = 1;
        params.ab = 1e-8;
        params.nfreq = 100;

        // 设置简单的测试值
        params.temp[0] = 10000.0;
        params.elec[0] = 1e12;

        // 设置 KIJ 映射
        for i in 0..100 {
            params.kij[i] = 100 - i; // 简单的反向映射
            params.ijorig[i] = i + 1;
            params.freq[i] = 1e14 * (i + 1) as f64;
            params.dlnfr[i] = 0.1;
            params.bnus[i] = 1e-10;
            params.sigec[i] = 0.0;
            params.delj[i][0] = 0.5;

            // 初始化导数数组
            params.cder1m[i] = 0.0;
            params.cder10[i] = 0.0;
            params.cder1p[i] = 0.0;
            params.cder2m[i] = 0.0;
            params.cder20[i] = 0.0;
            params.cder2p[i] = 0.0;

            // 初始化辐射场
            params.rad[i][0] = 1e10;
        }

        params
    }

    #[test]
    fn test_compt0_iji_equals_1() {
        let mut params = Compt0Params::default();
        // IJI = NFREQ - KIJ(IJ) + 1
        // 当 kij[i] = 100 - i 时:
        // ij = 1 → kij[0] = 100 → iji = 100 - 100 + 1 = 1
        params.ij = 1;
        params.nfreq = 100;
        for i in 0..100 {
            params.kij[i] = 100 - i;
        }

        let result = compt0(&mut params);

        // 当 IJI = 1 时，所有输出应为 0
        assert!((result.compa).abs() < 1e-15);
        assert!((result.compb).abs() < 1e-15);
        assert!((result.compc).abs() < 1e-15);
        assert!((result.compd).abs() < 1e-15);
        assert!((result.compe).abs() < 1e-15);
        assert!((result.comps).abs() < 1e-15);
    }

    #[test]
    fn test_compt0_basic() {
        let mut params = create_test_params();
        params.ichcoo = 0;
        params.icomde = 1;
        params.icomst = 1;
        params.icompt = 1;

        let result = compt0(&mut params);

        // 检查结果有限
        assert!(result.compa.is_finite());
        assert!(result.compb.is_finite());
        assert!(result.compc.is_finite());
        assert!(result.compd.is_finite());
        assert!(result.compe.is_finite());
        assert!(result.comps.is_finite());
    }

    #[test]
    fn test_compt0_icomde_zero() {
        let mut params = create_test_params();
        params.icomde = 0; // 禁用密度导数

        let result = compt0(&mut params);

        // ICOMDE = 0 时，COMPA, COMPB, COMPC, COMPD 应为 0
        assert!((result.compa).abs() < 1e-15);
        assert!((result.compb).abs() < 1e-15);
        assert!((result.compc).abs() < 1e-15);
        assert!((result.compd).abs() < 1e-15);
    }

    #[test]
    fn test_compt0_icompt_3() {
        let mut params = create_test_params();
        params.icomde = 1;
        params.icompt = 3; // 全禁用模式

        let result = compt0(&mut params);

        // ICOMPT = 3 时，COMPA, COMPB, COMPC 应为 0
        assert!((result.compa).abs() < 1e-15);
        assert!((result.compb).abs() < 1e-15);
        assert!((result.compc).abs() < 1e-15);
    }

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

        let result = compt0(&mut params);

        // 检查结果有限
        assert!(result.compa.is_finite());
        assert!(result.compb.is_finite());
        assert!(result.compc.is_finite());
        assert!(result.compd.is_finite());
        assert!(result.compe.is_finite());
        assert!(result.comps.is_finite());
    }

    #[test]
    fn test_constants() {
        // 验证常量
        assert!((XCON - 8.0935e-21).abs() < 1e-30);
        assert!((YCON - 1.68638e-10).abs() < 1e-20);
    }
}