SpectraRust/src/tlusty/math/partition/partf.rs

//! 配分函数计算（氢到锌，中性及前四个电离级）。
//!
//! 重构自 TLUSTY `PARTF` 子程序。
//!
//! # 功能
//!
//! - 计算 Z=1-30 元素的配分函数
//! - 基于 Traving, Baschek, Holweger 公式
//! - 特殊处理：Fe, Ni, 重元素
//!
//! # 参考
//!
//! Traving, Baschek, and Holweger, Abhand. Hamburg. Sternwarte. Band VIII, Nr. 1 (1966)

use crate::tlusty::data;
use crate::tlusty::io::{Result, IoError};

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

/// ln(10)
const LN10: f64 = std::f64::consts::LN_10;

// ============================================================================
// 配分函数模式
// ============================================================================

/// 配分函数计算模式
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum PartfMode {
    /// 标准公式 (Traving-Baschek-Holweger)
    Standard,
    /// 用户自定义 (PFSPEC)
    UserDefined,
    /// Opacity Project 数据 (OPFRAC)
    OpacityProject,
}

// ============================================================================
// 输入参数结构体
// ============================================================================

/// PARTF 输入参数。
pub struct PartfParams {
    /// 原子序数
    pub iat: i32,
    /// 离子序数 (1=中性, 2=一次电离, ...)
    pub izi: i32,
    /// 温度 (K)
    pub t: f64,
    /// 电子密度 (cm⁻³)
    pub ane: f64,
    /// 最高束缚态主量子数
    pub xmax: f64,
    /// 计算模式
    pub mode: PartfMode,
    /// Irwin 模式标志 (0=不使用, >0=使用)
    pub iirwin: i32,
}

/// PARTF 输出结果。
#[derive(Debug, Clone)]
pub struct PartfOutput {
    /// 配分函数
    pub u: f64,
    /// dU/dT 导数
    pub dut: f64,
    /// dU/d(ANE) 导数
    pub dun: f64,
}

// ============================================================================
// 辅助数组初始化
// ============================================================================

/// INDEXS 数组：每个离子的起始索引
static mut INDEXS: [i32; 123] = [0; 123];
/// INDEXM 数组：每个能级的起始索引
static mut INDEXM: [i32; 222] = [0; 222];
static mut ICOMP: i32 = 0;

/// 初始化辅助数组
fn init_arrays() {
    unsafe {
        if ICOMP != 0 {
            return;
        }

        // 构建 INDEXS
        let mut ind = 1i32;
        let is_data = [
            &data::PARTF_IS1[..], &data::PARTF_IS2[..]
        ];
        let is_combined: Vec<f64> = is_data.iter().flat_map(|x| x.iter().copied()).collect();

        for k in 0..123 {
            INDEXS[k] = ind;
            ind += is_combined[k] as i32;
        }

        // 构建 INDEXM
        ind = 1;
        let im_data = [
            &data::PARTF_IM1[..], &data::PARTF_IM2[..]
        ];
        let im_combined: Vec<f64> = im_data.iter().flat_map(|x| x.iter().copied()).collect();

        for k in 0..222.min(im_combined.len()) {
            INDEXM[k] = ind;
            ind += im_combined[k] as i32;
        }

        ICOMP = 1;
    }
}

// ============================================================================
// 核心计算函数
// ============================================================================

/// 执行 PARTF 计算（纯计算部分）。
///
/// # 参数
/// - `params`: 输入参数
///
/// # 返回
/// 配分函数及其导数
pub fn partf(params: &PartfParams) -> PartfOutput {
    // 初始化辅助数组
    init_arrays();

    let iat = params.iat;
    let izi = params.izi;
    let t = params.t;
    let ane = params.ane;
    let xmax = params.xmax;

    // 检查有效性
    if izi <= 0 || iat <= 0 {
        return partf_user_defined(params);
    }

    // 高电离级处理 (IZI > 5)
    if izi > 5 {
        if iat < izi {
            return PartfOutput { u: 1.0, dut: 0.0, dun: 0.0 };
        }
        if iat > 8 && iat <= 28 {
            // 使用 IGLE 数组
            let idx = (iat - izi + 1) as usize;
            if idx > 0 && idx <= data::PARTF_IGLE.len() {
                return PartfOutput {
                    u: data::PARTF_IGLE[idx - 1],
                    dut: 0.0,
                    dun: 0.0,
                };
            }
        }
        // CNO 高电离级
        return partf_cno(iat, izi, t, ane);
    }

    // Irwin 配分函数模式 (iirwin > 0 且 T < 16000K)
    if params.iirwin > 0 && t < 16000.0 {
        if izi <= 2 {
            return partf_irwin(params);
        }
    } else if iat > 30 && izi <= 3 {
        // 重元素 (Z > 30, 低电离级)
        return partf_heavy(params);
    }

    // 根据模式选择计算方法
    match params.mode {
        PartfMode::UserDefined => partf_user_defined(params),
        PartfMode::OpacityProject => partf_opacity_project(params),
        PartfMode::Standard => {
            // 特殊处理：Fe 和 Ni
            if iat == 26 && izi >= 4 && izi <= 9 {
                return partf_fe(params);
            }
            if iat == 28 && izi >= 4 && izi <= 9 {
                return partf_ni(params);
            }
            // 重元素 (Z > 30)
            if iat > 30 && izi <= 3 {
                return partf_heavy(params);
            }

            // 标准 Traving-Baschek-Holweger 公式
            partf_standard(params)
        }
    }
}

/// 标准 Traving-Baschek-Holweger 配分函数计算
fn partf_standard(params: &PartfParams) -> PartfOutput {
    let iat = params.iat;
    let izi = params.izi;
    let t = params.t;
    let ane = params.ane;
    let xmax = params.xmax;

    // 获取 INDEX0 (II1 或 II2)
    let i0 = if iat >= 1 && iat <= 15 && izi >= 1 && izi <= 5 {
        data::PARTF_II1[(izi - 1) as usize][(iat - 1) as usize]
    } else if iat >= 16 && iat <= 30 && izi >= 1 && izi <= 5 {
        data::PARTF_II2[(izi - 1) as usize][(iat - 16) as usize]
    } else {
        0.0
    };

    if i0 <= 0.0 {
        // 固定配分函数值 (负数表示固定值)
        return PartfOutput {
            u: (-i0),
            dut: 0.0,
            dun: 0.0,
        };
    }

    let i0_idx = (i0 - 1.0) as usize;

    // Traving-Baschek-Holweger 公式
    let qz = izi as f64;
    let thet = 5.0404e3 / t;
    let a = 31.321 * qz * qz * thet;
    let xmax2 = xmax * xmax;

    let sixth = 1.0 / 6.0;
    let half = 0.5;
    let third = 1.0 / 3.0;
    let trha = data::PARTF_TRHA;

    let qq = xmax / 4.0 * (xmax2 + xmax + sixth + a + a * a * half / xmax2);
    let qas1 = xmax * third * (xmax2 + trha * xmax + half);

    // 获取基态统计权重
    let ig0 = if i0_idx < data::PARTF_IG01.len() {
        data::PARTF_IG01[i0_idx]
    } else if i0_idx < data::PARTF_IG01.len() + data::PARTF_IG02.len() {
        data::PARTF_IG02[i0_idx - data::PARTF_IG01.len()]
    } else {
        1.0
    };

    // 获取 IS (能级数)
    let is_val = if i0_idx < data::PARTF_IS1.len() {
        data::PARTF_IS1[i0_idx] as i32
    } else if i0_idx < data::PARTF_IS1.len() + data::PARTF_IS2.len() {
        data::PARTF_IS2[i0_idx - data::PARTF_IS1.len()] as i32
    } else {
        1
    };

    // 获取起始索引
    let is0 = unsafe { INDEXS[i0_idx] } as usize;

    // 合并 XL 数组
    let xl_combined: Vec<f64> = data::PARTF_XL1.iter()
        .chain(data::PARTF_XL2.iter())
        .copied()
        .collect();

    // 合并 CHION 数组
    let chion_combined: Vec<f64> = data::PARTF_CH1.iter()
        .chain(data::PARTF_CH2.iter())
        .chain(data::PARTF_CH3.iter())
        .chain(data::PARTF_CH4.iter())
        .copied()
        .collect();

    // 合并 IGP 数组
    let igp_combined: Vec<f64> = data::PARTF_IGP1.iter()
        .chain(data::PARTF_IGP2.iter())
        .copied()
        .collect();

    // 合并 ALF 和 GAM 数组
    // ALF 由各元素的 A 数组组成
    let alf_combined: Vec<f64> = data::PARTF_AHH.iter()
        .chain(data::PARTF_ALB.iter())
        .chain(data::PARTF_AB.iter())
        .chain(data::PARTF_AC.iter())
        .chain(data::PARTF_AN.iter())
        .chain(data::PARTF_AO.iter())
        .chain(data::PARTF_AF.iter())
        .chain(data::PARTF_ANN.iter())
        .chain(data::PARTF_ANA.iter())
        .chain(data::PARTF_AMG.iter())
        .chain(data::PARTF_AAL.iter())
        .chain(data::PARTF_ASI.iter())
        .chain(data::PARTF_AP.iter())
        .chain(data::PARTF_AS.iter())
        .chain(data::PARTF_ACL.iter())
        .chain(data::PARTF_AAR.iter())
        .chain(data::PARTF_AK.iter())
        .chain(data::PARTF_ACA.iter())
        .chain(data::PARTF_ASC.iter())
        .chain(data::PARTF_ATI.iter())
        .chain(data::PARTF_AV.iter())
        .chain(data::PARTF_ACR.iter())
        .chain(data::PARTF_AMN.iter())
        .chain(data::PARTF_AFE.iter())
        .chain(data::PARTF_ACO.iter())
        .chain(data::PARTF_ANI.iter())
        .chain(data::PARTF_ACU.iter())
        .chain(data::PARTF_AZN.iter())
        .copied()
        .collect();

    let gam_combined: Vec<f64> = data::PARTF_GHH.iter()
        .chain(data::PARTF_GLB.iter())
        .chain(data::PARTF_GB.iter())
        .chain(data::PARTF_GC.iter())
        .chain(data::PARTF_GN.iter())
        .chain(data::PARTF_GO.iter())
        .chain(data::PARTF_GF.iter())
        .chain(data::PARTF_GNN.iter())
        .chain(data::PARTF_GNA.iter())
        .chain(data::PARTF_GMG.iter())
        .chain(data::PARTF_GAL.iter())
        .chain(data::PARTF_GSI.iter())
        .chain(data::PARTF_GP.iter())
        .chain(data::PARTF_GS.iter())
        .chain(data::PARTF_GCL.iter())
        .chain(data::PARTF_GAR.iter())
        .chain(data::PARTF_GK.iter())
        .chain(data::PARTF_GCA.iter())
        .chain(data::PARTF_GSC.iter())
        .chain(data::PARTF_GTI.iter())
        .chain(data::PARTF_GV.iter())
        .chain(data::PARTF_GCR.iter())
        .chain(data::PARTF_GMN.iter())
        .chain(data::PARTF_GFE.iter())
        .chain(data::PARTF_GCO.iter())
        .chain(data::PARTF_GNI.iter())
        .chain(data::PARTF_GCU.iter())
        .chain(data::PARTF_GZN.iter())
        .copied()
        .collect();

    // 主循环：遍历能级
    let mut su1 = 0.0_f64;
    let mut su2 = 0.0_f64;
    let mut sqa = 0.0_f64;
    let mut sqq = 0.0_f64;
    let mut sqt = 0.0_f64;
    let mut sq2 = 0.0_f64;

    for k in is0..(is0 + is_val as usize).min(xl_combined.len()).min(chion_combined.len()).min(igp_combined.len()) {
        let xxl = xl_combined[k];
        let gpr = igp_combined[k];
        let ch = chion_combined[k];

        let x = ch * thet;
        let ex = if x < 30.0 { (-x * LN10).exp() } else { 0.0 };

        sqq += gpr * ex;
        let qas = (qas1 - xxl * third * (xxl * xxl + trha * xxl + half)
            + (xmax - xxl) * (1.0 + a * half / xxl / xmax) * a) * gpr * ex;
        sqa += qas;
        sq2 += qas * ch;
        sqt += gpr * (xmax - xxl) * (1.0 + a / xmax / xxl) * ex;

        // 获取 IM (项数)
        let im_val = if k < data::PARTF_IM1.len() {
            data::PARTF_IM1[k] as i32
        } else if k < data::PARTF_IM1.len() + data::PARTF_IM2.len() {
            data::PARTF_IM2[k - data::PARTF_IM1.len()] as i32
        } else {
            1
        };

        let m0 = unsafe { INDEXM[k] } as usize;

        let mut al1 = 0.0_f64;
        let mut al2 = 0.0_f64;

        for m in m0..(m0 + im_val as usize).min(alf_combined.len()).min(gam_combined.len()) {
            let xg = gam_combined[m] * thet;
            if xg <= 20.0 {
                let xm = (-xg * LN10).exp() * alf_combined[m];
                al1 += xm;
                al2 += gam_combined[m] * xm;
            }
        }

        su1 += al1;
        su2 += al2;
    }

    // 计算最终配分函数
    let mut u = ig0 + su1 + sqa;
    if u < 0.0 {
        u = ig0;
    }

    let dut = (LN10 * thet * (su2 + sq2) + qq * sqq - a * sqt) / t;
    let dun = -qq * sqq / ane;

    PartfOutput { u, dut, dun }
}

/// Irwin 配分函数 (MPARTF)
fn partf_irwin(params: &PartfParams) -> PartfOutput {
    use crate::tlusty::math::partition::mpartf;

    let result = mpartf(params.iat as usize, params.izi as usize, 0, params.t);
    let u0 = result.u;
    let du0 = result.dulog;
    let mut dut = 0.0;
    if u0 > 0.0 && du0 > 0.0 {
        dut = u0 / params.t * du0;
    }
    PartfOutput {
        u: u0,
        dut,
        dun: 0.0,
    }
}

/// 用户自定义配分函数
fn partf_user_defined(params: &PartfParams) -> PartfOutput {
    // 调用 PFSPEC
    use crate::tlusty::math::pfspec;

    let (u, dut, dun) = pfspec(params.iat, params.izi, params.t, params.ane);
    PartfOutput { u, dut, dun }
}

/// Opacity Project 配分函数
fn partf_opacity_project(params: &PartfParams) -> PartfOutput {
    // 调用 OPFRAC
    use crate::tlusty::math::{opfrac, OpfracParams, PfOptB};

    let opfrac_params = OpfracParams {
        iat: params.iat,
        ion: params.izi,
        t: params.t,
        ane: params.ane,
    };

    let pfoptb = PfOptB::new();
    let result = opfrac(&opfrac_params, &pfoptb);

    PartfOutput {
        u: result.pf,
        dut: 0.0,
        dun: 0.0,
    }
}

/// Fe 配分函数
fn partf_fe(params: &PartfParams) -> PartfOutput {
    // 调用 PFFE
    use crate::tlusty::math::pffe;

    let (u, dut, dun) = pffe(params.izi, params.t, params.ane);
    PartfOutput { u, dut, dun }
}

/// Ni 配分函数
fn partf_ni(params: &PartfParams) -> PartfOutput {
    // 调用 PFNI
    use crate::tlusty::math::pfni;

    let (u, dut, dun) = pfni(params.izi, params.t);
    PartfOutput { u, dut, dun }
}

/// 重元素配分函数 (Z > 30)
fn partf_heavy(params: &PartfParams) -> PartfOutput {
    // 调用 PFHEAV
    use crate::tlusty::math::{pfheav, PfheavParams};

    let pfheav_params = PfheavParams {
        iiz: params.iat,
        jnion: params.izi,
        mode: 3,
        t: params.t,
        ane: params.ane,
    };

    let result = pfheav(&pfheav_params);
    PartfOutput {
        u: result.u,
        dut: 0.0,
        dun: 0.0,
    }
}

/// CNO 高电离级配分函数
fn partf_cno(iat: i32, izi: i32, t: f64, ane: f64) -> PartfOutput {
    // 调用 PFCNO
    use crate::tlusty::math::pfcno;

    let pf = pfcno(iat as usize, izi as usize, t, ane);
    PartfOutput {
        u: pf,
        dut: 0.0,
        dun: 0.0,
    }
}

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

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

    #[test]
    fn test_partf_hydrogen() {
        let params = PartfParams {
            iat: 1,
            izi: 1,
            t: 10000.0,
            ane: 1.0e12,
            xmax: 10.0,
            mode: PartfMode::Standard,
            iirwin: 0,
        };

        let result = partf(&params);
        assert!(result.u > 0.0, "Partition function should be positive");
    }

    #[test]
    fn test_partf_helium() {
        let params = PartfParams {
            iat: 2,
            izi: 1,
            t: 10000.0,
            ane: 1.0e12,
            xmax: 8.0,
            mode: PartfMode::Standard,
            iirwin: 0,
        };

        let result = partf(&params);
        assert!(result.u > 0.0, "Partition function should be positive");
    }

    #[test]
    fn test_partf_iron() {
        let params = PartfParams {
            iat: 26,
            izi: 4,
            t: 10000.0,
            ane: 1.0e12,
            xmax: 7.0,
            mode: PartfMode::Standard,
            iirwin: 0,
        };

        let result = partf(&params);
        assert!(result.u > 0.0, "Iron partition function should be positive");
    }

    #[test]
    fn test_partf_nickel() {
        let params = PartfParams {
            iat: 28,
            izi: 4,
            t: 10000.0,
            ane: 1.0e12,
            xmax: 7.0,
            mode: PartfMode::Standard,
            iirwin: 0,
        };

        let result = partf(&params);
        assert!(result.u > 0.0, "Nickel partition function should be positive");
    }

    #[test]
    fn test_partf_high_ionization() {
        // 高电离级 (IZI > 5)
        let params = PartfParams {
            iat: 8,
            izi: 7,
            t: 50000.0,
            ane: 1.0e15,
            xmax: 5.0,
            mode: PartfMode::Standard,
            iirwin: 0,
        };

        let result = partf(&params);
        assert!(result.u > 0.0, "High ionization partition function should be positive");
    }

    #[test]
    fn test_partf_opacity_project() {
        let params = PartfParams {
            iat: 1,
            izi: 1,
            t: 10000.0,
            ane: 1.0e12,
            xmax: 10.0,
            mode: PartfMode::OpacityProject,
            iirwin: 0,
        };

        let result = partf(&params);
        assert!(result.u > 0.0, "Opacity Project partition function should be positive");
    }

    #[test]
    fn test_data_arrays() {
        // 验证 data.rs 中的数组
        assert!(data::PARTF_IGLE.len() == 28);
        assert!(data::PARTF_AHH.len() == 6);
        assert!(data::PARTF_GHH.len() == 6);
        assert!(data::PARTF_XL1.len() == 99);
        assert!(data::PARTF_XL2.len() == 123);
    }
}