SpectraRust/src/tlusty/math/utils/newdmt.rs

//! 基于温度变化的新深度网格计算。
//!
//! 重构自 TLUSTY `newdmt.f`
//!
//! # 功能
//!
//! 计算新的 m-scale（柱质量密度），基于新网格更好地表示温度变化。
//! 使用 GRIDP 函数找到最优网格点分布。

use crate::tlusty::math::gridp;
use crate::tlusty::math::interp;
use crate::tlusty::state::constants::{HALF, MDEPTH, TWO, UN};
// f2r_depends: GRIDP, HESOLV, TEMPER

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

/// ln(10) = 2.3025851...
const LN10: f64 = 2.30258509299404568402;

// ============================================================================
// 辅助参数结构体
// ============================================================================

/// PRSAUX COMMON 块参数
#[derive(Debug, Clone)]
pub struct NewdmtPrsAux {
    /// 声速平方 [深度]
    pub vsnd2: Vec<f64>,
    /// 表面气压标高
    pub hg1: f64,
    /// 辐射气压标高
    pub hr1: f64,
    /// 辐射/气压标高比
    pub rr1: f64,
}

impl Default for NewdmtPrsAux {
    fn default() -> Self {
        Self {
            vsnd2: vec![0.0; MDEPTH],
            hg1: 0.0,
            hr1: 0.0,
            rr1: 0.0,
        }
    }
}

/// FACTRS COMMON 块参数
#[derive(Debug, Clone)]
pub struct NewdmtFactrs {
    /// 辐射因子 [深度]
    pub gamj: Vec<f64>,
    /// 几何稀释因子
    pub gamh: f64,
    pub fak0: f64,
}

impl Default for NewdmtFactrs {
    fn default() -> Self {
        Self {
            gamj: vec![1.0; MDEPTH],
            gamh: 1.0,
            fak0: 0.0,
        }
    }
}

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

/// NEWDMT 配置参数。
#[derive(Debug, Clone)]
pub struct NewdmtConfig {
    /// 深度点数
    pub nd: usize,
    /// 粘性深度参数
    pub dmvisc: f64,
    /// ZETA0 参数
    pub zeta0: f64,
    /// ZETA1 参数
    pub zeta1: f64,
    /// 粘性分数
    pub fractv: f64,
    /// 收敛迭代控制
    pub nconit: i32,
    /// 打印控制
    pub ipring: i32,
}

impl Default for NewdmtConfig {
    fn default() -> Self {
        Self {
            nd: 50,
            dmvisc: 0.0,
            zeta0: 0.0,
            zeta1: 0.0,
            fractv: 0.0,
            nconit: 0,
            ipring: 0,
        }
    }
}

// ============================================================================
// 模型状态结构体
// ============================================================================

/// NEWDMT 模型状态（可变）。
#[derive(Debug, Clone)]
pub struct NewdmtModelState {
    /// 深度 (柱质量密度, g/cm²) [深度]
    pub dm: Vec<f64>,
    /// 温度 (K) [深度]
    pub temp: Vec<f64>,
    /// 电子密度 (cm⁻³) [深度]
    pub elec: Vec<f64>,
    /// 总粒子密度 (cm⁻³) [深度]
    pub dens: Vec<f64>,
    /// 深度变量 (cm) [深度]
    pub zd: Vec<f64>,
    /// Rosseland 光学深度 [深度]
    pub tauros: Vec<f64>,
    /// 热光学深度 [深度]
    pub tauthe: Vec<f64>,
    /// Theta 函数 [深度]
    pub theta: Vec<f64>,
    /// 粘性系数 [深度]
    pub viscd: Vec<f64>,
    /// Rosseland 平均不透明度 [深度]
    pub abrosd: Vec<f64>,
    /// Planck 平均不透明度 [深度]
    pub abplad: Vec<f64>,
    /// 平均分子量 [深度]
    pub wmm: Vec<f64>,
    /// 总压力 [深度]
    pub ptotal: Vec<f64>,
}

impl Default for NewdmtModelState {
    fn default() -> Self {
        Self {
            dm: vec![0.0; MDEPTH],
            temp: vec![0.0; MDEPTH],
            elec: vec![0.0; MDEPTH],
            dens: vec![0.0; MDEPTH],
            zd: vec![0.0; MDEPTH],
            tauros: vec![0.0; MDEPTH],
            tauthe: vec![0.0; MDEPTH],
            theta: vec![0.0; MDEPTH],
            viscd: vec![0.0; MDEPTH],
            abrosd: vec![0.0; MDEPTH],
            abplad: vec![0.0; MDEPTH],
            wmm: vec![1.0; MDEPTH],
            ptotal: vec![0.0; MDEPTH],
        }
    }
}

// ============================================================================
// 主计算函数
// ============================================================================

/// NEWDMT 纯计算函数（网格计算部分）。
///
/// 计算新的深度网格，基于温度变化使用 GRIDP。
///
/// # 参数
///
/// * `config` - 配置参数
/// * `state` - 模型状态（输入/输出）
/// * `prs_aux` - PRSAUX 参数（输出）
/// * `factrs` - FACTRS 参数（输出）
///
/// # 注意
///
/// 完整实现需要调用 TEMPER 和 HESOLV，这里只实现网格计算部分。
pub fn newdmt(
    config: &NewdmtConfig,
    state: &mut NewdmtModelState,
    prs_aux: &mut NewdmtPrsAux,
    factrs: &mut NewdmtFactrs,
) {
    let nd = config.nd;
    let nd1 = nd - 1;

    // 工作数组
    let mut dm0 = vec![0.0; MDEPTH]; // log10(dm) 原始
    let mut dm11 = vec![0.0; MDEPTH]; // 新 log10(dm)
    let mut t0 = vec![0.0; MDEPTH]; // log10(tauros) 原始
    let mut t1 = vec![0.0; MDEPTH]; // 新 log10(tauros)
    let mut elec0 = vec![0.0; MDEPTH]; // 原始 elec
    let mut dens0 = vec![0.0; MDEPTH]; // 原始 dens
    let mut zd0 = vec![0.0; MDEPTH]; // 原始 zd
    let mut pt0 = vec![0.0; MDEPTH]; // 原始 ptotal
    let mut abrs0 = vec![0.0; MDEPTH]; // 原始 abrosd
    let mut abpl0 = vec![0.0; MDEPTH]; // 原始 abplad

    // ========================================================================
    // Step 1: 保存原始值并计算 log10
    // ========================================================================

    for id in 0..nd {
        dm0[id] = if state.dm[id] > 0.0 {
            state.dm[id].log10()
        } else {
            -10.0
        };
        t0[id] = if state.tauros[id] > 0.0 {
            state.tauros[id].log10()
        } else {
            -10.0
        };
        elec0[id] = state.elec[id];
        dens0[id] = state.dens[id];
        pt0[id] = state.ptotal[id];
        zd0[id] = state.zd[id];
        abrs0[id] = state.abrosd[id];
        abpl0[id] = state.abplad[id];
    }

    // ========================================================================
    // Step 2: 使用 GRIDP 计算新网格
    // ========================================================================

    // GRIDP 需要 nd1 个点（不包括最后一个点）
    gridp(&dm0, &t0, &mut dm11, &mut t1, nd1);

    // 最后一个点保持不变
    dm11[nd - 1] = dm0[nd - 1];
    t1[nd - 1] = t0[nd - 1];

    // ========================================================================
    // Step 3: 转换回线性尺度
    // ========================================================================

    for id in 0..nd {
        state.dm[id] = (LN10 * dm11[id]).exp();
        state.tauros[id] = (LN10 * t1[id]).exp();
    }

    // ========================================================================
    // Step 4: 插值各种物理量到新网格
    // ========================================================================

    // 需要可变数组用于 INTERP（INTERP 会修改输入数组）
    let mut dm0_mut = dm0.clone();
    let mut elec0_mut = elec0.clone();
    let mut dm11_mut = dm11.clone();
    let mut elec_out = vec![0.0; MDEPTH];
    interp(
        &mut dm0_mut,
        &mut elec0_mut,
        &mut dm11_mut,
        &mut elec_out,
        nd,
        nd,
        2, // 线性插值
        0, // 不对 x 取对数（已经是对数）
        1, // 对 y 取对数
    );
    for id in 0..nd {
        state.elec[id] = elec_out[id];
    }

    // INTERP(DM0, DENS0, DM11, DENS, ND, ND, 2, 0, 1)
    let mut dm0_mut = dm0.clone();
    let mut dens0_mut = dens0.clone();
    let mut dm11_mut = dm11.clone();
    let mut dens_out = vec![0.0; MDEPTH];
    interp(
        &mut dm0_mut,
        &mut dens0_mut,
        &mut dm11_mut,
        &mut dens_out,
        nd,
        nd,
        2,
        0,
        1,
    );
    for id in 0..nd {
        state.dens[id] = dens_out[id];
    }

    // INTERP(DM0, PT0, DM11, PTOTAL, ND, ND, 2, 0, 1)
    let mut dm0_mut = dm0.clone();
    let mut pt0_mut = pt0.clone();
    let mut dm11_mut = dm11.clone();
    let mut ptotal_out = vec![0.0; MDEPTH];
    interp(
        &mut dm0_mut,
        &mut pt0_mut,
        &mut dm11_mut,
        &mut ptotal_out,
        nd,
        nd,
        2,
        0,
        1,
    );
    for id in 0..nd {
        state.ptotal[id] = ptotal_out[id];
    }

    // INTERP(DM0, ZD0, DM11, ZD, ND, ND, 2, 0, 0)
    let mut dm0_mut = dm0.clone();
    let mut zd0_mut = zd0.clone();
    let mut dm11_mut = dm11.clone();
    let mut zd_out = vec![0.0; MDEPTH];
    interp(
        &mut dm0_mut,
        &mut zd0_mut,
        &mut dm11_mut,
        &mut zd_out,
        nd,
        nd,
        2,
        0,
        0, // 不对 y 取对数
    );
    for id in 0..nd {
        state.zd[id] = zd_out[id];
    }

    // INTERP(DM0, ABRS0, DM11, ABROSD, ND, ND, 2, 0, 1)
    let mut dm0_mut = dm0.clone();
    let mut abrs0_mut = abrs0.clone();
    let mut dm11_mut = dm11.clone();
    let mut abrosd_out = vec![0.0; MDEPTH];
    interp(
        &mut dm0_mut,
        &mut abrs0_mut,
        &mut dm11_mut,
        &mut abrosd_out,
        nd,
        nd,
        2,
        0,
        1,
    );
    for id in 0..nd {
        state.abrosd[id] = abrosd_out[id];
    }

    // INTERP(DM0, ABPL0, DM11, ABPLAD, ND, ND, 2, 0, 1)
    let mut dm0_mut = dm0.clone();
    let mut abpl0_mut = abpl0.clone();
    let mut dm11_mut = dm11.clone();
    let mut abplad_out = vec![0.0; MDEPTH];
    interp(
        &mut dm0_mut,
        &mut abpl0_mut,
        &mut dm11_mut,
        &mut abplad_out,
        nd,
        nd,
        2,
        0,
        1,
    );
    for id in 0..nd {
        state.abplad[id] = abplad_out[id];
    }

    // ========================================================================
    // Step 5: 计算声速平方 VSND2 = PTOTAL / DENS
    // ========================================================================

    for id in 0..nd {
        if state.dens[id] > 0.0 {
            prs_aux.vsnd2[id] = state.ptotal[id] / state.dens[id];
        }
    }

    // ========================================================================
    // Step 6: 计算新的 Rosseland 不透明度和 theta/tauthe 函数
    // ========================================================================

    let amuv0 = if config.dmvisc > 0.0 && state.dm[nd - 1] > 0.0 {
        config.dmvisc.powf(config.zeta0 + UN)
    } else {
        0.0
    };
    let amuv1 = UN - amuv0;

    for id in 0..nd {
        // 计算 viscd 和 theta
        if config.dmvisc > 0.0 && state.dm[nd - 1] > 0.0 {
            let ratio = state.dm[id] / state.dm[nd - 1];
            if ratio <= config.dmvisc {
                // 内部区域：dm <= dmvisc * dm[nd-1]
                state.viscd[id] = (UN - config.fractv) * (config.zeta1 + UN)
                    / config.dmvisc.powf(config.zeta1 + UN)
                    * ratio.powf(config.zeta1);
                state.theta[id] =
                    (UN - config.fractv) * (ratio / config.dmvisc).powf(config.zeta1 + UN);
            } else {
                // 外部区域：dm > dmvisc * dm[nd-1]
                state.viscd[id] =
                    config.fractv * (config.zeta0 + UN) / amuv1 * ratio.powf(config.zeta0);
                state.theta[id] = (UN - config.fractv)
                    + config.fractv * (ratio.powf(config.zeta0 + UN) - amuv0) / amuv1;
            }
        } else {
            // 当 dmvisc = 0 或 dm[nd-1] = 0 时，theta = 1.0（无粘性）
            state.theta[id] = UN;
            state.viscd[id] = 0.0;
        }

        factrs.gamj[id] = UN;

        // 计算 tauros 和 tauthe
        if id == 0 {
            state.tauros[id] = state.dm[id] * state.abrosd[id];
            state.tauthe[id] = state.tauros[id] * state.theta[id] / (config.zeta1 + TWO);
            // 原始代码中计算 ANEREL，但不使用
        } else {
            let ddm = state.dm[id] - state.dm[id - 1];
            state.tauros[id] =
                state.tauros[id - 1] + ddm * HALF * (state.abrosd[id - 1] + state.abrosd[id]);

            let zetad = if config.dmvisc > 0.0
                && state.dm[nd - 1] > 0.0
                && state.dm[id] <= config.dmvisc * state.dm[nd - 1]
            {
                config.zeta1
            } else {
                config.zeta0
            };

            let a0 = (state.abrosd[id - 1] * state.dm[id] - state.abrosd[id] * state.dm[id - 1])
                / ddm
                / (zetad + TWO);
            let a1 =
                (state.abrosd[id] - state.abrosd[id - 1]) / ddm / (zetad + 3.0);

            state.tauthe[id] = state.tauthe[id - 1]
                + a0 * (state.theta[id] * state.dm[id] - state.theta[id - 1] * state.dm[id - 1])
                + a1 * (state.theta[id] * state.dm[id] * state.dm[id]
                    - state.theta[id - 1] * state.dm[id - 1] * state.dm[id - 1]);
        }

        // 注意：原始代码中这里调用 TEMPER
        // CALL TEMPER(ID, TAUR, 1)
    }

    // ========================================================================
    // Step 7: 调用 HESOLV（如果 nconit >= 0）
    // ========================================================================
    // 注意：原始代码中这里调用 HESOLV
    // if(nconit.ge.0) CALL HESOLV

    // ========================================================================
    // Step 8: 再次计算温度和平均不透明度
    // ========================================================================
    // 注意：原始代码中再次调用 TEMPER 和 HESOLV

    // ========================================================================
    // Step 9: 打印输出（如果 ipring >= 1）
    // ========================================================================

    if config.ipring >= 1 {
        eprintln!("\n NEW DEPTH GRID ESTABLISHED, NEW MODEL:");
        eprintln!(" --------------------------------------");
        eprintln!(" ID    DM     TAUROSS   TEMP     NE       P        ZD     ROSS.MEAN  PLANCK");
        for id in 0..nd {
            eprintln!(
                " {:3}{:9.2e}{:9.2e}{:8.0}{:9.2e}{:9.2e}{:9.2e}  {:9.2e}{:9.2e}",
                id + 1,
                state.dm[id],
                state.tauros[id],
                state.temp[id],
                state.elec[id],
                state.ptotal[id],
                state.zd[id],
                state.abrosd[id],
                state.abplad[id]
            );
        }
    }
}

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

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

    fn create_test_state(nd: usize) -> NewdmtModelState {
        let mut state = NewdmtModelState::default();

        // 创建简单的测试网格
        for i in 0..nd {
            let x = (i + 1) as f64 / nd as f64;
            state.dm[i] = 1e-4 * x.powf(2.0);
            state.dens[i] = 1e13 * (1.0 - 0.5 * x);
            state.abrosd[i] = 1e-4 * (1.0 + x);
            state.abplad[i] = 2e-4 * (1.0 + x);
            state.tauros[i] = state.dm[i] * state.abrosd[i];
            state.elec[i] = 1e10 * x;
            state.ptotal[i] = 1e4 * (1.0 + x);
            state.zd[i] = 1e8 * x;
            state.wmm[i] = 1.0;
        }

        state
    }

    #[test]
    fn test_newdmt_basic() {
        let config = NewdmtConfig {
            nd: 50,
            dmvisc: 0.0,
            zeta0: 0.0,
            zeta1: 0.0,
            fractv: 0.0,
            nconit: 0,
            ipring: 0,
        };

        let mut state = create_test_state(config.nd);
        let mut prs_aux = NewdmtPrsAux::default();
        let mut factrs = NewdmtFactrs::default();

        newdmt(&config, &mut state, &mut prs_aux, &mut factrs);

        // 验证基本属性
        assert!(state.dm[0] > 0.0, "dm[0] should be positive");
        assert!(state.tauros[0] >= 0.0, "tauros[0] should be non-negative");

        // 验证 dm 是单调递增的
        for i in 1..config.nd {
            assert!(
                state.dm[i] >= state.dm[i - 1],
                "dm should be monotonically increasing at i={}",
                i
            );
        }

        // 验证 theta = 1.0（无粘性时）
        for i in 0..config.nd {
            assert!(
                (state.theta[i] - 1.0).abs() < 1e-10,
                "theta[{}] should be 1.0 when dmvisc=0",
                i
            );
        }
    }

    #[test]
    fn test_newdmt_with_viscosity() {
        let config = NewdmtConfig {
            nd: 50,
            dmvisc: 0.5,
            zeta0: 0.5,
            zeta1: 1.0,
            fractv: 0.1,
            nconit: 0,
            ipring: 0,
        };

        let mut state = create_test_state(config.nd);
        let mut prs_aux = NewdmtPrsAux::default();
        let mut factrs = NewdmtFactrs::default();

        newdmt(&config, &mut state, &mut prs_aux, &mut factrs);

        // 验证 theta 在有效范围内
        for i in 0..config.nd {
            assert!(
                state.theta[i] >= 0.0,
                "theta[{}] = {} is negative",
                i,
                state.theta[i]
            );
            assert!(
                state.theta[i] <= 2.0,
                "theta[{}] = {} is too large",
                i,
                state.theta[i]
            );
        }

        // 验证 vsnd2 已计算
        for i in 0..config.nd {
            if state.dens[i] > 0.0 {
                assert!(
                    prs_aux.vsnd2[i] > 0.0,
                    "vsnd2[{}] should be positive",
                    i
                );
            }
        }
    }

    #[test]
    fn test_newdmt_tauros_consistency() {
        let config = NewdmtConfig {
            nd: 50,
            dmvisc: 0.0,
            zeta0: 0.0,
            zeta1: 0.0,
            fractv: 0.0,
            nconit: 0,
            ipring: 0,
        };

        let mut state = create_test_state(config.nd);
        let mut prs_aux = NewdmtPrsAux::default();
        let mut factrs = NewdmtFactrs::default();

        newdmt(&config, &mut state, &mut prs_aux, &mut factrs);

        // 验证 tauros 是单调递增的
        for i in 1..config.nd {
            assert!(
                state.tauros[i] >= state.tauros[i - 1],
                "tauros should be monotonically increasing at i={}",
                i
            );
        }

        // 验证 tauthe 是单调递增的
        for i in 1..config.nd {
            assert!(
                state.tauthe[i] >= state.tauthe[i - 1],
                "tauthe should be monotonically increasing at i={}",
                i
            );
        }
    }
}