SpectraRust/src/math/contmd.rs

//! 盘模型对流温度确定。
//!
//! 重构自 TLUSTY `CONTMD.f`
//!
//! # 功能
//!
//! LTEGRD 的辅助过程，用于确定盘模型中对流不稳定层的温度。
//! 通过求解能量平衡方程 F(rad) + F(conv) = F(mech) 来计算，
//! 这产生一个关于对数温度梯度的三次方程。
//!
//! # 物理背景
//!
//! 在盘模型中，对流层的温度由以下平衡决定：
//! - 辐射通量 F(rad)
//! - 对流通量 F(conv)
//! - 机械通量 F(mech)
//!
//! 求解得到的 DELTA（对数温度梯度）用于更新温度结构。

use crate::state::constants::{HALF, PCK, SIG4P, UN};
use super::convec::{convec, ConvecConfig, ConvecParams};
use super::cubic::{cubic, CubicCon};
use super::conout::format_conout_header;

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

/// 温度收敛容差
const ERRT: f64 = 1e-3;

/// 最大内层迭代次数
const MAX_INNER_ITER: usize = 10;

// ============================================================================
// 配置结构体
// ============================================================================

/// CONTMD 配置参数。
#[derive(Debug, Clone)]
pub struct ContmdConfig {
    /// 混合长度参数 (HMIX0)
    pub hmix0: f64,
    /// 对流常数 A (ACONML)
    pub aconml: f64,
    /// 对流常数 B (BCONML)
    pub bconml: f64,
    /// 对流常数 C (CCONML)
    pub cconml: f64,
    /// 打印控制 (IPRING)
    pub ipring: i32,
    /// 最大对流迭代次数 (NCONIT)
    pub nconit: usize,
    /// 辐射压标志 (IFPRAD)
    /// - 0: 忽略辐射压
    /// - 1: 考虑辐射压
    pub ifprad: i32,
    /// 盘模式引力参数 (QGRAV)
    pub qgrav: f64,
}

impl Default for ContmdConfig {
    fn default() -> Self {
        Self {
            hmix0: 1.0,
            aconml: 1.0,
            bconml: 1.0,
            cconml: 1.0,
            ipring: 0,
            nconit: 20,
            ifprad: 1,
            qgrav: 1e-10,
        }
    }
}

// ============================================================================
// 输入/输出结构体
// ============================================================================

/// CONTMD 输入参数。
pub struct ContmdParams<'a> {
    /// 深度点数 (ND)
    pub nd: usize,
    /// 有效温度 (TEFF)
    pub teff: f64,
    /// 频率点数 (NFREQ)
    pub nfreq: usize,
    /// 配置
    pub config: ContmdConfig,

    // 模型状态数组 (长度 nd)
    /// 温度 (TEMP) - 可变
    pub temp: &'a mut [f64],
    /// 电子密度 (ELEC)
    pub elec: &'a [f64],
    /// 总粒子密度 (DENS) - 可变
    pub dens: &'a mut [f64],
    /// 总压力 (PTOTAL)
    pub ptotal: &'a [f64],
    /// 气压 (PGS)
    pub pgs: &'a [f64],
    /// 湍流速度 (VTURB)
    pub vturb: &'a [f64],
    /// 深度变量 (ZD)
    pub zd: &'a [f64],
    /// THETA 参数
    pub theta: &'a [f64],
    /// 辐射压 (PRADT) - 可变
    pub pradt: &'a mut [f64],

    // 不透明度数组
    /// Rosseland 不透明度/密度 (ABROSD) - 可变
    pub abrosd: &'a mut [f64],
    /// Planck 不透明度/密度 (ABPLAD) - 可变
    pub abplad: &'a mut [f64],

    // CUBCON 数据 (用于三次方程)
    /// 三次方程参数
    pub cubcon: &'a CubconData,

    // PRSAUX 数据
    /// 声速平方 (VSND2)
    pub vsnd2: &'a [f64],
    /// 辐射压尺度高度 (HR1)
    pub hr1: f64,
}

/// CUBCON 公共块数据。
#[derive(Debug, Clone, Default)]
pub struct CubconData {
    pub a: f64,
    pub b: f64,
    pub del: f64,
    pub grdadb: f64,
    pub delmde: f64,
    pub rho: f64,
    pub flxtot: f64,
    pub gravd: f64,
}

/// CONTMD 输出结果。
#[derive(Debug, Clone)]
pub struct ContmdOutput {
    /// 对流迭代次数
    pub iconit: usize,
    /// 最大温度相对变化
    pub chantm: f64,
    /// 各深度点的对流标志 (1 = 对流不稳定)
    pub iconv: Vec<i32>,
    /// 各深度点的温度变化
    pub delta_temp: Vec<f64>,
}

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

/// 计算盘模型对流层的温度 (CONTMD)。
///
/// # 参数
///
/// * `params` - 输入参数
///
/// # 返回值
///
/// 返回 `ContmdOutput`，包含迭代次数、温度变化等信息。
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE CONTMD
///   INCLUDE 'IMPLIC.FOR'
///   INCLUDE 'BASICS.FOR'
///   INCLUDE 'ATOMIC.FOR'
///   INCLUDE 'MODELQ.FOR'
///   INCLUDE 'ALIPAR.FOR'
///   COMMON ESEMAT(MLEVEL,MLEVEL),BESE(MLEVEL),
///   *       DEPTH(MDEPTH),DEPTH0(MDEPTH),TAU(MDEPTH),TAU0(MDEPTH),
///   *       TEMP0(MDEPTH),ELEC0(MDEPTH),DENS0(MDEPTH),DM0(MDEPTH)
///   DIMENSION DELTR(MDEPTH),TEMPR(MDEPTH),ICON0(MDEPTH)
///   COMMON/CUBCON/A,B,DEL,GRDADB,DELMDE,RHO,FLXTOT,GRAVD
///   COMMON/PRSAUX/VSND2(MDEPTH),HG1,HR1,RR1
///   ...
/// END
/// ```
pub fn contmd_pure(params: &mut ContmdParams) -> ContmdOutput {
    let nd = params.nd;

    // 初始化输出
    let mut iconv = vec![0; nd];
    let mut delta_temp = vec![0.0; nd];

    // 存储辐射平衡温度和梯度
    let mut tempr = vec![0.0; nd];
    let mut deltr = vec![0.0; nd];

    // 计算总通量
    let t4 = params.teff.powi(4);
    let flxto0 = SIG4P * t4;

    // 辐射压
    let mut dprad = 1.891204931e-15 * t4;
    if params.config.ifprad == 0 {
        dprad = 0.0;
    }
    let _prad0 = dprad / 1.732;

    // 存储初始温度和计算辐射梯度
    for id in 0..nd {
        tempr[id] = params.temp[id];
        if id == 0 {
            deltr[id] = 0.0;
        } else {
            // DELTR = d(ln T)/d(ln P)
            let p_plus = params.ptotal[id] + params.ptotal[id - 1];
            let p_minus = params.ptotal[id] - params.ptotal[id - 1];
            if p_minus.abs() > 0.0 && params.temp[id] + params.temp[id - 1] > 0.0 {
                deltr[id] = (params.temp[id] - params.temp[id - 1]) / p_minus
                    * p_plus
                    / (params.temp[id] + params.temp[id - 1]);
            } else {
                deltr[id] = 0.0;
            }
        }
    }

    // 初始化辅助变量
    let mut iconbe = 0;
    let mut deltc = 0.0;
    let hr1 = params.hr1;

    // 全局迭代循环
    let mut iconit = 0;
    let mut chantm = 0.0;

    loop {
        iconit += 1;
        iconbe = 0;

        // 辐射压尺度高度
        let _hr1_val = flxto0 * PCK * params.abrosd[0] / params.config.qgrav;

        chantm = 0.0;
        let mut pradm = if nd > 0 { params.pradt[0] } else { 0.0 };

        // 遍历所有深度点
        for id in 0..nd {
            let mut t = params.temp[id];
            let ptot = params.ptotal[id];
            let pgas = params.pgs[id];
            let pturb = HALF * params.dens[id] * params.vturb[id].powi(2);
            let prad = params.pradt[id];
            let flxtot = flxto0 * (UN - params.theta[id]);
            let gravd = params.zd[id] * params.config.qgrav;

            iconv[id] = 0;
            let mut delt0 = 0.0;

            if id == 0 {
                // 表面层：直接更新
                delt0 = params.temp[id] - t;
            } else {
                // 内部层：迭代求解对流温度
                let mut j = 0;

                // 初始温度估计
                if iconit == 1 {
                    t = t - tempr[id - 1] + params.temp[id - 1];
                }

                let tm = params.temp[id - 1];
                if t < 0.0 {
                    t = tm;
                }

                let pgm = params.pgs[id - 1];
                let ptotm = params.ptotal[id - 1];
                let pt0 = HALF * (ptot + ptotm);
                let delr = deltr[id];

                // 内层迭代循环
                loop {
                    j += 1;
                    let told = t;

                    let t0 = HALF * (t + tm);
                    let pg0 = HALF * (pgas + pgm);
                    let pr0 = HALF * (prad + pradm);
                    let ab0 = HALF * (params.abrosd[id] + params.abrosd[id - 1]);

                    // 检查是否需要计算对流
                    if id >= nd - 2 && iconbe == 0 {
                        // 接近底部且尚未开始对流，跳过
                        delt0 = params.temp[id] - t;
                        break;
                    }

                    // 计算对流通量
                    let convec_config = ConvecConfig {
                        hmix0: params.config.hmix0,
                        aconml: params.config.aconml,
                        bconml: params.config.bconml,
                        cconml: params.config.cconml,
                        idisk: 1, // 盘模式
                        ioptab: 0,
                        flxtot,
                        gravd,
                        grav: params.config.qgrav,
                    };

                    let convec_params = ConvecParams {
                        id: id + 1, // 1-based
                        t: t0,
                        ptot: pt0,
                        pg: pg0,
                        prad: pr0,
                        abros: ab0,
                        delta: delr,
                        taurs: 0.0, // 简化
                        config: convec_config,
                        trmder_config: None,
                        therm_tables: None,
                    };

                    let convec_out = convec(&convec_params);
                    let flxcnv = convec_out.flxcnv;
                    let vcon = convec_out.vconv;

                    if flxcnv == 0.0 {
                        // 无对流
                        delt0 = params.temp[id] - t;
                        break;
                    }

                    iconv[id] = 1;
                    iconbe = 1;

                    // 检查是否在底部
                    if id == nd - 1 {
                        // 底部边界：使用简单公式
                        let p_diff = ptot - ptotm;
                        if p_diff.abs() > 1e-30 {
                            let pip = (ptot + ptotm) / p_diff;
                            let denom = pip - delr;
                            if denom.abs() > 1e-30 {
                                t = tm * (pip + delr) / denom;
                            }
                        }
                        if !t.is_finite() || t <= 0.0 {
                            t = tm;
                        }
                        delt0 = params.temp[id] - t;
                        break;
                    }

                    // 使用三次方程求解 DELTA
                    let cubcon = CubicCon {
                        a: params.cubcon.a,
                        b: params.cubcon.b,
                        del: params.cubcon.del,
                        grdadb: convec_out.grdadb,
                        rho: convec_out.rho,
                        flxtot,
                        gravd,
                    };

                    let delta0 = cubic(&cubcon);

                    // 计算新的温度
                    let p_sum = ptot + ptotm;
                    let fac = if p_sum.abs() > 1e-30 {
                        delta0 * (ptot - ptotm) / p_sum
                    } else {
                        0.0
                    };

                    let denom = UN - fac;
                    if denom.abs() > 1e-30 {
                        t = tm * (UN + fac) / denom;
                    }

                    if !t.is_finite() || t < tm {
                        t = tm;
                    }

                    // 收敛检查
                    let rel_change = if told != 0.0 {
                        (UN - t / told).abs()
                    } else {
                        0.0
                    };

                    if rel_change <= ERRT || j >= MAX_INNER_ITER {
                        delt0 = params.temp[id] - t;
                        break;
                    }
                }
            }

            // 存储最终量
            if id > 0 && iconv[id] == 0 && iconv[id - 1] == 1 {
                deltc = delt0;
            }

            if id == nd - 1 {
                let ptotm = if id > 0 { params.ptotal[id - 1] } else { ptot };
                let tm = if id > 0 { params.temp[id - 1] } else { t };
                let delr = if id > 0 { deltr[id] } else { 0.0 };

                let p_diff = ptot - ptotm;
                if p_diff.abs() > 1e-30 {
                    let pip = (ptot + ptotm) / p_diff;
                    let denom = pip - delr;
                    if denom.abs() > 1e-30 {
                        let t_new = tm * (pip + delr) / denom;
                        if t_new.is_finite() && t_new > 0.0 {
                            t = t_new;
                        }
                    }
                }
            }

            delt0 = params.temp[id] - t;

            // 确保 t 是有效的
            if !t.is_finite() || t <= 0.0 {
                t = params.temp[id]; // 保持原值
                delt0 = 0.0;
            }

            // 更新辐射压
            if params.temp[id].abs() > 1e-30 && t.is_finite() && t > 0.0 {
                params.pradt[id] = params.pradt[id] * (t / params.temp[id]).powi(4);
            }

            // 更新密度
            if t.is_finite() && t > 0.0 && params.temp[id].abs() > 1e-30 {
                params.dens[id] = params.dens[id] * (params.temp[id] / t);
            }

            // 计算温度相对变化
            let chant0 = if params.temp[id] != 0.0 {
                (t - params.temp[id]).abs() / params.temp[id]
            } else {
                0.0
            };

            if chant0 > chantm {
                chantm = chant0;
            }

            // 更新温度
            delta_temp[id] = t - params.temp[id];
            params.temp[id] = t;

            // 处理对流区边缘
            if iconit > 1 && iconv[id] == 0 && iconbe == 1 {
                params.temp[id] = t - deltc;
            }

            pradm = params.pradt[id];
        }

        // 收敛检查
        if chantm <= ERRT || iconit >= params.config.nconit {
            break;
        }
    }

    ContmdOutput {
        iconit,
        chantm,
        iconv,
        delta_temp,
    }
}

// ============================================================================
// I/O 函数
// ============================================================================

/// 格式化 CONTMD 迭代信息。
pub fn format_contmd_iter(iconit: usize) -> String {
    format!("\n\n CONVECTIVE FLUX: AT CONTMD, ITER={:2}\n", iconit)
}

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

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

    const ND: usize = 50;

    fn create_test_params() -> ContmdParams<'static> {
        let config = ContmdConfig::default();

        // 创建测试数据
        let mut temp = vec![10000.0; ND];
        let elec = vec![1e12; ND];
        let mut dens = vec![1e-7; ND];
        let ptotal = vec![1e5; ND];
        let pgs = vec![1e5; ND];
        let vturb = vec![0.0; ND];
        let zd = vec![1e10; ND];
        let theta = vec![0.0; ND];
        let mut pradt = vec![0.0; ND];
        let mut abrosd = vec![0.1; ND];
        let mut abplad = vec![0.1; ND];
        let vsnd2 = vec![1e10; ND];
        let cubcon = CubconData::default();

        // 设置温度梯度
        for i in 0..ND {
            temp[i] = 10000.0 - i as f64 * 100.0;
        }

        ContmdParams {
            nd: ND,
            teff: 35000.0,
            nfreq: 100,
            config,
            temp: Box::leak(temp.into_boxed_slice()),
            elec: Box::leak(elec.into_boxed_slice()),
            dens: Box::leak(dens.into_boxed_slice()),
            ptotal: Box::leak(ptotal.into_boxed_slice()),
            pgs: Box::leak(pgs.into_boxed_slice()),
            vturb: Box::leak(vturb.into_boxed_slice()),
            zd: Box::leak(zd.into_boxed_slice()),
            theta: Box::leak(theta.into_boxed_slice()),
            pradt: Box::leak(pradt.into_boxed_slice()),
            abrosd: Box::leak(abrosd.into_boxed_slice()),
            abplad: Box::leak(abplad.into_boxed_slice()),
            cubcon: Box::leak(Box::new(cubcon)),
            vsnd2: Box::leak(vsnd2.into_boxed_slice()),
            hr1: 1e10,
        }
    }

    #[test]
    fn test_contmd_basic() {
        let mut params = create_test_params();
        let output = contmd_pure(&mut params);

        // 验证迭代次数在合理范围内
        assert!(output.iconit <= params.config.nconit);
        assert!(output.iconit > 0);

        // 验证输出数组长度
        assert_eq!(output.iconv.len(), ND);
        assert_eq!(output.delta_temp.len(), ND);
    }

    #[test]
    fn test_contmd_no_convection() {
        let mut params = create_test_params();

        // 禁用对流
        params.config.hmix0 = -1.0;

        let output = contmd_pure(&mut params);

        // 禁用对流时不应该有对流区
        for &iconv in &output.iconv {
            assert_eq!(iconv, 0);
        }
    }

    #[test]
    fn test_contmd_temperature_update() {
        let mut params = create_test_params();

        // 保存原始温度
        let orig_temp = params.temp.to_vec();

        let _output = contmd_pure(&mut params);

        // 温度可能被更新
        // 检查温度仍然是有限值
        for &t in params.temp.iter() {
            assert!(t.is_finite());
            assert!(t > 0.0);
        }
    }

    #[test]
    fn test_format_contmd_iter() {
        let msg = format_contmd_iter(5);
        assert!(msg.contains("5"));
        assert!(msg.contains("ITER"));
    }

    #[test]
    fn test_cubcon_data() {
        let cubcon = CubconData {
            a: 1.0,
            b: 2.0,
            del: 0.1,
            grdadb: 0.4,
            delmde: 0.0,
            rho: 1e-7,
            flxtot: 1e10,
            gravd: 1e4,
        };

        assert!((cubcon.a - 1.0).abs() < 1e-10);
        assert!((cubcon.b - 2.0).abs() < 1e-10);
    }

    #[test]
    fn test_config_default() {
        let config = ContmdConfig::default();

        assert!((config.hmix0 - 1.0).abs() < 1e-10);
        assert_eq!(config.nconit, 20);
        assert_eq!(config.ifprad, 1);
    }

    #[test]
    fn test_small_nd() {
        // 测试小深度点数情况
        let nd = 3;

        let config = ContmdConfig::default();
        let temp = vec![10000.0, 9000.0, 8000.0];
        let elec = vec![1e12; nd];
        let dens = vec![1e-7; nd];
        let ptotal = vec![1e5; nd];
        let pgs = vec![1e5; nd];
        let vturb = vec![0.0; nd];
        let zd = vec![1e10; nd];
        let theta = vec![0.0; nd];
        let pradt = vec![0.0; nd];
        let abrosd = vec![0.1; nd];
        let abplad = vec![0.1; nd];
        let vsnd2 = vec![1e10; nd];
        let cubcon = CubconData::default();

        let mut params = ContmdParams {
            nd,
            teff: 35000.0,
            nfreq: 100,
            config,
            temp: Box::leak(temp.into_boxed_slice()),
            elec: Box::leak(elec.into_boxed_slice()),
            dens: Box::leak(dens.into_boxed_slice()),
            ptotal: Box::leak(ptotal.into_boxed_slice()),
            pgs: Box::leak(pgs.into_boxed_slice()),
            vturb: Box::leak(vturb.into_boxed_slice()),
            zd: Box::leak(zd.into_boxed_slice()),
            theta: Box::leak(theta.into_boxed_slice()),
            pradt: Box::leak(pradt.into_boxed_slice()),
            abrosd: Box::leak(abrosd.into_boxed_slice()),
            abplad: Box::leak(abplad.into_boxed_slice()),
            cubcon: Box::leak(Box::new(cubcon)),
            vsnd2: Box::leak(vsnd2.into_boxed_slice()),
            hr1: 1e10,
        };

        let output = contmd_pure(&mut params);

        assert_eq!(output.iconv.len(), nd);
        assert!(output.iconit > 0);
    }
}