TORAX：基于 JAX 实现的高速可微型托卡马克输运模拟器
等离子体物理
2024-12-10 v3
摘要
我们提出TORAX，这是一个全新的、开源的、可微型托卡马克核心输运模拟器，采用 Python 实现并基于 JAX 框架开发。TORAX 求解离子热输运、电子热输运、粒子输运和电流扩散的耦合方程，集成模块化的基于物理的模型和机器学习模型。JAX 的即时编译确保了快速的运行时性能，其自动微分功能则使梯度优化工作流程成为可能，并简化了雅可比矩阵求解偏微分方程的数值方法。与机器学习模型的耦合 greatly facilitated by JAX's intrinsic support for neural network development and inference. TORAX 通过与 established RAPTOR code 进行验证，展示了在模拟等离子体剂量图谱方面的一致性。TORAX 提供了一个强大且灵活的工具，用于加速托卡马克情景建模、脉冲设计和控制研究。
引用
@article{arxiv.2406.06718,
title = {TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX},
author = {Jonathan Citrin and Ian Goodfellow and Akhil Raju and Jeremy Chen and Jonas Degrave and Craig Donner and Federico Felici and Philippe Hamel and Andrea Huber and Dmitry Nikulin and David Pfau and Brendan Tracey and Martin Riedmiller and Pushmeet Kohli},
journal= {arXiv preprint arXiv:2406.06718},
year = {2024}
}
备注
16 pages, 7 figures