SpectraRust/src/math/meanopt.rs

//! Rosseland 和 Planck 平均不透明度计算（使用 OPCTAB）。
//!
//! 重构自 TLUSTY `meanopt.f`
//!
//! 通过调用 OPCTAB 动态计算每个频率点的吸收和散射系数，
//! 然后计算 Rosseland 和 Planck 平均不透明度。

use super::opctab::{opctab, OpctabOutput};
use crate::state::constants::{HK, UN};

/// MEANOPT 输入参数
pub struct MeanoptParams<'a> {
    /// 温度 (K)
    pub t: f64,
    /// 深度索引 (1-indexed)
    pub id: usize,
    /// 密度 (g cm^-3)
    pub rho: f64,
    /// 迭代次数
    pub iter: i32,
    /// Rayleigh 散射标志 (<0: 简单公式, >0: 调用 rayleigh, =0: 关闭)
    pub ifrayl: i32,
    /// 选项表标志 (<0: 添加电子散射)
    pub ioptab: i32,
    /// Rayleigh 参数 (当 ifrayl > 0 时需要)
    pub rayleigh_params: Option<&'a super::rayleigh::RayleighParams<'a>>,
}

/// MEANOPT 模型状态
pub struct MeanoptModelState<'a> {
    /// 频率数组 (nfreq)
    pub freq: &'a [f64],
    /// Planck 函数 (nfreq)
    pub bnue: &'a [f64],
    /// 权重 (nfreq)
    pub w: &'a [f64],
}

/// MEANOPT 输出
pub struct MeanoptOutput {
    /// Rosseland 平均不透明度 (per gram)
    pub opros: f64,
    /// Planck 平均不透明度 (per gram)
    pub oppla: f64,
}

/// 计算 Rosseland 和 Planck 平均不透明度。
///
/// 通过调用 OPCTAB 为每个频率点动态计算吸收和散射系数。
///
/// # 参数
///
/// * `params` - 输入参数
/// * `model` - 模型状态（频率、Planck 函数、权重）
/// * `opctab_params_base` - OPCTAB 基础参数（不包括 fr, ij）
/// * `opctab_table` - 不透明度表数据
/// * `opctab_model` - OPCTAB 模型状态
///
/// # 返回值
///
/// 包含 Rosseland 和 Planck 平均不透明度的结构体
pub fn meanopt(
    params: &MeanoptParams,
    model: &MeanoptModelState,
    opctab_table: &super::opctab::OpctabTableData,
    opctab_model: &mut super::opctab::OpctabModelState,
) -> MeanoptOutput {
    let t = params.t;
    let hkt = HK / t;

    let mut abr = 0.0;
    let mut sumdb = 0.0;
    let mut abp = 0.0;
    let mut sumb = 0.0;

    let nfreq = model.freq.len();

    for ij in 0..nfreq {
        let fr = model.freq[ij];
        let ex = (hkt * fr).exp();
        let e1 = UN / (ex - UN);
        let plan = model.bnue[ij] * e1 * model.w[ij];
        let dplan = plan * hkt * fr * ex * e1;

        // 调用 OPCTAB 获取吸收和散射系数
        let opctab_params = super::opctab::OpctabParams {
            fr,
            ij: ij + 1, // 1-indexed
            id: params.id,
            t,
            rho: params.rho,
            igram: 1, // 返回每克
            iter: params.iter,
            ifrayl: params.ifrayl,
            ioptab: params.ioptab,
            rayleigh_params: params.rayleigh_params,
        };

        let OpctabOutput { ab, sct, .. } = opctab(&opctab_params, opctab_table, opctab_model);

        abr += dplan / (ab + sct);
        abp += plan * ab;
        sumdb += dplan;
        sumb += plan;
    }

    let opros = sumdb / abr;
    let oppla = abp / sumb;

    MeanoptOutput { opros, oppla }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::opctab::{OpctabTableData, OpctabModelState};
    use approx::assert_relative_eq;

    #[test]
    fn test_meanopt_function_basic() {
        let numtemp = 1;
        let nd = 1;
        let nfreq = 3;
        let max_numrh = 1;

        let tempvec = vec![9.2103];
        let numrh = vec![1];
        let rhomat = vec![-16.1181];
        let absopac = vec![0.0, 0.0, 0.0];
        let raysc = vec![0.0];
        let freq = vec![1e14, 3e14, 1e15];

        let table = OpctabTableData {
            numtemp,
            nd,
            nfreq,
            max_numrh,
            tempvec: &tempvec,
            ttab1: 8.0,
            ttab2: 11.0,
            numrh: &numrh,
            rhomat: &rhomat,
            absopac: &absopac,
            raysc: &raysc,
            freq: &freq,
            sige: 0.0,
        };

        let t = 10000.0;
        let rho = 1e-7;

        // 简化的 Planck 函数
        let bnue: Vec<f64> = freq.iter().map(|f| 1.4743e-2 * f.powi(3)).collect();
        let w: Vec<f64> = vec![0.3, 0.4, 0.3]; // 权重

        let params = MeanoptParams {
            t,
            id: 1,
            rho,
            iter: 1,
            ifrayl: 0,
            ioptab: 0,
            rayleigh_params: None,
        };

        let model = MeanoptModelState {
            freq: &freq,
            bnue: &bnue,
            w: &w,
        };

        let elec = vec![1e-10];
        let dens = vec![rho];
        let mut opctab_model = OpctabModelState {
            elec: &elec,
            dens: &dens,
            raysct: None,
            eospar: None,
        };

        // 调用函数
        let result = meanopt(&params, &model, &table, &mut opctab_model);

        // 验证输出
        assert!(result.opros > 0.0, "Rosseland mean should be positive");
        assert!(result.oppla > 0.0, "Planck mean should be positive");
        assert!(result.opros.is_finite(), "Rosseland mean should be finite");
        assert!(result.oppla.is_finite(), "Planck mean should be finite");
    }

    #[test]
    fn test_meanopt_function_uniform_opacity() {
        let numtemp = 1;
        let nd = 1;
        let nfreq = 3;
        let max_numrh = 1;

        let tempvec = vec![9.2103];
        let numrh = vec![1];
        let rhomat = vec![-16.1181];
        let absopac = vec![0.0, 0.0, 0.0];
        let raysc = vec![0.0];
        let freq = vec![1e14, 3e14, 1e15];

        let table = OpctabTableData {
            numtemp,
            nd,
            nfreq,
            max_numrh,
            tempvec: &tempvec,
            ttab1: 8.0,
            ttab2: 11.0,
            numrh: &numrh,
            rhomat: &rhomat,
            absopac: &absopac,
            raysc: &raysc,
            freq: &freq,
            sige: 0.0,
        };

        let t = 10000.0;
        let rho = 1e-7;

        // 使用相同的权重
        let bnue: Vec<f64> = vec![1.0, 1.0, 1.0];
        let w: Vec<f64> = vec![1.0/3.0, 1.0/3.0, 1.0/3.0];

        let params = MeanoptParams {
            t,
            id: 1,
            rho,
            iter: 1,
            ifrayl: 0,
            ioptab: 0,
            rayleigh_params: None,
        };

        let model = MeanoptModelState {
            freq: &freq,
            bnue: &bnue,
            w: &w,
        };

        let elec = vec![1e-10];
        let dens = vec![rho];
        let mut opctab_model = OpctabModelState {
            elec: &elec,
            dens: &dens,
            raysct: None,
            eospar: None,
        };

        let result = meanopt(&params, &model, &table, &mut opctab_model);

        // 验证结果合理
        assert!(result.opros > 0.0 && result.opros < 10.0,
            "Rosseland mean should be reasonable: got {}", result.opros);
        assert!(result.oppla > 0.0 && result.oppla < 10.0,
            "Planck mean should be reasonable: got {}", result.oppla);
    }

    #[test]
    fn test_meanopt_function_temperature_dependence() {
        let numtemp = 1;
        let nd = 1;
        let nfreq = 3;
        let max_numrh = 1;

        let tempvec = vec![9.2103];
        let numrh = vec![1];
        let rhomat = vec![-16.1181];
        let absopac = vec![0.0, 0.0, 0.0];
        let raysc = vec![0.0];
        let freq = vec![1e14, 3e14, 1e15];

        let table = OpctabTableData {
            numtemp,
            nd,
            nfreq,
            max_numrh,
            tempvec: &tempvec,
            ttab1: 8.0,
            ttab2: 11.0,
            numrh: &numrh,
            rhomat: &rhomat,
            absopac: &absopac,
            raysc: &raysc,
            freq: &freq,
            sige: 0.0,
        };

        let bnue: Vec<f64> = freq.iter().map(|f| 1.4743e-2 * f.powi(3)).collect();
        let w: Vec<f64> = vec![0.3, 0.4, 0.3];

        let rho = 1e-7;
        let elec = vec![1e-10];
        let dens = vec![rho];

        // 测试不同温度
        let temperatures = [5000.0, 10000.0, 20000.0];
        let mut results = Vec::new();

        for t in &temperatures {
            let params = MeanoptParams {
                t: *t,
                id: 1,
                rho,
                iter: 1,
                ifrayl: 0,
                ioptab: 0,
                rayleigh_params: None,
            };

            let model = MeanoptModelState {
                freq: &freq,
                bnue: &bnue,
                w: &w,
            };

            let mut opctab_model = OpctabModelState {
                elec: &elec,
                dens: &dens,
                raysct: None,
                eospar: None,
            };

            results.push(meanopt(&params, &model, &table, &mut opctab_model));
        }

        // 所有结果应该是有效的
        for (i, r) in results.iter().enumerate() {
            assert!(r.opros > 0.0 && r.opros.is_finite(),
                "Rosseland mean at T={} should be valid", temperatures[i]);
            assert!(r.oppla > 0.0 && r.oppla.is_finite(),
                "Planck mean at T={} should be valid", temperatures[i]);
        }
    }
}