什么是 TORAX？

TORAX 是一个可微分的托卡马克芯部输运模拟器，旨在实现快速且精确的正向建模、脉冲设计、轨迹优化以及控制器设计工作流。TORAX 使用 JAX 和 Python 编写，其设计动机如下：

• 开源且可扩展，便于灵活的工作流耦合。
• JAX 提供自动微分能力和代码编译功能，从而实现快速运行。可微性使得基于梯度的非线性偏微分方程求解器能够进行快速准确的建模，并对仿真结果相对于任意参数输入的敏感性进行分析，进而支持诸如轨迹优化和半经验模型的数据驱动参数识别等应用。自动微分特性还允许通过添加新的物理模型（包括机器学习代理模型）或新的参数输入来轻松扩展这些应用，而无需手动推导雅可比矩阵。
• Python-JAX 是一种天然适合与物理模型的机器学习代理模型进行耦合的框架。

如需更全面的文档，请参阅我们的 readthedocs 页面。

在 v1.0.0 版本中，TORAX 具备以下物理和数值功能集：

• 离子和电子热输运、电子粒子输运以及电流扩散的耦合偏微分方程组
  • 有限体积法离散化
  • 多种求解器选项：带有 Pereverzev-Corrigan 项的线性求解器及预测校正法，牛顿-拉夫森法的非线性求解器，以及使用 jaxopt 库进行优化的非线性求解器
  • 基于最后闭合磁通面处总电流或环电压的极向磁通边界条件
• 欧姆加热功率、离子-电子热交换、聚变功率、轫致辐射、杂质谱线辐射，以及一项 [ICRH 机器学习代理模型]（目前仅覆盖有限的工况）。
• 基于 Sauter 解析模型的新生代自举电流和电导率。
• 与 [QLKNN_7_11] 和 QLKNN10D [van de Plassche 等人, Phys. Plasmas 2020] 的 QuaLiKiz 神经网络代理模型相耦合，用于描述基于物理机制的湍流输运。
• 支持通用几何形状，可通过 CHEASE、FBT 或 EQDSK 平衡文件提供。
  • 为测试和演示目的，data/third_party/geo 目录下提供了一个单一的 CHEASE 平衡文件。该文件对应于 ITER 混合场景下的平衡态，基于 [Citrin 等人, Nucl. Fusion 2010] 的模拟结果，并从 PINT 获取。PINT 许可证文件位于 data/third_party/geo 中。
  • 通过提供一系列随时间变化的几何文件，可以支持时变几何。
• 使用局部自适应源设置内部边界条件的简单台基模型。
• 锯齿振荡触发及剖面重分布。

用户还可以通过自定义的解析模型或预设数据来引入额外的加热和电流驱动源。

模型的实现已通过将仿真输出与 RAPTOR [Felici 等人, Plasma Phys. Control. Fusion 2012] 托卡马克输运模拟器的结果直接对比进行了验证。

本项目并非 Google 官方支持的产品。

开发路线图

我们的 readthedocs 页面中概述了 开发路线图。

安装指南

系统要求

请安装 Python 3.11 或更高版本。

确保已安装 tkinter：

sudo apt-get install python3-tk

安装方法

准备虚拟环境

安装 Virtualenv（如果尚未安装）：

pip install --upgrade pip

pip install virtualenv

创建并激活虚拟环境：

python3 -m venv toraxvenv
source toraxvenv/bin/activate

从 PyPI 安装

使用 TORAX 最简单的方式是通过 PyPI 进行安装：

pip install torax

您可以检查一切是否正常运行：

run_torax --config=examples/basic_config.py --quit

如果您计划参与 TORAX 的开发，则需要克隆仓库，详情请参阅 贡献指南 和 贡献提示。

运行示例

以下命令将使用默认配置文件 examples/basic_config.py 运行 TORAX：

run_torax --config='examples/basic_config.py'

仿真进度会通过终端进度条显示当前时间和已完成百分比。

若要运行更复杂的、受 ITER 启发的仿真，可以执行：

run_torax --config='examples/iterhybrid_rampup.py'

以及

run_torax --config='examples/iterhybrid_predictor_corrector.py'

此外，还可以通过附加到上述运行命令的标志以及环境变量来进一步配置。例如，若需提高输出的详细程度，可以使用 --log_progress 标志。

run_torax  --config='examples/iterhybrid_rampup.py' --log_progress

设置环境变量

错误检查

若设置为 True，将在内部例程中启用错误检查。主要用于调试。默认值为 False，因为这与持久化编译缓存不兼容。

$ export TORAX_ERRORS_ENABLED=<True/False>

JAX 编译与缓存

若设置为 False，JAX 将不会编译 TORAX 的内部函数。主要用于调试。默认值为 False。

$ export JAX_DISABLE_JIT=<True/False>

以下设置启用了 JAX 的持久化缓存，会使 JAX 将编译后的程序存储到文件系统中，在某些情况下可减少重新编译的时间：

$ export JAX_COMPILATION_CACHE_DIR=<您选择的路径，例如 ~/jax_cache>
$ export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1
$ export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0.0

设置标志

log_progress - 输出每个时间步的仿真时间、dt 以及求解器迭代次数（对于非线性求解器，还包括 dt 回溯次数）。

run_torax \
   --config=examples/iterhybrid_predictor_corrector.py \
   --log_progress

output_dir - 覆盖默认输出目录。

run_torax \
   --config=examples/iterhybrid_predictor_corrector.py \
   --output_dir=/path/to/output/dir

运行包含完整断言检查的测试

要运行启用所有断言检查的 pytest 测试套件，可以使用以下命令：

TORAX_JAXTYPING=True TORAX_ERRORS_ENABLED=True pytest . -n auto

模拟后处理

模拟完成后，模拟状态及其派生量的时间历程将被写入一个带时间戳的文件，文件格式为 state_history_%Y%m%d_%H%M%S.nc。输出目录可由用户配置，默认路径为 /tmp/torax_results。

为了充分利用内存中的（非持久化）缓存，进程在模拟结束时并不会终止。您可以修改 runtime_params，切换 log_progress 和 plot_progress 标志，并重新运行模拟。以下修改会触发重新编译：

• 网格分辨率
• 演化变量（即所求解的方程）
• 更改内部使用的函数，例如输运模型、时间步长计算器、源项模型、求解器模型或pedestal模型。

清理工作

您可以通过取消激活虚拟环境来退出 Python 虚拟环境：

deactivate

（可选）安装 QLKNN-hyper

TORAX 中默认的输运模型是 QLKNN_7_11，它作为 TORAX 的依赖项被安装。另一种选择是使用 QLKNN-hyper-10D，也称为 QLKNN10D K.L. van de Plassche PoP 2020。QLKNN_7_11 基于 QuaLiKiz 2.8.1，其碰撞算子相比 QLKNN10D 的训练数据集有所改进。QLKNN_7_11 的训练数据还包括杂质密度梯度作为输入特征，并且对近 LCFS 区域的覆盖范围优于 QLKNN-hyper-10D。然而，在某些特定场景下，例如与其他模拟器进行基准测试，或与 QLKNN_7_11 进行对比时，使用 QLKNN10D 仍然可能具有价值。

下载 QLKNN-hyper 的依赖项：

git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git

要使用 QLKNN10D，请将您的 TORAX 配置文件中 transport 部分的 model_path 设置为克隆后的仓库路径。

模拟教程

请参阅我们 readthedocs 页面 中的教程部分。

引用 TORAX

关于 TORAX 的论文已在 arXiv 上发布：arXiv:2406.06718。引用该论文即可正确引用 TORAX：

@article{torax2024arxiv,
  title={{TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX}},
  author={Citrin, Jonathan and Goodfellow, Ian and Raju, Akhil and Chen, Jeremy
  and Degrave, Jonas and Donner, Craig and Felici, Federico and Hamel, Philippe
  and Huber, Andrea and Nikulin, Dmitry and Pfau, David and Tracey, Brendan, and
  Riedmiller, Martin and Kohli, Pushmeet},
  journal={arXiv preprint arXiv:2406.06718},
  year={2024}
}

TORAX 快速上手指南

TORAX 是一个基于 Python 和 JAX 开发的开源可微托卡马克核心输运模拟器，专为快速正向建模、脉冲设计、轨迹优化及控制器设计而打造。利用 JAX 的自动微分和编译加速特性，它支持基于梯度的非线性偏微分方程求解及灵敏度分析。

环境准备

系统要求

• Python 版本：3.11 或更高版本
• 操作系统：Linux / macOS (Windows 需通过 WSL 运行)

前置依赖

确保系统已安装 tkinter 库（用于部分可视化功能）：

sudo apt-get install python3-tk

提示：国内用户建议使用国内镜像源加速 pip 安装过程（如清华源、阿里源）。

安装步骤

1. 创建虚拟环境

推荐使用 venv 隔离开发环境：

python3 -m venv toraxvenv
source toraxvenv/bin/activate

升级 pip 并安装 virtualenv（如未安装）：

pip install --upgrade pip
pip install virtualenv

2. 安装 TORAX

从 PyPI 直接安装是最简便的方式：

pip install torax -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 验证安装

运行以下命令检查是否安装成功：

run_torax --config=examples/basic_config.py --quit

若无报错退出，则说明安装正常。

基本使用

运行基础示例

使用默认配置文件运行一个简单的模拟：

run_torax --config='examples/basic_config.py'

运行时终端会显示进度条，包含当前时间和完成百分比。

运行进阶示例（ITER 混合场景）

尝试更复杂的 ITER 混合场景模拟：

run_torax --config='examples/iterhybrid_rampup.py'

或使用预测 - 校正求解器：

run_torax --config='examples/iterhybrid_predictor_corrector.py'

常用参数配置

可以通过命令行标志自定义运行行为：

• 输出详细日志（显示时间步长、求解器迭代次数等）：

  run_torax --config='examples/iterhybrid_rampup.py' --log_progress
  

• 指定输出目录：

  run_torax --config='examples/iterhybrid_predictor_corrector.py' --output_dir=/path/to/output/dir
  

模拟结果

模拟完成后，状态历史数据将保存为 NetCDF 格式文件（文件名格式：state_history_YYYYMMDD_HHMMSS.nc），默认存储路径为 /tmp/torax_results。

注意：若需修改网格分辨率、演化变量或物理模型等核心参数，重新运行时会触发 JAX 重新编译。