[论文解读] TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

J. Citrin, Ian Goodfellow+12 more|arXiv (Cornell University)|Jun 10, 2024

Magnetic confinement fusion research被引用 6

一句话总结

TORAX 是一个开源的、可微分的单端 tokamak 核心输运仿真器，用 Python 与 JAX 构建，能够快速运行和基于梯度的优化，并且与 RAPTOR 进行基准对比。

ABSTRACT

We present TORAX, a new, open-source, differentiable tokamak core transport simulator implemented in Python using the JAX framework. TORAX solves the coupled equations for ion heat transport, electron heat transport, particle transport, and current diffusion, incorporating modular physics-based and ML models. JAX's just-in-time compilation ensures fast runtimes, while its automatic differentiation capability enables gradient-based optimization workflows and simplifies the use of Jacobian-based PDE solvers. Coupling to ML-surrogates of physics models is greatly facilitated by JAX's intrinsic support for neural network development and inference. TORAX is verified against the established RAPTOR code, demonstrating agreement in simulated plasma profiles. TORAX provides a powerful and versatile tool for accelerating research in tokamak scenario modeling, pulse design, and control.

研究动机与目标

• 提供一个在 Python/JAX 中实现的可微分、快速的核心 tokamak 输运仿真器。
• 实现与 ML 代理和基于梯度的优化工作流的耦合。
• 对 TORAX 进行与成熟代码（如 RAPTOR）的基准比较，以验证准确性和可靠性。
• 提供模块化的物理模型，涵盖几何、输运、源项和 neoclassical 物理，以支持情景建模和控制。
• 为扩展以包含更多物理和动态平衡奠定路线图。

提出的方法

• 将离子/电子热、电子密度和电流扩散的耦合一维输运偏微分方程在归一化通量坐标中进行表述。
• 在均匀的一维网格上使用有限体积法离散，并对面值应用幂律 Peclet 加权。
• 使用 theta-方法进行时间步进求解，并提供线性、牛顿-拉夫森和优化器求解器，借助 JAX 进行自动微分。
• 整合模块化物理模型，包括几何（CHEASE 或 Circular）、输运（Constant、CGM、QLKNN10D）和 neoclassical 物理（Sauter 模型）。
• 实现 ML-代理耦合（QLKNN10D），配有专门的 JAX 推理代码和用于求解器稳定性的平滑选项。

Figure 1: Poloidal cross section of the core region of a toroidal tokamak plasma, with closed nested magnetic flux surfaces. R is the major radius from the center of the torus, Z is the vertical direction. ITER-like magnitudes are shown. The flux surfaces are elliptical in the poloidal plane for ill

实验结果

研究问题

• RQ1 TORAX 是否能够以适用于基于梯度的工作流的可微分、快速求解器再现核心 tokamak 输运物理？
• RQ2 TORAX 与成熟的输运代码（如 RAPTOR）在等离子体轮廓方面的一致性如何？
• RQ3 模块化物理模型（几何、输运、 neoclassical）对仿真速度、稳定性和可扩展性有何影响？
• RQ4 在保持可微分性和数值鲁棒性的前提下，TORAX 在多大程度上能够整合用于输运模型的 ML 代理？
• RQ5 计划有哪些路线扩展以扩大 TORAX 的物理覆盖范围和情景？

主要发现

• TORAX 提供一个在 Python 中使用 JAX 的可微分、快速的核心输运仿真器，能够进行基于梯度的优化和 ML 耦合。
• TORAX 已通过与 RAPTOR 的对比验证，在模拟等离子体轮廓方面显示出一致性。
• TORAX 在模块化框架内支持多种物理模块（几何、输运模型 CGM/QLKNN10D、neoclassical Sauter）。
• 求解器提供线性、牛顿-拉夫森和基于优化器的选项，具有自适应时间步长和用于鲁棒性的线搜索/回溯。
• JAX 实现自动微分和后端灵活性（CPU/GPU），便于 ML 代理和多工作流集成。
• 代码库是开源的，计划扩展以包含更多物理和动态平衡。

Figure 2: 1D spatial discretization with the finite volume method, with $N=5$ and $\hat{\rho}=0.2$ . The cell volume centered at $\hat{\rho}=0.5$ is highlighted. Keeping with 3D FVM formalism, a cell volume is defined with $V=d\hat{\rho}A$ , with arbitrary face area $A$ . The corresponding cell face