TLUSTY NN：用神经网络秒级预测恒星大气模型的开源工具

TLUSTY NN 是一个基于深度学习的恒星大气模型预测工具，能够在毫秒级时间内生成传统数值模拟需要数小时才能计算的50层大气结构参数，为天体物理学研究提供了革命性的加速方案。

在天体物理学研究中，理解恒星的大气结构是分析恒星光谱、推算物理参数的基础。TLUSTY（Turkish Light Understanding Stellar Atmosphere Code）是业界广泛使用的恒星大气数值模拟软件，能够精确计算恒星大气的温度分布、电子密度、质量密度以及原子能级布居等关键参数。

然而，传统的 TLUSTY 数值模拟存在一个显著瓶颈：单次完整的大气结构计算往往需要数小时甚至更长时间。当研究人员需要扫描大量参数空间（例如研究不同温度、表面重力、化学组成的恒星样本）时，这种计算成本成为严重制约。

TLUSTY NN 项目正是为解决这一痛点而生——它利用深度神经网络的强大拟合能力，将数小时的数值模拟压缩到毫秒级的推理时间，同时保持足够的物理精度。

TLUSTY NN 的神经网络接受三个核心恒星参数作为输入：

• 有效温度（Teff）：范围 10,000 - 100,000 K，反映恒星表面的热辐射强度
• 表面重力（logg）：范围 1.5 - 10.0（log₁₀(cm/s²)），与恒星质量和半径密切相关
• 氦丰度（log(n_He / n_H)）：范围 -4.0 到 0.0（dex），表征恒星的化学组成

基于这三个参数，网络能够预测完整的 50 层大气结构，每层包含 58 个物理量：

• 光学深度（tau）
• 温度分布（T）
• 电子数密度（ne）
• 质量密度（rho）
• 55 个能级布居参数（level_1 至 level_55）

项目采用了物理信息神经网络的训练范式，在标准的数据驱动损失函数之外，额外引入了物理约束损失项。这意味着网络不仅学习训练数据的模式，还必须遵守基本的物理定律（如流体静力学平衡、能量守恒等）。这种设计显著提升了模型在参数空间边界处的泛化能力，并确保预测结果的天体物理合理性。

TLUSTY NN 提供了灵活的输出格式选择：

CSV 格式：便于数据分析和可视化，包含清晰的列标题和结构化的表格数据，与原始训练数据格式保持一致。

TLUSTY .7 格式：标准的 TLUSTY fort.7 输出格式，可直接作为其他天体物理软件的输入，实现与现有工具链的无缝集成。

TLUSTY .5 格式：生成 TLUSTY 的输入模型文件（fort.5 格式），支持非局部热动平衡（non-LTE）计算配置。

项目基于 Python 3.9+ 和 PyTorch 2.0+ 构建，安装过程简洁明了：

# 克隆仓库
git clone https://github.com/JqRambo/tlustynn.git
cd tlustynn

# 安装依赖
pip install .

安装完成后，即可通过 import tlustynn 在任何 Python 环境中调用。

以下代码演示如何预测一个特定恒星参数的大气模型：

from tlustynn import predict_atmosphere

# 预测并保存为 CSV
df, csv_path = predict_atmosphere(
    teff=10000,      # 有效温度 10000 K
    logg=3.7,        # 表面重力 logg=3.7
    mh=0.0,          # 太阳金属丰度
    output_dir="./predictions",
    output_format='csv'
)

print(f"结果已保存至: {csv_path}")
print(f"数据维度: {df.shape}")  # 输出: (50, 58)，即 50 层 × 58 参数

对于需要进一步进行非局部热动平衡计算的场景，可以生成标准的 TLUSTY 输入文件：

from tlustynn import create_ff_model

create_ff_model(
    output_dir='/path/to/workdir',
    teff=10000,
    logg=3.7,
    mh=0.0,
    lte_flag='F',      # 非 LTE 模式
    ltgray_flag='F',
    nstmode='nst',
    frequency=2000,
    natoms_num=8
)