SpectraRust/src/math/alifrk.rs

//! ALI 频率相关计算 - Kantorovich 迭代简化版本。
//!
//! 重构自 TLUSTY `alifrk.f`
//!
//! 这是 ALIFR1 的简化版本，用于 Kantorovich 迭代。
//! 计算流体静力学和辐射平衡量 - ALI 点的总加热和冷却率对
//! 温度、电子密度和占据数的导数。

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

/// ALIFRK 输入参数
pub struct AlifrkParams {
    /// 频率索引 (1-indexed)
    pub ij: usize,
    /// 深度点数
    pub nd: usize,
    /// ALI 模式
    pub ifali: i32,
}

/// ALIFRK 需要的模型状态
pub struct AlifrkState<'a> {
    // 深度相关 (MDEPTH)
    /// 深度差分 DELDMZ
    pub deldmz: &'a [f64],
    /// 总吸收 ABSOT
    pub absot: &'a [f64],

    // 辐射相关 (MDEPTH)
    /// 辐射强度 RAD1
    pub rad1: &'a [f64],
    /// 吸收系数 × 辐射 FAK1
    pub fak1: &'a [f64],

    // 平衡相关
    /// 辐射等效积分 REINT
    pub reint: &'a [f64],

    // 频率相关
    /// 外辐射 EXTRAD
    pub extrad: &'a [f64],
    /// 氦外辐射 HEXTRD
    pub hextrd: &'a [f64],
    /// 频率权重 FH
    pub fh: &'a [f64],
    /// 频率权重 W
    pub w: &'a [f64],

    // 不透明度 (MDEPTH)
    /// 吸收系数 ABSO1
    pub abso1: &'a [f64],
    /// 发射系数 EMIS1
    pub emis1: &'a [f64],
    /// 散射系数 SCAT1
    pub scat1: &'a [f64],

    // 频率权重 WC
    pub wc: &'a [f64],

    // 跳过标志 (MDEPTH × MFREQ)
    /// LSKIP(ID, IJ) - .TRUE. 表示跳过该频率-深度点
    pub lskip: &'a [Vec<i32>],

    // 输出累积变量 (MDEPTH)
    /// 辐射压力导数 FPRD
    pub fprd: &'a mut [f64],
    /// 固定辐射通量 FLFIX
    pub flfix: &'a mut [f64],
    /// 辐射通量红翼 FLRD
    pub flrd: &'a mut [f64],
    /// 冷却率积分 FCOOLI
    pub fcooli: &'a mut [f64],
}

/// ALI 频率相关计算 - Kantorovich 迭代简化版本。
///
/// 这是 ALIFR1 的简化版本，计算流体静力学和辐射平衡量。
///
/// # 参数
///
/// * `params` - 输入参数
/// * `state` - 模型状态 (包含输入和输出)
///
/// # Fortran 索引说明
///
/// - IJ 是频率索引 (1-indexed)
/// - ID 是深度索引 (1-indexed)
///
/// # 算法说明
///
/// 1. 如果 IFALI <= 1，直接返回
/// 2. 对第一个深度点 (ID=1) 特殊处理
/// 3. 对 ID=2 到 ND 循环计算
pub fn alifrk(params: &AlifrkParams, state: &mut AlifrkState) {
    // 如果 IFALI <= 1，直接返回
    if params.ifali <= 1 {
        return;
    }

    let ij = params.ij;
    let nd = params.nd;

    // Fortran 1-indexed 转 Rust 0-indexed
    let ij_idx = ij - 1;

    // 权重因子
    let ww = state.wc[ij_idx];

    // 工作数组 WFL(MDEPTH) - 用于存储中间结果
    let mut wfl = vec![0.0; MDEPTH];

    // ========================================================================
    // 第一个深度点 (ID=1) 的特殊处理
    // ========================================================================
    let id = 1;
    let id_idx = id - 1; // 0

    // LNSKIP = .NOT. LSKIP(ID, IJ)
    // Fortran LSKIP(ID, IJ) 是 true 表示跳过
    // Rust lskip[depth][freq]，如果 lskip[id_idx][ij_idx] != 0 表示跳过
    let lnskip = state.lskip[id_idx][ij_idx] == 0;

    // WF = WW * (FH(IJ) * RAD1(ID) - HEXTRD(IJ))
    let wf = ww * (state.fh[ij_idx] * state.rad1[id_idx] - state.hextrd[ij_idx]);

    // IF(LNSKIP) FPRD(ID) = FPRD(ID) + WF * ABSO1(ID)
    if lnskip {
        state.fprd[id_idx] += wf * state.abso1[id_idx];
    }

    // FLFIX(ID) = FLFIX(ID) + WF
    state.flfix[id_idx] += wf;

    // FLRD(ID) = FLRD(ID) + W(IJ) * (FH(IJ) * RAD1(ID) - HALF * EXTRAD(IJ))
    state.flrd[id_idx] += state.w[ij_idx]
        * (state.fh[ij_idx] * state.rad1[id_idx] - HALF * state.extrad[ij_idx]);

    // IF(REINT(ID).GT.0) THEN
    if state.reint[id_idx] > 0.0 {
        // ABST = ABSO1(ID) - SCAT1(ID)
        let abst = state.abso1[id_idx] - state.scat1[id_idx];
        // FCOOLI(ID) = FCOOLI(ID) + WW * (EMIS1(ID) - ABST * RAD1(ID))
        state.fcooli[id_idx] += ww * (state.emis1[id_idx] - abst * state.rad1[id_idx]);
    }

    // ========================================================================
    // 循环处理 ID = 2 到 ND
    // ========================================================================
    for id in 2..=nd {
        let id_idx = id - 1;
        let id_prev_idx = id - 2; // ID - 1 in 0-indexed

        // LNSKIP = .NOT. LSKIP(ID, IJ)
        let lnskip = state.lskip[id_idx][ij_idx] == 0;

        // DT = UN / ((ABSOT(ID) + ABSOT(ID-1)) * DELDMZ(ID-1))
        // 注意：Fortran DELDMZ(ID-1) 对应 Rust deldmz[id_prev_idx]
        let dt = UN / ((state.absot[id_idx] + state.absot[id_prev_idx]) * state.deldmz[id_prev_idx]);

        // FL = RAD1(ID) * FAK1(ID) - RAD1(ID-1) * FAK1(ID-1)
        let fl = state.rad1[id_idx] * state.fak1[id_idx]
            - state.rad1[id_prev_idx] * state.fak1[id_prev_idx];

        // WFL(ID) = WW * FL
        wfl[id_idx] = ww * fl;

        // IF(LNSKIP) FPRD(ID) = FPRD(ID) + WFL(ID)
        if lnskip {
            state.fprd[id_idx] += wfl[id_idx];
        }

        // FLFIX(ID) = FLFIX(ID) + WFL(ID) * DT
        state.flfix[id_idx] += wfl[id_idx] * dt;

        // FLRD(ID) = FLRD(ID) + FL * W(IJ) * DT
        state.flrd[id_idx] += fl * state.w[ij_idx] * dt;

        // IF(REINT(ID).GT.0) THEN
        if state.reint[id_idx] > 0.0 {
            // ABST = ABSO1(ID) - SCAT1(ID)
            let abst = state.abso1[id_idx] - state.scat1[id_idx];
            // FCOOLI(ID) = FCOOLI(ID) + WW * (EMIS1(ID) - ABST * RAD1(ID))
            state.fcooli[id_idx] += ww * (state.emis1[id_idx] - abst * state.rad1[id_idx]);
        }
    }
}

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

    fn create_test_state() -> (
        AlifrkParams,
        Vec<f64>,  // deldmz
        Vec<f64>,  // absot
        Vec<f64>,  // rad1
        Vec<f64>,  // fak1
        Vec<f64>,  // reint
        Vec<f64>,  // extrad
        Vec<f64>,  // hextrd
        Vec<f64>,  // fh
        Vec<f64>,  // w
        Vec<f64>,  // abso1
        Vec<f64>,  // emis1
        Vec<f64>,  // scat1
        Vec<f64>,  // wc
        Vec<Vec<i32>>,  // lskip
        Vec<f64>,  // fprd
        Vec<f64>,  // flfix
        Vec<f64>,  // flrd
        Vec<f64>,  // fcooli
    ) {
        let nd = 5;
        let params = AlifrkParams {
            ij: 1,
            nd,
            ifali: 2,
        };

        // 初始化数组
        let deldmz = vec![0.1; MDEPTH];
        let absot = vec![1.0; MDEPTH];
        let rad1 = vec![0.5; MDEPTH];
        let fak1 = vec![1.0; MDEPTH];
        let reint = vec![1.0; MDEPTH];
        let extrad = vec![0.1; MFREQ];
        let hextrd = vec![0.05; MFREQ];
        let fh = vec![1.0; MFREQ];
        let w = vec![1.0; MFREQ];
        let abso1 = vec![2.0; MDEPTH];
        let emis1 = vec![1.0; MDEPTH];
        let scat1 = vec![0.5; MDEPTH];
        let wc = vec![1.0; MFREQ];
        let lskip = vec![vec![0; MFREQ]; MDEPTH]; // 全部不跳过

        // 输出数组
        let fprd = vec![0.0; MDEPTH];
        let flfix = vec![0.0; MDEPTH];
        let flrd = vec![0.0; MDEPTH];
        let fcooli = vec![0.0; MDEPTH];

        (
            params, deldmz, absot, rad1, fak1, reint, extrad, hextrd, fh, w, abso1, emis1, scat1,
            wc, lskip, fprd, flfix, flrd, fcooli,
        )
    }

    #[test]
    fn test_alifrk_ifali_le_1() {
        let (
            _,
            deldmz,
            absot,
            rad1,
            fak1,
            reint,
            extrad,
            hextrd,
            fh,
            w,
            abso1,
            emis1,
            scat1,
            wc,
            lskip,
            mut fprd,
            mut flfix,
            mut flrd,
            mut fcooli,
        ) = create_test_state();

        let params = AlifrkParams {
            ij: 1,
            nd: 5,
            ifali: 1, // <= 1，应该直接返回
        };

        let mut state = AlifrkState {
            deldmz: &deldmz,
            absot: &absot,
            rad1: &rad1,
            fak1: &fak1,
            reint: &reint,
            extrad: &extrad,
            hextrd: &hextrd,
            fh: &fh,
            w: &w,
            abso1: &abso1,
            emis1: &emis1,
            scat1: &scat1,
            wc: &wc,
            lskip: &lskip,
            fprd: &mut fprd,
            flfix: &mut flfix,
            flrd: &mut flrd,
            fcooli: &mut fcooli,
        };

        alifrk(&params, &mut state);

        // 所有输出应该保持为 0
        for i in 0..5 {
            assert_eq!(state.fprd[i], 0.0);
            assert_eq!(state.flfix[i], 0.0);
            assert_eq!(state.flrd[i], 0.0);
            assert_eq!(state.fcooli[i], 0.0);
        }
    }

    #[test]
    fn test_alifrk_basic() {
        let (
            params,
            deldmz,
            absot,
            rad1,
            fak1,
            reint,
            extrad,
            hextrd,
            fh,
            w,
            abso1,
            emis1,
            scat1,
            wc,
            lskip,
            mut fprd,
            mut flfix,
            mut flrd,
            mut fcooli,
        ) = create_test_state();

        let mut state = AlifrkState {
            deldmz: &deldmz,
            absot: &absot,
            rad1: &rad1,
            fak1: &fak1,
            reint: &reint,
            extrad: &extrad,
            hextrd: &hextrd,
            fh: &fh,
            w: &w,
            abso1: &abso1,
            emis1: &emis1,
            scat1: &scat1,
            wc: &wc,
            lskip: &lskip,
            fprd: &mut fprd,
            flfix: &mut flfix,
            flrd: &mut flrd,
            fcooli: &mut fcooli,
        };

        alifrk(&params, &mut state);

        // 验证 ID=1 的计算
        // wf = 1.0 * (1.0 * 0.5 - 0.05) = 0.45
        // fprd[0] += 0.45 * 2.0 = 0.9
        // flfix[0] += 0.45
        // flrd[0] += 1.0 * (1.0 * 0.5 - 0.5 * 0.1) = 0.45
        // abst = 2.0 - 0.5 = 1.5
        // fcooli[0] += 1.0 * (1.0 - 1.5 * 0.5) = 0.25
        assert!((state.fprd[0] - 0.9).abs() < 1e-10);
        assert!((state.flfix[0] - 0.45).abs() < 1e-10);
        assert!((state.flrd[0] - 0.45).abs() < 1e-10);
        assert!((state.fcooli[0] - 0.25).abs() < 1e-10);

        // 验证 ID=2 的计算
        // dt = 1.0 / ((1.0 + 1.0) * 0.1) = 5.0
        // fl = 0.5 * 1.0 - 0.5 * 1.0 = 0.0
        // wfl[1] = 1.0 * 0.0 = 0.0
        // flfix[1] += 0.0 * 5.0 = 0.0
        // flrd[1] += 0.0 * 1.0 * 5.0 = 0.0
        // fcooli[1] += 1.0 * (1.0 - 1.5 * 0.5) = 0.25
        assert!((state.flfix[1]).abs() < 1e-10);
        assert!((state.flrd[1]).abs() < 1e-10);
        assert!((state.fcooli[1] - 0.25).abs() < 1e-10);
    }

    #[test]
    fn test_alifrk_skip() {
        let (
            params,
            deldmz,
            absot,
            rad1,
            fak1,
            reint,
            extrad,
            hextrd,
            fh,
            w,
            abso1,
            emis1,
            scat1,
            wc,
            mut lskip,
            mut fprd,
            mut flfix,
            mut flrd,
            mut fcooli,
        ) = create_test_state();

        // 设置 lskip[0][0] = 1，表示跳过 ID=1, IJ=1
        lskip[0][0] = 1;

        let mut state = AlifrkState {
            deldmz: &deldmz,
            absot: &absot,
            rad1: &rad1,
            fak1: &fak1,
            reint: &reint,
            extrad: &extrad,
            hextrd: &hextrd,
            fh: &fh,
            w: &w,
            abso1: &abso1,
            emis1: &emis1,
            scat1: &scat1,
            wc: &wc,
            lskip: &lskip,
            fprd: &mut fprd,
            flfix: &mut flfix,
            flrd: &mut flrd,
            fcooli: &mut fcooli,
        };

        alifrk(&params, &mut state);

        // fprd[0] 不应该被更新，因为 lskip[0][0] = 1
        // 但 flfix[0], flrd[0], fcooli[0] 仍然更新
        assert!((state.fprd[0]).abs() < 1e-10); // 跳过，不更新
        assert!((state.flfix[0] - 0.45).abs() < 1e-10);
        assert!((state.flrd[0] - 0.45).abs() < 1e-10);
        assert!((state.fcooli[0] - 0.25).abs() < 1e-10);
    }

    #[test]
    fn test_alifrk_reint_zero() {
        let (
            params,
            deldmz,
            absot,
            rad1,
            fak1,
            mut reint,
            extrad,
            hextrd,
            fh,
            w,
            abso1,
            emis1,
            scat1,
            wc,
            lskip,
            mut fprd,
            mut flfix,
            mut flrd,
            mut fcooli,
        ) = create_test_state();

        // 设置 reint[0] = 0，fcooli[0] 不应该更新
        reint[0] = 0.0;

        let mut state = AlifrkState {
            deldmz: &deldmz,
            absot: &absot,
            rad1: &rad1,
            fak1: &fak1,
            reint: &reint,
            extrad: &extrad,
            hextrd: &hextrd,
            fh: &fh,
            w: &w,
            abso1: &abso1,
            emis1: &emis1,
            scat1: &scat1,
            wc: &wc,
            lskip: &lskip,
            fprd: &mut fprd,
            flfix: &mut flfix,
            flrd: &mut flrd,
            fcooli: &mut fcooli,
        };

        alifrk(&params, &mut state);

        // fcooli[0] 不应该更新
        assert!((state.fcooli[0]).abs() < 1e-10);
        // fcooli[1] 应该更新
        assert!((state.fcooli[1] - 0.25).abs() < 1e-10);
    }
}