[論文レビュー] 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はPythonとJAXで構築されたオープンソースの微分可能な1Dトカマクコア輸送シミュレータで、高速な実行と勾配ベースの最適化を可能にし、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で実装された微分可能で高速なコアトカマク輸送シミュレータを提供する。
• MLサロゲートと勾配ベースの最適化ワークフローとの結合を可能にする。
• 精度と信頼性を検証するため、RAPTOR などの既存コードと TORAX をベンチマークする。
• ジオメトリ、輸送、ソース、ネオクラシカル物理を含むモジュール型の物理モデルを提供し、シナリオモデリングと制御をサポートする。
• より多くの物理と動的平衡を含む拡張のロードマップを示す。

提案手法

• 正規化フラックス座標で、イオン/電子熱、電子密度、電流拡散を対象とする1D輸送PDEの結合系を定式化する。
• 均一な1Dグリッド上で有限体積法を用いて離散化し、面値にはべき法則のペックレ重みを適用する。
• θ-法に基づく時間ステップで時間発展を解き、JAXによる自動微分を利用した線形、ニュートン-ラフソン、最適化子ソルバーを提供する。
• ジオメトリ（CHEASE または Circular）、輸送（Constant, CGM, QLKNN10D）、ネオクラシカル物理（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はコアトカマク輸送物理を再現できるか。
• RQ2 TORAXはプラズマプロファイルに対して、RAPTOR などの既存の輸送コードとどの程度一致するか。
• RQ3 モジュール化された物理モデル（ジオメトリ、輸送、ネオクラシカル）がシミュレーション速度、安定性、拡張性に与える影響はどれくらいか。
• RQ4 微分可能性と数値的頑健性を維持しつつ、TORAXは輸送モデルのMLサロゲートをどの程度統合できるか。
• RQ5 TORAXの物理カバー範囲とシナリオを拡張する計画ロードマップは何か。

主な発見

• TORAXはJAXを用いたPythonの微分可能で高速なコア輸送シミュレータを提供し、勾配ベースの最適化とML結合を可能にする。
• TORAXはRAPTORと比較検証され、シミュレートされたプラズマプロファイルの一致を示している。
• TORAXはモジュール型フレームワーク内で複数の物理モジュール（ジオメトリ、輸送モデルCGM/QLKNN10D、ネオクラシカル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