503 lines
16 KiB
Rust
503 lines
16 KiB
Rust
//! 输出计算好的大气模型。
|
||
//!
|
||
//! 重构自 TLUSTY `output.f`
|
||
//! 将模型写入文件，可作为后续运行的输入模型。
|
||
|
||
use crate::tlusty::io::{FortranWriter, Result, format_exp_fortran};
|
||
use crate::tlusty::state::config::TlustyConfig;
|
||
use crate::tlusty::state::model::ModelState;
|
||
use std::io::Write;
|
||
|
||
/// OUTPUT 参数结构体
|
||
pub struct OutputParams<'a> {
|
||
/// 配置
|
||
pub config: &'a TlustyConfig,
|
||
/// 模型状态
|
||
pub model: &'a ModelState,
|
||
}
|
||
|
||
/// 输出模型到文件。
|
||
///
|
||
/// # 参数
|
||
/// * `writer7` - 主输出写入器 (fort.7)
|
||
/// * `params` - 输出参数
|
||
/// * `writer17` - 诊断输出写入器 (fort.17)，IPRIND>0 时使用
|
||
/// * `writer20` - BFAC 数据写入器 (fort.20)，NLTE 且 IPRIND>0 时使用
|
||
///
|
||
/// # 返回值
|
||
/// 成功返回 Ok(())
|
||
pub fn output<W: Write>(
|
||
writer7: &mut FortranWriter<W>,
|
||
params: &OutputParams,
|
||
writer17: Option<&mut FortranWriter<W>>,
|
||
mut writer20: Option<&mut FortranWriter<W>>,
|
||
) -> Result<()> {
|
||
let config = params.config;
|
||
let model = params.model;
|
||
|
||
let nd = config.basnum.nd as usize;
|
||
let nlevel = config.basnum.nlevel as usize;
|
||
let idisk = config.basnum.idisk;
|
||
let ifmol = config.basnum.ifmol;
|
||
let lte = config.inppar.lte;
|
||
let iprinp = config.prints.iprinp;
|
||
let iprind = config.prints.iprind;
|
||
|
||
// 计算 NUMLT 和 NUMPAR
|
||
let mut numlt: i32 = 3;
|
||
if idisk == 1 {
|
||
numlt = 4;
|
||
}
|
||
if ifmol == 1 {
|
||
numlt += 1;
|
||
}
|
||
|
||
let mut numpar: i32 = nlevel as i32 + numlt;
|
||
if lte && iprinp == 0 {
|
||
numpar = numlt;
|
||
}
|
||
if ifmol > 0 {
|
||
numpar = -numpar;
|
||
}
|
||
|
||
// 写入头部: ND, NUMPAR (unit 7)
|
||
writer7.write_raw(&format!("{:5}{:5}", nd, numpar))?;
|
||
writer7.write_newline()?;
|
||
|
||
// 写入 DM 数组 (6 个一组, FORMAT 1P6E13.6)
|
||
write_dm_array(writer7, &model.modpar.dm, nd)?;
|
||
|
||
// 写入每个深度点的数据 (unit 7)
|
||
write_depth_data(writer7, model, nd, nlevel, idisk, ifmol, lte, iprinp)?;
|
||
|
||
// IPRIND > 0: 写诊断输出到 unit 17 和 unit 20
|
||
if iprind > 0 {
|
||
if let Some(w17) = writer17 {
|
||
w17.write_raw(&format!("{:5}{:5}", nd, numpar))?;
|
||
w17.write_newline()?;
|
||
write_dm_array(w17, &model.modpar.dm, nd)?;
|
||
|
||
if idisk == 0 {
|
||
if lte {
|
||
for id in 0..nd {
|
||
write_depth_line(w17, &[
|
||
model.modpar.temp[id],
|
||
model.modpar.elec[id],
|
||
model.modpar.dens[id],
|
||
])?;
|
||
}
|
||
} else {
|
||
// NLTE: 写 unit 20 头部和 DM
|
||
if let Some(w20) = writer20.as_mut() {
|
||
w20.write_raw(&format!("{:5}{:5}", nd, numpar))?;
|
||
w20.write_newline()?;
|
||
write_dm_array(w20, &model.modpar.dm, nd)?;
|
||
}
|
||
for id in 0..nd {
|
||
write_depth_line_with_popul(
|
||
w17, model.modpar.temp[id], model.modpar.elec[id],
|
||
model.modpar.dens[id], None,
|
||
&model.levpop.popul, id, nlevel,
|
||
)?;
|
||
if let Some(w20) = writer20.as_mut() {
|
||
write_depth_line_with_popul(
|
||
w20, model.modpar.temp[id], model.modpar.elec[id],
|
||
model.modpar.dens[id], None,
|
||
&model.levpop.bfac, id, nlevel,
|
||
)?;
|
||
}
|
||
}
|
||
}
|
||
} else {
|
||
// 圆盘模型
|
||
if lte {
|
||
for id in 0..nd {
|
||
write_depth_line(w17, &[
|
||
model.modpar.temp[id],
|
||
model.modpar.elec[id],
|
||
model.modpar.dens[id],
|
||
model.modpar.zd[id],
|
||
])?;
|
||
}
|
||
} else {
|
||
if let Some(w20) = writer20.as_mut() {
|
||
w20.write_raw(&format!("{:5}{:5}", nd, numpar))?;
|
||
w20.write_newline()?;
|
||
write_dm_array(w20, &model.modpar.dm, nd)?;
|
||
}
|
||
for id in 0..nd {
|
||
write_depth_line_with_popul_disk(
|
||
w17, model.modpar.temp[id], model.modpar.elec[id],
|
||
model.modpar.dens[id], None, model.modpar.zd[id],
|
||
&model.levpop.popul, id, nlevel,
|
||
)?;
|
||
if let Some(w20) = writer20.as_mut() {
|
||
write_depth_line_with_popul_disk(
|
||
w20, model.modpar.temp[id], model.modpar.elec[id],
|
||
model.modpar.dens[id], None, model.modpar.zd[id],
|
||
&model.levpop.bfac, id, nlevel,
|
||
)?;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// 写入 DM 数组 (6 个一组, FORMAT 1P6E13.6)
|
||
fn write_dm_array<W: Write>(writer: &mut FortranWriter<W>, dm: &[f64], nd: usize) -> Result<()> {
|
||
let mut dm_line = String::new();
|
||
for (i, &dm_val) in dm.iter().take(nd).enumerate() {
|
||
dm_line.push_str(&format_exp_fortran(dm_val, 13, 6, false));
|
||
if (i + 1) % 6 == 0 || i == nd - 1 {
|
||
writer.write_raw(&dm_line)?;
|
||
writer.write_newline()?;
|
||
dm_line.clear();
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// 写入每个深度点的数据（unit 7 主输出）
|
||
fn write_depth_data<W: Write>(
|
||
writer: &mut FortranWriter<W>,
|
||
model: &ModelState,
|
||
nd: usize,
|
||
nlevel: usize,
|
||
idisk: i32,
|
||
ifmol: i32,
|
||
lte: bool,
|
||
iprinp: i32,
|
||
) -> Result<()> {
|
||
for id in 0..nd {
|
||
let temp = model.modpar.temp[id];
|
||
let elec = model.modpar.elec[id];
|
||
let dens = model.modpar.dens[id];
|
||
|
||
if idisk == 0 {
|
||
// 平面平行模型
|
||
if lte && iprinp == 0 {
|
||
if ifmol == 0 {
|
||
write_depth_line(writer, &[temp, elec, dens])?;
|
||
} else {
|
||
let totn = model.modpar.totn[id];
|
||
write_depth_line(writer, &[temp, elec, dens, totn])?;
|
||
}
|
||
} else {
|
||
if ifmol == 0 {
|
||
write_depth_line_with_popul(writer, temp, elec, dens, None, &model.levpop.popul, id, nlevel)?;
|
||
} else {
|
||
let totn = model.modpar.totn[id];
|
||
write_depth_line_with_popul(writer, temp, elec, dens, Some(totn), &model.levpop.popul, id, nlevel)?;
|
||
}
|
||
}
|
||
} else {
|
||
// 圆盘模型
|
||
let zd = model.modpar.zd[id];
|
||
if lte && iprinp == 0 {
|
||
if ifmol == 0 {
|
||
write_depth_line(writer, &[temp, elec, dens, zd])?;
|
||
} else {
|
||
let totn = model.modpar.totn[id];
|
||
write_depth_line(writer, &[temp, elec, dens, totn, zd])?;
|
||
}
|
||
} else {
|
||
if ifmol == 0 {
|
||
write_depth_line_with_popul_disk(writer, temp, elec, dens, None, zd, &model.levpop.popul, id, nlevel)?;
|
||
} else {
|
||
let totn = model.modpar.totn[id];
|
||
write_depth_line_with_popul_disk(writer, temp, elec, dens, Some(totn), zd, &model.levpop.popul, id, nlevel)?;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// 写入深度点数据（无占据数）
|
||
/// 使用 Fortran 格式 1P5E15.6
|
||
fn write_depth_line<W: Write>(writer: &mut FortranWriter<W>, values: &[f64]) -> Result<()> {
|
||
let mut line = String::new();
|
||
for &val in values {
|
||
line.push_str(&format_exp_fortran(val, 15, 6, false));
|
||
}
|
||
writer.write_raw(&line)?;
|
||
writer.write_newline()?;
|
||
Ok(())
|
||
}
|
||
|
||
/// 写入深度点数据（带占据数，平面平行）
|
||
/// 使用 Fortran 格式 1P5E15.6
|
||
fn write_depth_line_with_popul<W: Write>(
|
||
writer: &mut FortranWriter<W>,
|
||
temp: f64,
|
||
elec: f64,
|
||
dens: f64,
|
||
totn: Option<f64>,
|
||
popul: &[Vec<f64>],
|
||
id: usize,
|
||
nlevel: usize,
|
||
) -> Result<()> {
|
||
// Fortran: WRITE(7,503) TEMP(ID),ELEC(ID),DENS(ID),(POPUL(J,ID),J=1,NLEVEL)
|
||
// FORMAT(1P5E15.6) - 5 values per line, flowing continuously
|
||
// So first line has temp, elec, dens, pop[0], pop[1]
|
||
// then subsequent lines have 5 pop values each
|
||
|
||
let mut line = String::new();
|
||
let mut col = 0;
|
||
|
||
// Write temp, elec, dens
|
||
for &val in &[temp, elec, dens] {
|
||
line.push_str(&format_exp_fortran(val, 15, 6, false));
|
||
col += 1;
|
||
}
|
||
if let Some(t) = totn {
|
||
line.push_str(&format_exp_fortran(t, 15, 6, false));
|
||
col += 1;
|
||
}
|
||
|
||
// Write populations, continuing from where temp/elec/dens left off
|
||
for j in 0..nlevel {
|
||
line.push_str(&format_exp_fortran(popul[j][id], 15, 6, false));
|
||
col += 1;
|
||
if col % 5 == 0 || j == nlevel - 1 {
|
||
writer.write_raw(&line)?;
|
||
writer.write_newline()?;
|
||
line.clear();
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// 写入深度点数据（带占据数，圆盘）
|
||
/// 使用 Fortran 格式 1P5E15.6
|
||
fn write_depth_line_with_popul_disk<W: Write>(
|
||
writer: &mut FortranWriter<W>,
|
||
temp: f64,
|
||
elec: f64,
|
||
dens: f64,
|
||
totn: Option<f64>,
|
||
zd: f64,
|
||
popul: &[Vec<f64>],
|
||
id: usize,
|
||
nlevel: usize,
|
||
) -> Result<()> {
|
||
// Same flow as planar: values flow continuously, 5 per line
|
||
let mut line = String::new();
|
||
let mut col = 0;
|
||
|
||
for &val in &[temp, elec, dens] {
|
||
line.push_str(&format_exp_fortran(val, 15, 6, false));
|
||
col += 1;
|
||
}
|
||
if let Some(t) = totn {
|
||
line.push_str(&format_exp_fortran(t, 15, 6, false));
|
||
col += 1;
|
||
}
|
||
line.push_str(&format_exp_fortran(zd, 15, 6, false));
|
||
col += 1;
|
||
|
||
for j in 0..nlevel {
|
||
line.push_str(&format_exp_fortran(popul[j][id], 15, 6, false));
|
||
col += 1;
|
||
if col % 5 == 0 || j == nlevel - 1 {
|
||
writer.write_raw(&line)?;
|
||
writer.write_newline()?;
|
||
line.clear();
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn create_test_state() -> (TlustyConfig, ModelState) {
|
||
let mut config = TlustyConfig::default();
|
||
config.basnum.nd = 5;
|
||
config.basnum.nlevel = 3;
|
||
config.basnum.idisk = 0;
|
||
config.basnum.ifmol = 0;
|
||
config.inppar.lte = true;
|
||
config.prints.iprinp = 0;
|
||
|
||
let mut model = ModelState::new();
|
||
for i in 0..5 {
|
||
model.modpar.dm[i] = (i + 1) as f64 * 1e2;
|
||
model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
|
||
model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
|
||
model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
|
||
}
|
||
|
||
(config, model)
|
||
}
|
||
|
||
#[test]
|
||
fn test_output_lte_plane_parallel() {
|
||
let (config, model) = create_test_state();
|
||
|
||
let mut writer = FortranWriter::to_memory();
|
||
|
||
let params = OutputParams {
|
||
config: &config,
|
||
model: &model,
|
||
};
|
||
|
||
let result = output(&mut writer, ¶ms, None, None);
|
||
assert!(result.is_ok());
|
||
|
||
let output_str = writer.into_string().unwrap();
|
||
|
||
// 验证头部
|
||
let lines: Vec<&str> = output_str.lines().collect();
|
||
assert_eq!(lines[0], " 5 3"); // ND=5, NUMPAR=3 (LTE, plane-parallel)
|
||
|
||
// 验证 DM 数组 (6 个一组)
|
||
assert!(lines[1].contains("1.000000E+02"));
|
||
|
||
// 验证深度数据
|
||
assert!(lines[2].contains("1.000000E+04")); // temp
|
||
assert!(lines[2].contains("1.000000E+12")); // elec
|
||
assert!(lines[2].contains("1.000000E-07")); // dens
|
||
}
|
||
|
||
#[test]
|
||
fn test_output_nlte_with_populations() {
|
||
let mut config = TlustyConfig::default();
|
||
config.basnum.nd = 3;
|
||
config.basnum.nlevel = 7; // 测试多行占据数
|
||
config.basnum.idisk = 0;
|
||
config.basnum.ifmol = 0;
|
||
config.inppar.lte = false; // NLTE
|
||
config.prints.iprinp = 1; // 需要输出占据数
|
||
|
||
let mut model = ModelState::new();
|
||
for i in 0..3 {
|
||
model.modpar.dm[i] = (i + 1) as f64 * 1e2;
|
||
model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
|
||
model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
|
||
model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
|
||
}
|
||
// 设置占据数
|
||
for j in 0..7 {
|
||
for i in 0..3 {
|
||
model.levpop.popul[j][i] = (j + 1) as f64 * 1e10 + i as f64;
|
||
}
|
||
}
|
||
|
||
let mut writer = FortranWriter::to_memory();
|
||
|
||
let params = OutputParams {
|
||
config: &config,
|
||
model: &model,
|
||
};
|
||
|
||
let result = output(&mut writer, ¶ms, None, None);
|
||
assert!(result.is_ok());
|
||
|
||
let output_str = writer.into_string().unwrap();
|
||
let lines: Vec<&str> = output_str.lines().collect();
|
||
|
||
// 验证头部: NUMPAR = NLEVEL + NUMLT = 7 + 3 = 10
|
||
assert_eq!(lines[0], " 3 10");
|
||
|
||
// 验证 DM 数组
|
||
assert!(lines[1].contains("1.000000E+02"));
|
||
|
||
// 验证第一深度点的数据
|
||
// Line 2: temp, elec, dens (3 个值)
|
||
assert!(lines[2].contains("1.000000E+04")); // temp
|
||
assert!(lines[2].contains("1.000000E+12")); // elec
|
||
assert!(lines[2].contains("1.000000E-07")); // dens
|
||
|
||
// Line 3-4: 占据数 (7 个，每行最多 5 个，但第一行只有 2 个空位)
|
||
// 实际上，由于第一行已经写了 3 个值，占据数从新行开始
|
||
assert!(lines[3].contains("1.000000E+10")); // popul[0][0]
|
||
assert!(lines[3].contains("2.000000E+10")); // popul[1][0]
|
||
}
|
||
|
||
#[test]
|
||
fn test_output_disk_model() {
|
||
let mut config = TlustyConfig::default();
|
||
config.basnum.nd = 3;
|
||
config.basnum.nlevel = 2;
|
||
config.basnum.idisk = 1; // 圆盘模型
|
||
config.basnum.ifmol = 0;
|
||
config.inppar.lte = true;
|
||
config.prints.iprinp = 0;
|
||
|
||
let mut model = ModelState::new();
|
||
for i in 0..3 {
|
||
model.modpar.dm[i] = (i + 1) as f64 * 1e2;
|
||
model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
|
||
model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
|
||
model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
|
||
model.modpar.zd[i] = (i + 1) as f64 * 1e5;
|
||
}
|
||
|
||
let mut writer = FortranWriter::to_memory();
|
||
|
||
let params = OutputParams {
|
||
config: &config,
|
||
model: &model,
|
||
};
|
||
|
||
let result = output(&mut writer, ¶ms, None, None);
|
||
assert!(result.is_ok());
|
||
|
||
let output_str = writer.into_string().unwrap();
|
||
let lines: Vec<&str> = output_str.lines().collect();
|
||
|
||
// 验证头部: NUMPAR = NUMLT = 4 (LTE, disk)
|
||
assert_eq!(lines[0], " 3 4");
|
||
|
||
// 验证深度数据包含 ZD
|
||
assert!(lines[2].contains("1.000000E+05")); // zd
|
||
}
|
||
|
||
#[test]
|
||
fn test_output_with_molecules() {
|
||
let mut config = TlustyConfig::default();
|
||
config.basnum.nd = 3;
|
||
config.basnum.nlevel = 2;
|
||
config.basnum.idisk = 0;
|
||
config.basnum.ifmol = 1; // 有分子
|
||
config.inppar.lte = true;
|
||
config.prints.iprinp = 0;
|
||
|
||
let mut model = ModelState::new();
|
||
for i in 0..3 {
|
||
model.modpar.dm[i] = (i + 1) as f64 * 1e2;
|
||
model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
|
||
model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
|
||
model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
|
||
model.modpar.totn[i] = (i + 1) as f64 * 1e14;
|
||
}
|
||
|
||
let mut writer = FortranWriter::to_memory();
|
||
|
||
let params = OutputParams {
|
||
config: &config,
|
||
model: &model,
|
||
};
|
||
|
||
let result = output(&mut writer, ¶ms, None, None);
|
||
assert!(result.is_ok());
|
||
|
||
let output_str = writer.into_string().unwrap();
|
||
let lines: Vec<&str> = output_str.lines().collect();
|
||
|
||
// 验证头部: NUMPAR = -4 (负数表示有分子)
|
||
assert_eq!(lines[0], " 3 -4");
|
||
|
||
// 验证深度数据包含 TOTN
|
||
assert!(lines[2].contains("1.000000E+14")); // totn
|
||
}
|
||
}
|