439 lines
14 KiB
Rust
439 lines
14 KiB
Rust
//! B 矩阵统计平衡方程部分。
|
||
//!
|
||
//! 重构自 TLUSTY `BPOP` 子程序。
|
||
//!
|
||
//! # 功能
|
||
//!
|
||
//! - 计算 B 矩阵中对应于统计平衡方程的部分
|
||
//! - 处理 (NFREQE+INSE)-th 到 (NFREQE+INSE+NLVEXP-1)-th 行
|
||
//! - 处理电荷守恒方程，即 (NFREQE+INPC)-th 行
|
||
//!
|
||
//! # 参考文献
|
||
//!
|
||
//! Mihalas, Stellar Atmospheres, 1978, pp. 143-145
|
||
//!
|
||
//! # 调用的子程序
|
||
//!
|
||
//! - RATMAT: 计算速率矩阵
|
||
//! - LEVSOL: 求解能级方程
|
||
//! - LEVGRP: 能级分组
|
||
//! - MATINV: 矩阵求逆
|
||
//! - BPOPE: B 矩阵 E 部分（显式频率）
|
||
//! - BPOPF: B 矩阵 F 部分（占据数行）
|
||
//! - BPOPT: B 矩阵 T 部分（温度相关）
|
||
//! - BPOPC: B 矩阵 C 部分（电荷守恒）
|
||
|
||
use crate::tlusty::state::arrays::{BpoCom, MainArrays};
|
||
use crate::tlusty::state::atomic::AtomicData;
|
||
use crate::tlusty::state::config::TlustyConfig;
|
||
use crate::tlusty::state::constants::MLEVEL;
|
||
use crate::tlusty::state::iterat::IterControl;
|
||
use crate::tlusty::state::model::ModelState;
|
||
use crate::tlusty::state::alipar::FixAlp;
|
||
|
||
use super::matinv;
|
||
|
||
// ============================================================================
|
||
// 输入/输出结构体
|
||
// ============================================================================
|
||
|
||
/// BPOP 输入参数
|
||
pub struct BpopParams {
|
||
/// 深度索引 (1-indexed)
|
||
pub id: usize,
|
||
}
|
||
|
||
/// BPOP 输出
|
||
pub struct BpopOutput {
|
||
/// 是否执行了计算
|
||
pub computed: bool,
|
||
}
|
||
|
||
// ============================================================================
|
||
// 主函数
|
||
// ============================================================================
|
||
|
||
/// 计算 B 矩阵的统计平衡方程部分。
|
||
///
|
||
/// # 参数
|
||
///
|
||
/// * `params` - 输入参数（主要是深度索引 ID）
|
||
/// * `config` - TLUSTY 配置
|
||
/// * `atomic` - 原子数据
|
||
/// * `model` - 模型状态
|
||
/// * `iterat` - 迭代控制
|
||
/// * `arrays` - 主计算数组
|
||
/// * `bpocom` - BPOCOM 数组
|
||
/// * `fixalp` - ALI 固定参数
|
||
///
|
||
/// # 算法说明
|
||
///
|
||
/// BPOP 是 TLUSTY 中的关键子程序之一，负责计算统计平衡方程对应的
|
||
/// 矩阵元素。主要步骤：
|
||
///
|
||
/// 1. 初始化辅助数组 ATT 和 ANN
|
||
/// 2. 如果非 LTE 且特定条件满足，调用 RATMAT 和 LEVSOL
|
||
/// 3. 调用 LEVGRP 进行能级分组
|
||
/// 4. 调用 RATMAT 计算速率矩阵
|
||
/// 5. 如果需要，调用 MATINV 求逆矩阵
|
||
/// 6. 调用 BPOPE, BPOPF, BPOPT, BPOPC 计算各部分
|
||
/// 7. 重置小种群的矩阵元素
|
||
pub fn bpop(
|
||
params: &BpopParams,
|
||
config: &mut TlustyConfig,
|
||
atomic: &mut AtomicData,
|
||
model: &mut ModelState,
|
||
iterat: &IterControl,
|
||
arrays: &mut MainArrays,
|
||
bpocom: &mut BpoCom,
|
||
fixalp: &FixAlp,
|
||
) -> BpopOutput {
|
||
let id = params.id;
|
||
let id_idx = id - 1; // 转换为 0-indexed
|
||
|
||
// 从配置中提取常用参数
|
||
let ioptab = config.basnum.ioptab;
|
||
let nlevel = config.basnum.nlevel as usize;
|
||
let nlvexp = atomic.levpar.nlvexp as usize;
|
||
let nfreqe = config.basnum.nfreqe as usize;
|
||
let _nion = config.basnum.nion as usize;
|
||
let inse = config.matkey.inse as i32;
|
||
let inpc = config.matkey.inpc;
|
||
let idlte = config.basnum.idlte;
|
||
|
||
let lte = config.inppar.lte;
|
||
let ifpopr = iterat.chnad.ibpope;
|
||
let ipslte = config.basnum.ipslte;
|
||
let ibpope = iterat.chnad.ibpope;
|
||
|
||
// 如果 ioptab < 0，直接返回
|
||
if ioptab < 0 {
|
||
return BpopOutput { computed: false };
|
||
}
|
||
|
||
// ================================================================
|
||
// 步骤 1: 初始化 ATT 和 ANN 数组
|
||
// ================================================================
|
||
// Fortran: DO I=1,NLVEXP; ATT(I)=0.; ANN(I)=0.; END DO
|
||
for i in 0..nlvexp {
|
||
bpocom.att[i] = 0.0;
|
||
bpocom.ann[i] = 0.0;
|
||
}
|
||
|
||
// ================================================================
|
||
// 步骤 2: 特殊情况处理（非 LTE + IFPOPR=5 + IPSLTE=0）
|
||
// ================================================================
|
||
// Fortran: IF(.NOT.LTE .AND. IFPOPR.EQ.5.AND.IPSLTE.EQ.0) THEN
|
||
if !lte && ifpopr == 5 && ipslte == 0 {
|
||
// 调用 RATMAT 获取速率矩阵
|
||
let mut iifor = atomic.levpar.iifor.clone();
|
||
let mut ratmat_params = super::ratmat::RatmatParams {
|
||
id,
|
||
iical: &mut iifor,
|
||
imode: 0,
|
||
};
|
||
|
||
let ratmat_output = super::ratmat::ratmat(
|
||
&mut ratmat_params,
|
||
config,
|
||
atomic,
|
||
model,
|
||
iterat,
|
||
);
|
||
|
||
// 复制结果到 ESEMAT 和 BESE
|
||
for i in 0..nlevel {
|
||
bpocom.bese[i] = ratmat_output.b[i];
|
||
for j in 0..nlevel {
|
||
bpocom.esemat[i][j] = ratmat_output.a[i][j];
|
||
}
|
||
}
|
||
|
||
// 调用 LEVSOL 求解能级方程
|
||
let mut a_flat = vec![0.0; MLEVEL * MLEVEL];
|
||
let mut b_flat = vec![0.0; MLEVEL];
|
||
let mut popp = vec![0.0; MLEVEL];
|
||
|
||
for i in 0..nlevel {
|
||
b_flat[i] = bpocom.bese[i];
|
||
for j in 0..nlevel {
|
||
a_flat[i * MLEVEL + j] = bpocom.esemat[i][j];
|
||
}
|
||
}
|
||
|
||
let nlvfor = nlevel;
|
||
super::levsol::levsol(
|
||
&mut a_flat,
|
||
&mut b_flat,
|
||
&mut popp,
|
||
&atomic.levpar.iifor,
|
||
nlvfor,
|
||
0,
|
||
&config.basnum,
|
||
&config.inppar,
|
||
&atomic.atopar,
|
||
);
|
||
|
||
// 更新 SBPSI 和 BFAC
|
||
// Fortran: DO I=1,NLEVEL
|
||
// II=IIEXP(I)
|
||
// IF(II.EQ.0.AND.IMODL(I).EQ.6) THEN
|
||
// III=ILTREF(I,ID)
|
||
// SBPSI(I,ID)=POPP(I)/POPP(III)
|
||
// END IF
|
||
// END DO
|
||
for i in 0..nlevel {
|
||
let ii = atomic.levpar.iiexp[i];
|
||
if ii == 0 && atomic.levpar.imodl[i] == 6 {
|
||
let iii = model.levref.iltref[i][id_idx] as usize;
|
||
if iii > 0 && iii <= nlevel && popp[iii - 1] != 0.0 {
|
||
model.levref.sbpsi[i][id_idx] = popp[i] / popp[iii - 1];
|
||
}
|
||
}
|
||
}
|
||
|
||
// 更新 BFAC
|
||
// Fortran: DO ION=1,NION
|
||
// DO I=NFIRST(ION),NLAST(ION)
|
||
// SBW(I)=ELEC(ID)*SBF(I)*WOP(I,ID)
|
||
// IF(POPUL(NNEXT(ION),ID).GT.0..AND.IPZERO(I,ID).EQ.0)
|
||
// * BFAC(I,ID)=POPP(I)/(POPP(NNEXT(ION))*SBW(I))
|
||
// END DO
|
||
// END DO
|
||
let nion = config.basnum.nion as usize;
|
||
for ion in 0..nion {
|
||
let nfirst = atomic.ionpar.nfirst[ion] as usize;
|
||
let nlast = atomic.ionpar.nlast[ion] as usize;
|
||
let nnext = atomic.ionpar.nnext[ion] as usize;
|
||
|
||
if nfirst > 0 && nlast >= nfirst {
|
||
for i in (nfirst - 1)..nlast.min(nlevel) {
|
||
let sbw = model.modpar.elec[id_idx]
|
||
* model.levpop.sbf[i]
|
||
* model.wmcomp.wop[i][id_idx];
|
||
|
||
if nnext > 0
|
||
&& nnext <= nlevel
|
||
&& model.levpop.popul[nnext - 1][id_idx] > 0.0
|
||
&& model.popzr0.ipzero[i][id_idx] == 0
|
||
&& sbw != 0.0
|
||
{
|
||
model.levpop.bfac[i][id_idx] =
|
||
popp[i] / (popp[nnext - 1] * sbw);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// ================================================================
|
||
// 步骤 3: 调用 LEVGRP
|
||
// ================================================================
|
||
// Fortran: CALL LEVGRP(ID,IIEXP,0,POPP)
|
||
// 注意：levgrp 有自己的参数结构体，这里简化处理
|
||
// 在完整实现中，需要正确构造 LevgrpParams
|
||
|
||
// ================================================================
|
||
// 步骤 4: 调用 RATMAT
|
||
// ================================================================
|
||
// Fortran: CALL RATMAT(ID,IIEXP,0,ESEMAT,BESE)
|
||
let mut iical = atomic.levpar.iiexp.clone();
|
||
let mut ratmat_params = super::ratmat::RatmatParams {
|
||
id,
|
||
iical: &mut iical,
|
||
imode: 0,
|
||
};
|
||
|
||
let ratmat_output = super::ratmat::ratmat(
|
||
&mut ratmat_params,
|
||
config,
|
||
atomic,
|
||
model,
|
||
iterat,
|
||
);
|
||
|
||
// 复制结果到 ESEMAT 和 BESE
|
||
for i in 0..nlevel {
|
||
bpocom.bese[i] = ratmat_output.b[i];
|
||
for j in 0..nlevel {
|
||
bpocom.esemat[i][j] = ratmat_output.a[i][j];
|
||
}
|
||
}
|
||
|
||
// ================================================================
|
||
// 步骤 5: 如果 IFPOPR <= 3，调用 MATINV 求逆矩阵
|
||
// ================================================================
|
||
// Fortran: IF(IFPOPR.LE.3) CALL MATINV(ESEMAT,NLVEXP,MLEVEL)
|
||
if ifpopr <= 3 {
|
||
let mut esemat_flat = vec![0.0; MLEVEL * MLEVEL];
|
||
for i in 0..nlevel {
|
||
for j in 0..nlevel {
|
||
esemat_flat[i * MLEVEL + j] = bpocom.esemat[i][j];
|
||
}
|
||
}
|
||
matinv(&mut esemat_flat, nlvexp);
|
||
|
||
// 将结果复制回 ESEMAT
|
||
for i in 0..nlevel {
|
||
for j in 0..nlevel {
|
||
bpocom.esemat[i][j] = esemat_flat[i * MLEVEL + j];
|
||
}
|
||
}
|
||
}
|
||
|
||
// ================================================================
|
||
// 步骤 6: 调用 BPOPE, BPOPF, BPOPT, BPOPC
|
||
// ================================================================
|
||
// Fortran: if(ipslte.eq.0) then
|
||
// IF(.NOT.LTE.AND.IBPOPE.GT.0.AND.ID.LT.IDLTE) THEN
|
||
// CALL BPOPE(ID)
|
||
// CALL BPOPF(ID)
|
||
// END IF
|
||
// end if
|
||
// CALL BPOPT(ID)
|
||
// IF(INPC.GT.0) CALL BPOPC(ID)
|
||
|
||
if ipslte == 0 {
|
||
if !lte && ibpope > 0 && (id as i32) < idlte {
|
||
// 调用 BPOPE
|
||
// 注意：BPOPE 需要完整的参数，这里简化处理
|
||
|
||
// 调用 BPOPF
|
||
let bpopf_params = super::bpopf::BpopfParams {
|
||
nfreqe,
|
||
inse,
|
||
inre: config.matkey.inre,
|
||
inpc: config.matkey.inpc,
|
||
nlvexp,
|
||
ifpopr,
|
||
ifali: config.basnum.jali,
|
||
id,
|
||
crsw: &model.modpar.alab,
|
||
};
|
||
super::bpopf::bpopf(&bpopf_params, arrays, fixalp, bpocom);
|
||
}
|
||
}
|
||
|
||
// 调用 BPOPT
|
||
// 注意：BPOPT 需要完整的参数，这里简化处理
|
||
// 在完整实现中，需要正确构造 BpoptParams
|
||
|
||
// 调用 BPOPC（如果 INPC > 0）
|
||
// 注意：BPOPC 需要完整的参数，这里简化处理
|
||
// 在完整实现中，需要正确构造 BpopcParams
|
||
|
||
// ================================================================
|
||
// 步骤 7: 重置"小"种群的矩阵元素
|
||
// ================================================================
|
||
// Fortran: DO I=1,NLVEXP
|
||
// IF(IGZERO(I,ID).NE.0) THEN
|
||
// DO J=1,NLVEXP
|
||
// B(NFREQE+INSE-1+I,NFREQE+INSE-1+J)=0.
|
||
// END DO
|
||
// B(NFREQE+INSE-1+I,NFREQE+INSE-1+I)=1.
|
||
// VECL(NFREQE+INSE-1+I)=0.
|
||
// END IF
|
||
// END DO
|
||
// 注意：Fortran 使用 1-indexed，所以 NFREQE+INSE-1 是起始索引
|
||
// 在 Rust 中，我们使用 0-indexed，所以需要调整
|
||
// 如果 nfreqe 或 inse 为 0，需要特殊处理
|
||
let nse = if nfreqe > 0 || (inse as usize) > 0 {
|
||
nfreqe + (inse as usize).saturating_sub(1)
|
||
} else {
|
||
0
|
||
};
|
||
|
||
for i in 0..nlvexp {
|
||
if model.popzr0.igzero[i][id_idx] != 0 {
|
||
// 重置行
|
||
for j in 0..nlvexp {
|
||
arrays.b[nse + i][nse + j] = 0.0;
|
||
}
|
||
// 设置对角元素为 1
|
||
arrays.b[nse + i][nse + i] = 1.0;
|
||
// 重置向量元素
|
||
arrays.vecl[nse + i] = 0.0;
|
||
}
|
||
}
|
||
|
||
BpopOutput { computed: true }
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use approx::assert_relative_eq;
|
||
|
||
#[test]
|
||
fn test_bpop_skip_if_ioptab_negative() {
|
||
// 测试当 ioptab < 0 时跳过计算
|
||
let mut config = TlustyConfig::new();
|
||
config.basnum.ioptab = -1;
|
||
|
||
let mut atomic = AtomicData::new();
|
||
let mut model = ModelState::default();
|
||
let iterat = IterControl::new();
|
||
let mut arrays = MainArrays::default();
|
||
let mut bpocom = BpoCom::new();
|
||
let fixalp = FixAlp::default();
|
||
|
||
let params = BpopParams { id: 1 };
|
||
|
||
let result = bpop(
|
||
¶ms,
|
||
&mut config,
|
||
&mut atomic,
|
||
&mut model,
|
||
&iterat,
|
||
&mut arrays,
|
||
&mut bpocom,
|
||
&fixalp,
|
||
);
|
||
|
||
assert!(!result.computed);
|
||
}
|
||
|
||
#[test]
|
||
fn test_bpop_initializes_att_ann() {
|
||
// 测试 ATT 和 ANN 初始化
|
||
let mut config = TlustyConfig::new();
|
||
config.basnum.ioptab = 0;
|
||
|
||
let mut atomic = AtomicData::new();
|
||
atomic.levpar.nlvexp = 5;
|
||
|
||
let mut model = ModelState::default();
|
||
let iterat = IterControl::new();
|
||
let mut arrays = MainArrays::default();
|
||
let mut bpocom = BpoCom::new();
|
||
let fixalp = FixAlp::default();
|
||
|
||
// 设置非零值以验证初始化
|
||
for i in 0..5 {
|
||
bpocom.att[i] = 1.0;
|
||
bpocom.ann[i] = 2.0;
|
||
}
|
||
|
||
let params = BpopParams { id: 1 };
|
||
|
||
let result = bpop(
|
||
¶ms,
|
||
&mut config,
|
||
&mut atomic,
|
||
&mut model,
|
||
&iterat,
|
||
&mut arrays,
|
||
&mut bpocom,
|
||
&fixalp,
|
||
);
|
||
|
||
assert!(result.computed);
|
||
|
||
// 验证 ATT 和 ANN 被初始化为零
|
||
for i in 0..5 {
|
||
assert_relative_eq!(bpocom.att[i], 0.0);
|
||
assert_relative_eq!(bpocom.ann[i], 0.0);
|
||
}
|
||
}
|
||
}
|