SpectraRust/src/tlusty/math/solvers/rhsgen.rs

//! 线性化方程组右端向量计算。
//!
//! 重构自 TLUSTY `rhsgen.f`。
//!
//! 功能：
//! - 计算辐射传输组件
//! - 流体静力学平衡
//! - 辐射平衡
//! - 电荷守恒
//! - 对流贡献

use crate::tlusty::math::{convec, ConvecConfig, ConvecParams};
use crate::tlusty::state::constants::{BOLK, HALF, UN};

/// 常量
const XCON: f64 = 8.0935e-21;
const YCON: f64 = 1.68638e-10;
const SIXTH: f64 = 1.0 / 6.0;
const THIRD: f64 = 1.0 / 3.0;

/// RHSGEN 配置参数
#[derive(Debug, Clone)]
pub struct RhsgenConfig {
    /// 插值方法 (ISPLIN)
    pub isplin: i32,
    /// 盘模式标志 (IDISK)
    pub idisk: i32,
    /// 上边界条件 (IBC)
    pub ibc: i32,
    /// 氦方程标志 (INHE)
    pub inhe: i32,
    /// 能量方程标志 (INRE)
    pub inre: i32,
    /// 压力方程标志 (INPC)
    pub inpc: i32,
    /// DELTA 方程标志 (INDL)
    pub indl: i32,
    /// 统计平衡标志 (INSE)
    pub inse: i32,
    /// z-d 关系标志 (INZD)
    pub inzd: i32,
    /// 频率数 (NFREQE)
    pub nfreqe: usize,
    /// 总方程数 (NN)
    pub nn: usize,
    /// 不透明度缩放标志 (IZSCAL)
    pub izscal: i32,
    /// Compton 散射标志 (ICOMPT)
    pub icompt: i32,
    /// 混合长度参数 (HMIX0)
    pub hmix0: f64,
    /// 对流模式标志 (ICONV)
    pub iconv: i32,
    /// 有效温度 (TEFF)
    pub teff: f64,
    /// σTeff⁴/π (SIG4P)
    pub sig4p: f64,
    /// 压力常数 (PCK)
    pub pck: f64,
    /// 重力加速度 (GRAV)
    pub grav: f64,
    /// 重力缩放因子 (QGRAV)
    pub qgrav: f64,
}

impl Default for RhsgenConfig {
    fn default() -> Self {
        Self {
            isplin: 0,
            idisk: 0,
            ibc: 1,
            inhe: 1,
            inre: 1,
            inpc: 0,
            indl: 0,
            inse: 0,
            inzd: 0,
            nfreqe: 10,
            nn: 15,
            izscal: 0,
            icompt: 0,
            hmix0: -1.0,
            iconv: 0,
            teff: 10000.0,
            sig4p: 5.67e-5 / 3.14159265359,
            pck: 1.0,
            grav: 1e4,
            qgrav: 1e4,
        }
    }
}

/// RHSGEN 频率数据
pub struct RhsgenFreqData<'a> {
    /// 频率权重
    pub w: &'a [f64],
    /// 辐射强度
    pub rad: &'a [f64],
    /// 吸收系数
    pub abso: &'a [f64],
    /// 发射系数
    pub emis: &'a [f64],
    /// 散射系数
    pub scat: &'a [f64],
    /// FK 系数
    pub fk: &'a [f64],
}

/// RHSGEN 输入参数
pub struct RhsgenParams<'a> {
    /// 深度点索引 (1-indexed)
    pub id: usize,
    /// 总深度点数
    pub nd: usize,
    /// 温度数组
    pub temp: &'a [f64],
    /// 密度数组
    pub dens: &'a [f64],
    /// 柱密度数组
    pub dm: &'a [f64],
    /// 平均分子量数组
    pub wmm: &'a [f64],
    /// 湍流速度数组
    pub vturb: &'a [f64],
    /// 电子密度数组
    pub elec: &'a [f64],
    /// 几何因子数组
    pub zd: &'a [f64],
    /// Delta 参数数组
    pub delta: &'a mut [f64],
    /// 对流通量数组
    pub flxc: &'a mut [f64],
    /// 冷却通量数组
    pub fcool: &'a [f64],
    /// Rosseland 不透明度数组
    pub abrosd: &'a [f64],
    /// 微分方程权重
    pub redif: &'a [f64],
    /// 积分方程权重
    pub reint: &'a [f64],
    /// 当前点频率数据
    pub freq0: &'a RhsgenFreqData<'a>,
    /// 上一点频率数据
    pub freqm: &'a RhsgenFreqData<'a>,
    /// 下一点频率数据
    pub freqp: &'a RhsgenFreqData<'a>,
    /// 配置
    pub config: RhsgenConfig,
    /// CONVEC 配置
    pub convec_config: ConvecConfig,
}

/// RHSGEN 输出
pub struct RhsgenOutput {
    /// RHS 向量
    pub vecl: Vec<f64>,
}

/// 计算 RHS 向量。
pub fn rhsgen(params: &mut RhsgenParams) -> RhsgenOutput {
    let id = params.id;
    let nd = params.nd;

    // 提取配置值（避免借用冲突）
    let nfreqe = params.config.nfreqe;
    let nn = params.config.nn;
    let inhe = params.config.inhe;
    let inre = params.config.inre;
    let inpc = params.config.inpc;
    let indl = params.config.indl;
    let hmix0 = params.config.hmix0;

    // 初始化 RHS 向量
    let mut vecl = vec![0.0; nn];

    // 计算行索引
    let nhe = nfreqe + inhe as usize;
    let _nre = nfreqe + inre as usize;
    let npc = nfreqe + inpc as usize;
    let _ndel = nfreqe + indl as usize;

    // 辐射传输组件
    if nfreqe > 0 {
        if id == 1 {
            // 上边界条件
            compute_upper_boundary(params, &mut vecl);
        } else if id < nd {
            // 内部点
            compute_interior_point(params, &mut vecl);
        } else {
            // 下边界条件
            compute_lower_boundary(params, &mut vecl);
        }
    }

    // 流体静力学平衡
    if inhe > 0 && nhe < vecl.len() {
        compute_hydrostatic(params, &mut vecl, nhe);
    }

    // 电荷守恒
    if inpc > 0 && npc < vecl.len() {
        // 简化实现
        vecl[npc] = 0.0;
    }

    // 对流贡献
    if hmix0 > 0.0 && id > 1 && id < nd {
        compute_convection(params, &mut vecl);
    }

    RhsgenOutput { vecl }
}

/// 计算上边界条件
fn compute_upper_boundary(params: &mut RhsgenParams, vecl: &mut [f64]) {
    let cfg = &params.config;
    let ddp = 1e5; // DELDMZ(1) 简化值

    for ij in 0..cfg.nfreqe.min(params.freq0.abso.len()) {
        let abso0 = params.freq0.abso[ij];
        let absop = params.freqp.abso[ij];
        let dens = params.dens[0];

        let omeg0 = if cfg.izscal == 0 { abso0 / dens } else { abso0 };
        let omegp = if cfg.izscal == 0 { absop / dens } else { absop };

        let dzp = omeg0 + omegp;
        let dtaup = dzp * ddp;

        let fk0 = if ij < params.freq0.fk.len() { params.freq0.fk[ij] } else { 1.0 };
        let fkp = if ij < params.freqp.fk.len() { params.freqp.fk[ij] } else { 1.0 };
        let rad0 = if ij < params.freq0.rad.len() { params.freq0.rad[ij] } else { 0.0 };
        let radp = if ij < params.freqp.rad.len() { params.freqp.rad[ij] } else { 0.0 };

        let alf1 = (fk0 * rad0 - fkp * radp) / dtaup;

        let scat0 = if ij < params.freq0.scat.len() { params.freq0.scat[ij] } else { 0.0 };
        let emis0 = if ij < params.freq0.emis.len() { params.freq0.emis[ij] } else { 0.0 };

        let s0 = if abso0.abs() > 1e-30 {
            (emis0 + scat0 * rad0) / abso0
        } else {
            0.0
        };

        let bs = HALF * dtaup;
        let alf2 = bs * (rad0 - s0);

        if ij < vecl.len() {
            vecl[ij] = alf1 + alf2;
        }
    }
}

/// 计算内部点
fn compute_interior_point(params: &mut RhsgenParams, vecl: &mut [f64]) {
    let id = params.id;
    let cfg = &params.config;

    let ddm = 1e5; // 简化值
    let ddp = 1e5; // 简化值

    for ij in 0..cfg.nfreqe.min(params.freq0.abso.len()) {
        let dens_id = params.dens[id - 1];
        let dens_im = params.dens[id - 2];
        let dens_ip = params.dens[id];

        let abso0 = params.freq0.abso[ij];
        let absom = params.freqm.abso[ij];
        let absop = params.freqp.abso[ij];

        let omeg0 = if cfg.izscal == 0 { abso0 / dens_id } else { abso0 };
        let omegm = if cfg.izscal == 0 { absom / dens_im } else { absom };
        let omegp = if cfg.izscal == 0 { absop / dens_ip } else { absop };

        let dzp = omeg0 + omegp;
        let dzm = omeg0 + omegm;
        let dtaup = dzp * ddp;
        let dtaum = dzm * ddm;
        let dtau0 = HALF * (dtaup + dtaum);

        let fk0 = if ij < params.freq0.fk.len() { params.freq0.fk[ij] } else { 1.0 };
        let fkm = if ij < params.freqm.fk.len() { params.freqm.fk[ij] } else { 1.0 };
        let fkp = if ij < params.freqp.fk.len() { params.freqp.fk[ij] } else { 1.0 };
        let rad0 = if ij < params.freq0.rad.len() { params.freq0.rad[ij] } else { 0.0 };
        let radm = if ij < params.freqm.rad.len() { params.freqm.rad[ij] } else { 0.0 };
        let radp = if ij < params.freqp.rad.len() { params.freqp.rad[ij] } else { 0.0 };

        let frd = fk0 * rad0;
        let alf1 = (frd - fkp * radp) / dtaup / dtau0;
        let gam1 = (frd - fkm * radm) / dtaum / dtau0;
        let bet1 = alf1 + gam1;

        let scat0 = if ij < params.freq0.scat.len() { params.freq0.scat[ij] } else { 0.0 };
        let emis0 = if ij < params.freq0.emis.len() { params.freq0.emis[ij] } else { 0.0 };

        let s0 = if abso0.abs() > 1e-30 {
            (emis0 + scat0 * rad0) / abso0
        } else {
            0.0
        };

        let bet2 = UN * (rad0 - s0);

        if ij < vecl.len() {
            vecl[ij] = bet1 + bet2;
        }
    }
}

/// 计算下边界条件
fn compute_lower_boundary(params: &mut RhsgenParams, vecl: &mut [f64]) {
    let id = params.id;
    let cfg = &params.config;

    let t = params.temp[id - 1];
    let ddm = 1e5; // 简化值

    for ij in 0..cfg.nfreqe.min(params.freq0.abso.len()) {
        let dens_id = params.dens[id - 1];
        let dens_im = params.dens[id - 2];

        let abso0 = params.freq0.abso[ij];
        let absom = params.freqm.abso[ij];

        let omeg0 = if cfg.izscal == 0 { abso0 / dens_id } else { abso0 };
        let omegm = if cfg.izscal == 0 { absom / dens_im } else { absom };

        let dzm = omeg0 + omegm;
        let dtaum = dzm * ddm;

        let fk0 = if ij < params.freq0.fk.len() { params.freq0.fk[ij] } else { 1.0 };
        let fkm = if ij < params.freqm.fk.len() { params.freqm.fk[ij] } else { 1.0 };
        let rad0 = if ij < params.freq0.rad.len() { params.freq0.rad[ij] } else { 0.0 };
        let radm = if ij < params.freqm.rad.len() { params.freqm.rad[ij] } else { 0.0 };

        let gam1 = (fk0 * rad0 - fkm * radm) / dtaum;

        // Planck 函数简化
        let plan = 0.0; // 需要实际计算

        if ij < vecl.len() {
            vecl[ij] = gam1 - HALF * (plan - rad0);
        }
    }
}

/// 计算流体静力学平衡
fn compute_hydrostatic(params: &mut RhsgenParams, vecl: &mut [f64], nhe: usize) {
    let id = params.id;
    let cfg = &params.config;

    if id == 1 {
        // 上边界条件
        let mut grd = 0.0;
        if cfg.nfreqe > 0 {
            for ij in 0..cfg.nfreqe.min(params.freq0.abso.len()) {
                let w = if ij < params.freq0.w.len() { params.freq0.w[ij] } else { 1.0 };
                let rad0 = if ij < params.freq0.rad.len() { params.freq0.rad[ij] } else { 0.0 };
                let abso0 = params.freq0.abso[ij];
                grd += w * rad0 * abso0;
            }
        }

        let x1 = cfg.pck / params.dens[0];
        let vt0 = HALF * params.vturb[0].powi(2) / params.dm[0] * params.wmm[0];

        if nhe < vecl.len() {
            vecl[nhe] = cfg.grav - BOLK * params.temp[0] * 0.0 / params.dm[0]
                - x1 * grd - vt0 / params.wmm[0] * params.dens[0];
        }
    } else {
        // 内部点
        let mut grd = 0.0;
        if cfg.nfreqe > 0 {
            for ij in 0..cfg.nfreqe.min(params.freq0.abso.len()) {
                let w = if ij < params.freq0.w.len() { params.freq0.w[ij] } else { 1.0 };
                let fk0 = if ij < params.freq0.fk.len() { params.freq0.fk[ij] } else { 1.0 };
                let fkm = if ij < params.freqm.fk.len() { params.freqm.fk[ij] } else { 1.0 };
                let rad0 = if ij < params.freq0.rad.len() { params.freq0.rad[ij] } else { 0.0 };
                let radm = if ij < params.freqm.rad.len() { params.freqm.rad[ij] } else { 0.0 };
                grd += (fk0 * rad0 - fkm * radm) * w;
            }
        }

        let vt0 = HALF * params.vturb[id - 1].powi(2) * params.wmm[id - 1];
        let vtm = HALF * params.vturb[id - 2].powi(2) * params.wmm[id - 2];

        if nhe < vecl.len() {
            vecl[nhe] = cfg.grav * (params.dm[id - 1] - params.dm[id - 2])
                - BOLK * (params.temp[id - 1] - params.temp[id - 2])
                - cfg.pck * grd
                - vt0 / params.wmm[id - 1] * params.dens[id - 1]
                + vtm / params.wmm[id - 2] * params.dens[id - 2];
        }
    }
}

/// 计算对流贡献
fn compute_convection(params: &mut RhsgenParams, vecl: &mut [f64]) {
    let id = params.id;
    let cfg = &params.config;

    let t = params.temp[id - 1];
    let tm = params.temp[id - 2];

    let t0 = HALF * (t + tm);
    let dlt = if t0.abs() > 1e-30 { (t - tm) / t0 } else { 0.0 };

    params.delta[id - 1] = dlt;

    // 调用 CONVEC
    let convec_params = ConvecParams {
        id,
        t: t0,
        ptot: 1e5,
        pg: 1e5,
        prad: 0.0,
        abros: 0.4,
        delta: dlt,
        taurs: 0.0,
        config: params.convec_config.clone(),
        trmder_config: None,
        therm_tables: None,
    };

    let convec_out = convec(&convec_params);
    let flxcnv = convec_out.flxcnv;
    params.flxc[id - 1] = flxcnv;

    // 更新 RHS 向量
    let nre = cfg.nfreqe + cfg.inre as usize;
    if params.redif[id - 1] > 0.0 && nre < vecl.len() {
        vecl[nre] -= flxcnv * params.redif[id - 1];
    }
}

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

    fn create_test_freq_data(n: usize) -> (Vec<f64>, Vec<f64>, Vec<f64>, Vec<f64>, Vec<f64>, Vec<f64>) {
        let w = vec![1.0; n];
        let rad = vec![1e10; n];
        let abso = vec![0.1; n];
        let emis = vec![1e9; n];
        let scat = vec![0.01; n];
        let fk = vec![1.0; n];
        (w, rad, abso, emis, scat, fk)
    }

    #[test]
    fn test_rhsgen_upper_boundary() {
        let nd = 5;
        let (w0, rad0, abso0, emis0, scat0, fk0) = create_test_freq_data(10);
        let (wp, radp, absop, emisp, scatp, fkp) = create_test_freq_data(10);

        let freq0 = RhsgenFreqData {
            w: &w0, rad: &rad0, abso: &abso0, emis: &emis0, scat: &scat0, fk: &fk0,
        };
        let freqm = RhsgenFreqData {
            w: &[], rad: &[], abso: &[], emis: &[], scat: &[], fk: &[],
        };
        let freqp = RhsgenFreqData {
            w: &wp, rad: &radp, abso: &absop, emis: &emisp, scat: &scatp, fk: &fkp,
        };

        let temp = vec![10000.0, 9500.0, 9000.0, 8500.0, 8000.0];
        let dens = vec![1e-7; nd];
        let dm = vec![1e2; nd];
        let wmm = vec![1.4e-24; nd];
        let vturb = vec![2e5; nd];
        let elec = vec![1e-8; nd];
        let zd = vec![1.0; nd];
        let mut delta = vec![0.0; nd];
        let mut flxc = vec![0.0; nd];
        let fcool = vec![0.0; nd];
        let abrosd = vec![0.4; nd];
        let redif = vec![1.0; nd];
        let reint = vec![0.0; nd];

        let mut params = RhsgenParams {
            id: 1,
            nd,
            temp: &temp,
            dens: &dens,
            dm: &dm,
            wmm: &wmm,
            vturb: &vturb,
            elec: &elec,
            zd: &zd,
            delta: &mut delta,
            flxc: &mut flxc,
            fcool: &fcool,
            abrosd: &abrosd,
            redif: &redif,
            reint: &reint,
            freq0: &freq0,
            freqm: &freqm,
            freqp: &freqp,
            config: RhsgenConfig {
                hmix0: -1.0, // 禁用对流
                ..Default::default()
            },
            convec_config: ConvecConfig::default(),
        };

        let result = rhsgen(&mut params);

        assert!(result.vecl.len() > 0);
        // 上边界时对流应被禁用
        assert_eq!(params.flxc[0], 0.0);
    }

    #[test]
    fn test_rhsgen_interior_point() {
        let nd = 5;
        let (w0, rad0, abso0, emis0, scat0, fk0) = create_test_freq_data(10);
        let (wm, radm, absom, emism, scatm, fkm) = create_test_freq_data(10);
        let (wp, radp, absop, emisp, scatp, fkp) = create_test_freq_data(10);

        let freq0 = RhsgenFreqData {
            w: &w0, rad: &rad0, abso: &abso0, emis: &emis0, scat: &scat0, fk: &fk0,
        };
        let freqm = RhsgenFreqData {
            w: &wm, rad: &radm, abso: &absom, emis: &emism, scat: &scatm, fk: &fkm,
        };
        let freqp = RhsgenFreqData {
            w: &wp, rad: &radp, abso: &absop, emis: &emisp, scat: &scatp, fk: &fkp,
        };

        let temp = vec![10000.0, 9500.0, 9000.0, 8500.0, 8000.0];
        let dens = vec![1e-7; nd];
        let dm = vec![1e2; nd];
        let wmm = vec![1.4e-24; nd];
        let vturb = vec![2e5; nd];
        let elec = vec![1e-8; nd];
        let zd = vec![1.0; nd];
        let mut delta = vec![0.0; nd];
        let mut flxc = vec![0.0; nd];
        let fcool = vec![0.0; nd];
        let abrosd = vec![0.4; nd];
        let redif = vec![1.0; nd];
        let reint = vec![0.0; nd];

        let mut params = RhsgenParams {
            id: 3, // 内部点
            nd,
            temp: &temp,
            dens: &dens,
            dm: &dm,
            wmm: &wmm,
            vturb: &vturb,
            elec: &elec,
            zd: &zd,
            delta: &mut delta,
            flxc: &mut flxc,
            fcool: &fcool,
            abrosd: &abrosd,
            redif: &redif,
            reint: &reint,
            freq0: &freq0,
            freqm: &freqm,
            freqp: &freqp,
            config: RhsgenConfig {
                hmix0: -1.0,
                ..Default::default()
            },
            convec_config: ConvecConfig::default(),
        };

        let result = rhsgen(&mut params);

        assert!(result.vecl.len() > 0);
    }
}