SpectraRust/src/math/crosew.rs

//! 光致电离截面数组设置。
//!
//! 重构自 SYNSPEC `CROSET` 和 `CROSEW` 函数。
//!
//! # 功能
//!
//! 设置光致电离截面数组，用于辐射转移计算。

use super::sigk::{sigk, SigkParams};
use crate::state::atomic::AtomicData;
use crate::state::constants::{MCROSS, MFREQ};

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

/// INDEXP 值表示溶解能级（需要特殊处理）
const INDEXP_DISSOLVED: i32 = 5;

// ============================================================================
// CROSET - 使用 FREQ 数组设置截面
// ============================================================================

/// CROSET 输入参数。
pub struct CrosetParams<'a> {
    /// 频率数组 (FREQ)
    pub freq: &'a [f64],
    /// 能级数 (NLEVEL)
    pub nlevel: usize,
    /// 频率数 (NFREQ)
    pub nfreq: usize,
    /// 模式 (IMODE)
    pub imode: i32,
    /// INDEXP 数组 - 能级索引类型
    pub indexp: &'a [i32],
    /// FROPC 数组 - 阈值频率
    pub fropc: &'a [f64],
    /// 原子数据引用
    pub atomic: &'a AtomicData,
}

/// 设置光致电离截面数组 (使用 FREQ)。
///
/// # 参数
///
/// * `params` - 输入参数
///
/// # 返回值
///
/// CROSS 数组 [MCROSS][MFREQ]，其中 CROSS[itr][ij] 是能级 itr 在频率 ij 处的截面
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE CROSET(CROSS)
///   IJ0=2
///   IF(NFREQ.EQ.1) IJ0=1
///   IF(IMODE.EQ.2) IJ0=NFREQ
///   DO IJ=1,IJ0
///      DO IT=1,MCROSS
///         CROSS(IT,IJ)=0.
///      END DO
///   END DO
///   DO IT=1,NLEVEL
///      IF(INDEXP(IT).NE.5) THEN
///         DO IJ=1,IJ0
///            FR=FREQ(IJ)
///            CROSS(IT,IJ)=SIGK(FR,IT,0)
///         END DO
///      ELSE
///         DO IJ=1,IJ0
///            FR=FREQ(IJ)
///            CROSS(IT,IJ)=SIGK(FR,IT,1)
///            IF(FR.LT.FROPC(IT)) CROSS(IT,IJ)=0.
///         END DO
///      END IF
///   END DO
/// END
/// ```
pub fn croset(params: &CrosetParams) -> Vec<Vec<f64>> {
    let CrosetParams {
        freq,
        nlevel,
        nfreq,
        imode,
        indexp,
        fropc,
        atomic,
    } = *params;

    // 确定频率范围
    // Fortran: IJ0=2; IF(NFREQ.EQ.1) IJ0=1; IF(IMODE.EQ.2) IJ0=NFREQ
    let ij0 = if nfreq == 1 {
        1
    } else if imode == 2 {
        nfreq
    } else {
        2
    };

    // 初始化截面数组
    // 注意：Fortran 是 1-indexed，Rust 是 0-indexed
    let mut cross = vec![vec![0.0; ij0]; nlevel];

    // 计算每个能级在每个频率处的截面
    for it in 0..nlevel {
        let idxp = indexp[it];

        if idxp != INDEXP_DISSOLVED {
            // 普通能级：mode = 0（边缘长波方向截面为零）
            for ij in 0..ij0 {
                let fr = freq[ij];
                let sigk_params = SigkParams {
                    fr,
                    itr: it,
                    mode: 0,
                    atomic,
                    opdata: &super::topbas::OpData::default(),
                };
                cross[it][ij] = sigk(&sigk_params);
            }
        } else {
            // 溶解能级：mode = 1（边缘长波方向截面非零）
            for ij in 0..ij0 {
                let fr = freq[ij];
                let sigk_params = SigkParams {
                    fr,
                    itr: it,
                    mode: 1,
                    atomic,
                    opdata: &super::topbas::OpData::default(),
                };
                cross[it][ij] = sigk(&sigk_params);

                // 如果频率低于阈值，截面设为零
                // Fortran: IF(FR.LT.FROPC(IT)) CROSS(IT,IJ)=0.
                if fr < fropc[it] {
                    cross[it][ij] = 0.0;
                }
            }
        }
    }

    cross
}

// ============================================================================
// CROSEW - 使用 FREQC 数组设置截面
// ============================================================================

/// CROSEW 输入参数。
pub struct CrosewParams<'a> {
    /// 频率数组 (FREQC)
    pub freqc: &'a [f64],
    /// 能级数 (NLEVEL)
    pub nlevel: usize,
    /// 频率数 (NFREQC)
    pub nfreqc: usize,
    /// INDEXP 数组 - 能级索引类型
    pub indexp: &'a [i32],
    /// FROPC 数组 - 阈值频率
    pub fropc: &'a [f64],
    /// 原子数据引用
    pub atomic: &'a AtomicData,
}

/// 设置光致电离截面数组 (使用 FREQC)。
///
/// # 参数
///
/// * `params` - 输入参数
///
/// # 返回值
///
/// CROSS 数组 [MCROSS][MFREQC]
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE CROSEW(CROSS)
///   IJ0=NFREQC
///   DO IJ=1,IJ0
///      DO IT=1,MCROSS
///         CROSS(IT,IJ)=0.
///      END DO
///   END DO
///   DO IT=1,NLEVEL
///      IF(INDEXP(IT).NE.5) THEN
///         DO IJ=1,IJ0
///            FR=FREQC(IJ)
///            CROSS(IT,IJ)=SIGK(FR,IT,0)
///         END DO
///      ELSE
///         DO IJ=1,IJ0
///            FR=FREQC(IJ)
///            CROSS(IT,IJ)=SIGK(FR,IT,1)
///            IF(FR.LT.FROPC(IT)) CROSS(IT,IJ)=0.
///         END DO
///      END IF
///   END DO
/// END
/// ```
pub fn crosew(params: &CrosewParams) -> Vec<Vec<f64>> {
    let CrosewParams {
        freqc,
        nlevel,
        nfreqc,
        indexp,
        fropc,
        atomic,
    } = *params;

    // 初始化截面数组
    let mut cross = vec![vec![0.0; nfreqc]; nlevel];

    // 计算每个能级在每个频率处的截面
    for it in 0..nlevel {
        let idxp = indexp[it];

        if idxp != INDEXP_DISSOLVED {
            // 普通能级：mode = 0
            for ij in 0..nfreqc {
                let fr = freqc[ij];
                let sigk_params = SigkParams {
                    fr,
                    itr: it,
                    mode: 0,
                    atomic,
                    opdata: &super::topbas::OpData::default(),
                };
                cross[it][ij] = sigk(&sigk_params);
            }
        } else {
            // 溶解能级：mode = 1
            for ij in 0..nfreqc {
                let fr = freqc[ij];
                let sigk_params = SigkParams {
                    fr,
                    itr: it,
                    mode: 1,
                    atomic,
                    opdata: &super::topbas::OpData::default(),
                };
                cross[it][ij] = sigk(&sigk_params);

                // 如果频率低于阈值，截面设为零
                if fr < fropc[it] {
                    cross[it][ij] = 0.0;
                }
            }
        }
    }

    cross
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::state::atomic::AtomicData;

    fn create_test_atomic() -> AtomicData {
        let mut atomic = AtomicData::new();

        // 设置能级参数
        atomic.phoset.ibf[0] = 0; // 氢原子，Gaunt = 1
        atomic.phoset.ibf[1] = 0;
        atomic.phoset.ibf[2] = 0;

        // 设置电离能 (转换为频率)
        // H I 电离能 = 13.6 eV = 2.1785e-11 erg
        // 频率 = E/h = 2.1785e-11 / 6.6256e-27 = 3.288e15 Hz
        atomic.levpar.enion[0] = 3.288e15 * 6.6256e-27; // Hz * h = erg
        atomic.levpar.enion[1] = 3.288e15 * 6.6256e-27;
        atomic.levpar.enion[2] = 3.288e15 * 6.6256e-27;

        // 设置主量子数
        atomic.levpar.nquant[0] = 1;
        atomic.levpar.nquant[1] = 2;
        atomic.levpar.nquant[2] = 3;

        atomic
    }

    #[test]
    fn test_croset_basic() {
        let atomic = create_test_atomic();

        // 创建频率数组（在电离阈值之上）
        let freq = vec![4.0e15, 5.0e15, 6.0e15];
        let indexp = vec![0, 0, 0]; // 普通能级
        let fropc = vec![0.0; 3];

        let params = CrosetParams {
            freq: &freq,
            nlevel: 3,
            nfreq: 3,
            imode: 0,
            indexp: &indexp,
            fropc: &fropc,
            atomic: &atomic,
        };

        let cross = croset(&params);

        // 验证数组大小
        assert_eq!(cross.len(), 3);
        // ij0 = 2 (因为 nfreq != 1 且 imode != 2)
        assert_eq!(cross[0].len(), 2);
    }

    #[test]
    fn test_croset_single_freq() {
        let atomic = create_test_atomic();

        let freq = vec![4.0e15];
        let indexp = vec![0];
        let fropc = vec![0.0];

        let params = CrosetParams {
            freq: &freq,
            nlevel: 1,
            nfreq: 1,
            imode: 0,
            indexp: &indexp,
            fropc: &fropc,
            atomic: &atomic,
        };

        let cross = croset(&params);

        // 当 nfreq = 1 时，ij0 = 1
        assert_eq!(cross.len(), 1);
        assert_eq!(cross[0].len(), 1);
    }

    #[test]
    fn test_croset_imode_2() {
        let atomic = create_test_atomic();

        let freq = vec![4.0e15, 5.0e15, 6.0e15];
        let indexp = vec![0, 0, 0];
        let fropc = vec![0.0; 3];

        let params = CrosetParams {
            freq: &freq,
            nlevel: 3,
            nfreq: 3,
            imode: 2, // 使用所有频率
            indexp: &indexp,
            fropc: &fropc,
            atomic: &atomic,
        };

        let cross = croset(&params);

        // 当 imode = 2 时，ij0 = nfreq = 3
        assert_eq!(cross.len(), 3);
        assert_eq!(cross[0].len(), 3);
    }

    #[test]
    fn test_crosew_basic() {
        let atomic = create_test_atomic();

        let freqc = vec![4.0e15, 5.0e15, 6.0e15];
        let indexp = vec![0, 0, 0];
        let fropc = vec![0.0; 3];

        let params = CrosewParams {
            freqc: &freqc,
            nlevel: 3,
            nfreqc: 3,
            indexp: &indexp,
            fropc: &fropc,
            atomic: &atomic,
        };

        let cross = crosew(&params);

        // 验证数组大小
        assert_eq!(cross.len(), 3);
        assert_eq!(cross[0].len(), 3);
    }

    #[test]
    fn test_crosew_dissolved_level() {
        let atomic = create_test_atomic();

        let freqc = vec![4.0e15, 5.0e15];
        // INDEXP = 5 表示溶解能级
        let indexp = vec![5, 0];
        // 设置阈值频率
        let fropc = vec![3.5e15, 0.0];

        let params = CrosewParams {
            freqc: &freqc,
            nlevel: 2,
            nfreqc: 2,
            indexp: &indexp,
            fropc: &fropc,
            atomic: &atomic,
        };

        let cross = crosew(&params);

        // 验证数组大小
        assert_eq!(cross.len(), 2);
        assert_eq!(cross[0].len(), 2);

        // 能级 0 是溶解能级，频率 > fropc[0]，所以截面应该非零
        // 注意：实际值取决于 SIGK 的实现
    }
}