yancc: A GPU-accelerated, differentiable solver for neoclassical transport in tokamaks and stellarators

技術要約：yancc: トカマクおよびステラレータにおける新古典輸送のためのGPU加速・微分可能なソルバー

問題提起
新古典輸送理論は、磁気閉じ込め核融合装置（トカマクおよびステラレータ）における粒子、熱、および運動量の衝突輸送を記述するものである。ドリフト・キネティック方程式（DKE）を解くことは、径方向電場、ブートストラップ電流、および輸送係数を決定するために不可欠である。しかし、DKEの数値解法は、強い移流支配性、再循環流、内部境界層（特に捕捉粒子と通過粒子の間の境界）、深刻な異方性、および高次元性（空間2次元、速度2次元の計4次元）により、極めて困難であることが知られている。

既存のソルバーは、計算速度と物理的忠実度の間でトレードオフに直面している。単エネルギー近似コード（例：DKES, MONKES）は、性能のために種間結合やエネルギー散乱を犠牲にしている。フル4Dソルバー（例：SFINCS, NEO）は高い精度を提供するが、不規則なメモリアクセスを伴う疎行列の因数分解により、メモリ使用量が多く、GPUハードウェアへのスケーラビリティに乏しい。さらに、既存のコードには自動微分機能が欠けているため、ステラレータ設計や不確実性定量化においてますます必要とされている勾配ベースの最適化や随伴感度解析が困難である。

手法
著者らは、GPU加速と自動微分を活用するためにJAXフレームワーク上に構築された新しいソルバーであるyanccを提示する。このコードは、フル4D DKE（速度依存の衝突、エネルギー散乱、および完全な種間結合を保持）と、簡略化された単エネルギー形式の両方を解くことができる。

• 離散化:

  • 速度 ($x$): 重み付きマックスウェル多項式（$e^{-x^2}$に関して直交）を用いたコロケーション・グリッドを使用し、わずか5〜10点のみを必要とする。
  • ピッチ角 ($\alpha$)、ポロイダル角 ($\theta$)、およびトロイダル角 ($\zeta$): 差分離散化を用いる。
  • ステンシル修正: マルチグリッドの効率を高めるため、標準的なアップウィンド・ステンシルの代わりに、修正されたアップウィンド・ステンシル（例：2次精度では $\{0, 1, 4\}$、4次精度では $\{-2, 0, 1, 3, 4\}$ を使用）を採用している。これにより、主要な誤差定数を大幅に増大させることなく、結果として得られる線形システムの行列の対角優位性を高めている。
  • 衝突演算子: 場・粒子散乱（ローゼンブラット・ポテンシャル）は、ルジャンドル展開と不完全ガンマ関数を用いて、速度空間における関連するポアソン方程式をグリーン関数法を通じて解析的に処理しており、ポテンシャルを未知数として追加で解く必要を回避している。

• 線形ソルバー:

  • 離散化されたDKEは、大きな疎な線形システム $Af = s$ を形成する。
  • 疎な直接解法（GPU上でのスケーリングに失敗する）の代わりに、yanccは幾何学的マルチグリッド・スキームによって前処理されたKrylov部分空間法（内側GMRESを用いたGCROT）を使用する。
  • マルチグリッド戦略:
    • グリッド階層: セミ・コーサニング（半粗分化）を使用する。速度座標 ($x$) はすべてのレベルで最も細かい解像度を維持し、$\alpha, \theta, \zeta$ は因子2で粗分化される。
    • スムーシング: 演算子のブロック対角近似を使用し、一つの座標に沿った完全な結合を維持するように置換をサイクルさせる。修正されたアップウィンド・ステンシルは、このブロック対角スムーザーの効果を高める。
    • 粗いグリッド補正: メモリ使用量を低く抑え、行列フリーの粗いグリッド操作のコストを避けるために、直接再離散化が好まれる。
    • サイクル: フルDKEにはVサイクル（インデックス1）を使用し、単エネルギーの場合にはWサイクルまたはそれ以上の次数（インデックス3）が効率的であることが判明している。

主な貢献

1. GPU加速されたフル4Dソルバー: yanccは、フル4D DKE（完全な種間結合およびエネルギー散乱を含む）の解決と、効率的なGPU実行を組み合わせた初のソルバーである。これにより、SFINCSのような疎行列ベースのコードが持つメモリおよびスケーリングの制限を克服している。
2. 微分可能性: JAXで構築されているため、本ソルバーは完全に微分可能であり、磁気形状の勾配ベースの最適化や、輸送係数の手動導出を必要としない随伴感度解析を可能にする。
3. アルゴリズムの革新: 対角優位性を高めるための拡張アップウィンド・ステンシルの導入、およびDKEの特定の異方性に適したセミ・コーサニングとブロック対角スムーザーの使用。
4. 解析的ローゼンブラット・ポテンシャル: 事前計算されたグリーン関数とテンソル縮退を用いた、効率的な場・粒子衝突演算子の実装。これにより、グリッド上で補助的なポアソン方程式を解く計算オーバーヘッドを回避している。

結果とベンチマーク

• 精度: MONKES（単エネルギー）およびSFINCS（フル4D）との比較ベンチマークでは、幅広い衝突性（$\nu^* \in [10^{-4}, 10^2]$）、磁気形状（W7-X、NCSX様）、およびプラズマ構成（単一および多種）において1%以内の合致を示している。
• SFINCSに対するパフォーマンス:
  • 速度: yanccは、1スキャンあたりの実行時間においてSFINCSに対して1桁の高速化を達成している。SFINCSの実行時間は高衝突性において著しく増加する（前処理器の劣化による）が、yanccの実行時間はほぼ平坦である。
  • メモリ: yanccは、テストされたシナリオにおいて、SFINCS（>50 GB）と比較して1桁少ないメモリ（例：~6 GB）を使用しており、GPUベースのワークフローへの適用が可能である。
• MONKESに対するパフォーマンス: 単エネルギーのベンチマークにおいて、yanochはほとんどの衝突性範囲でJAXで再実装されたMONKESアルゴリズムよりも2〜4倍高速であるが、非常に単純な形状で低解像度を必要とする場合にはMONKESの方が高速になる可能性がある。

意義と主張
論文は、速度、低メモリ使用量、および微分可能性の組み合わせにより、yanccが現代のステラレータ最適化ワークフロー、不確実性定量化（モンテカルロ・サンプリング経由）、および統合プロファイル予測への統合に特異的に適していると主張している。事前計算されたデータベースやフィット関数に依存する現在のアプローチとは異なり、本コードは輸送係数をオンザフライで計算できる。入力パラメータに対する輸送係数の勾配を計算できる能力により、完全な構成空間における効率的な最適化が可能となる。著者らは、yanccを高度な物理学と設計・解析の計算要求との間のギャップを埋めるツールとして位置づけており、今後の課題として、自己整合的な径方向電場の計算および適応的メッシュ細分化を含めるとしている。