技術要約: JAX-FVM

問題提起

数値流体力学（CFD）は伝統的に順方向のシミュレーションに使用されてきたが、形状最適化、パラメータ同定、乱流モデリングといった多くの重要なアプリケーションでは、逆問題の解決が必要となる。これらのタスクには、入力（境界条件、物理パラメータ、初期状態など）に対するソルバー出力の勾配を計算することが不可欠である。これまで、これらには随伴変数法（adjoint methods）が用いられてきたが、これらはコードへの介入（code-intrusive）が必要であり、現代の機械学習（ML）エコシステムからは切り離されていることが多い。

近年の微分可能プログラミング（JAX、PyTorchなど）の発展により、エンドツーエンドの微分可能性が可能となり、手動で導出された随伴変数を用いずに、自動微分（AD）を通じて勾配を計算できるようになった。しかし、既存の微分可能なCFDフレームワーク（JAX-CFD、JAX-Fluids、ΦFlowなど）は、主に構造化格子またはスペクトル離散化に限定されている。ここに大きなギャップが存在する。すなわち、産業用CFDに不可欠な複雑な幾何学的形状を扱うための非構造格子は、アクセラレータのベクトル化実行モデルと相反する不規則な接続性を持つため、微分可能なシミュレーションのエコシステムからはほとんど欠落している。さらに、非構造格子上で圧縮性流体をロバストにシミュレートするには、非線形不安定性や衝撃波付近での偽の振動を防ぐためのエントロピー安定スキームが必要であるが、これは完全に微分可能なフレームワーク内で実装することがしばしば困難な特徴である。

手法

本論文では、2次元圧縮性オイラー方程式およびナビエ・ストークス方程式のための、非構造三角形メッシュを用いた、完全に微分可能な有限体積法（FVM）ソルバーであるJAX-FVMを提案する。このソルバーはすべてJAXで記述されており、メッシュの接続性やフラックス評価から、勾配制限（slope limiting）や時間積分に至るまで、あらゆる操作がジャストインタイム（JIT）コンパイル、ベクトル化、および微分可能であることを保証している。

数値離散化

• 支配方程式: ソルバーは、理想気体の状態方程式によって閉じられた、2次元圧縮性ナビエ・ストークス方程式の保存形式を扱う。
• メッシュ表現: ドメインはMeshPy/Triangleジェネレータを用いて三角形分割される。メッシュデータ構造は、頂点、セルの重心、面積、および面接続を格納し、周期境界条件および様々な境界条件マーカーをサポートする。
• 半離散形式: セル中心の更新は、発散定理を介して導出される。残差は、セル面上の数値フラックスの総和として計算される。
• 勾配再構成: 2次精度は、原始変数（primitive variables）のMUSCL再構成を通じて達成される。セルの勾配は、ベクトル化された線形ソルバー（jax.vmap）を介して解かれる重み付き最小二乗（LSQ）フィットを用いて計算される。勾配制限には、Venkatakrishnanリミッターまたは古典的なminmodリミッターのいずれかが適用される。
• 数値フラックス: 本スキームの核となるのは、エントロピー安定な定式化である。
  • ベースラインとなる対流フラックスは、エントロピー保存型の2点フラックス（Tadmor/Ismail-Roe）である。
  • 消散（dissipation）は、エントロピー変数を用いたRusanov（局所Lax-Friedrichs）またはRoe消散項によって追加される。
  • 決定的な点として、エントロピー変数のジャンプを保存変数に写像するために必要なヤコビアン作用は、解析的な手動導出ではなく、前方モードAD（jax.jvp）を用いて正確に計算される。
  • 粘性フラックスは、面平均の原始状態とLSQ勾配を用いて評価される。
• 時間積分: ソルバーは、明示的積分器（RK2, RK3, RK4）および行列フリーの陰解法（Backward Euler, SDIRK2）をサポートする。陰的なステージでは、Newton反復を用いて非線形システムを解く。ここで、ヤコビアン・ベクトル積は前方モードADを用いて行列フリーで計算され、GMRESソルバーに渡される。

ソフトウェアアーキテクチャ

• 微分可能性: 残差および時間ステップのロジック全体がjax.jitでラップされている。エントロピー消散および陰解法のヤコビアンには前方モードADが内部的に使用され、逆モードADは、任意のスカラ目的関数から入力への勾配の伝播を可能にする。
• ハードウェア非依存性: コードはCPUまたはGPU（CUDA対応のJAX経由）上で透過的に動作する。
• 統合: ライブラリはflaxおよびoptaxに依存しており、ソルバーの上に直接、最適化や逆設計のループを構築することを容易にしている。

主な貢献

1. 公開されている初の微分可能な圧縮性流体用非構造FVM: JAX-FVMは、構造化格子に限定されていた既存の微分可能なCFDツールの空白を埋めるものである。著者の知る限り、これは、ネイティブに非構造三角形メッシュ上で動作する、圧縮性オイラー/ナビエ・ストークス方程式のための最初の微分可能なソルバーの一つである。
2. ADによるエントロピー安定性: この実装は、エントロピー保存型フラックス理論とエントロピー変数消散を統合している。消散項のための必要なヤコビアン作用をADを用いて計算することで、本ソルバーは、複雑なヤコビアンの手動導出なしに、証明可能なエントロピー安定性を達成している。
3. エンドツーエンドの微分可能性: ソルバーは、初期条件、境界データ、物理パラメータ（粘性、熱伝導率、$\gamma$）、または組み込まれたニューラルモデルに対するフロー関数の勾配を直接計算することを可能にする。
4. 研究プラットフォーム: 本コードは、非構造幾何、エントロピー一貫的な数値計算、およびJAXベースのADを組み合わせた、コンパクトで研究指向の実装を提供し、勾配ベースの逆問題およびハイブリッド物理/MLモデリングの基礎として機能する。

結果と検証

JAX-FVMは、一連の標準的な検証ケースを通じて検証される：

• Mach 3 前方ステップ: 分離したボウショック、マッハステム、反射衝撃波、およびスリップラインを含む、典型的な非定常衝撃波問題。ソルバーは、エントロピー保存型フラックス、Venkatakrishnanリミッターを用いたMUSCL再構成、および明示的RK2時間ステップを用い、完全な非構造三角形分割上でこれらの特徴を正常に捉える。
• その他のベンチマーク: 配布物には、19種類の2次元リエマン構成（LaxおよびLiu）、ダブルマッハ反射、ケルビン・ヘルムホルツ不安定性、テイラー・グリーン渦、および様々な低マッハ数渦が含まれる。
• 診断: ソルバーは、構造保存特性を検証するために、全エントロピー、運動エネルギー、エンストロフィー、パリストロフィー、渦度、およびマッハ数を監視する。

意義と主張

本論文は、JAX-FVMを、複雑な幾何形状上の高忠実度CFDと現代的な微分可能プログラミングの架け橋として位置づけている。その意義は、非構造格子上の圧縮性流体に対する勾配ベースの逆設計およびPDE制約最適化を可能にする点にあり、これはこれまでオープンソースの微分可能なフレームワークでは利用できなかった能力である。

著者らは、非構造格子能力とエントロピー安定な数値計算、およびエンドツーエンドの微分可能性を組み合わせることで、JAX-FVMが以下のような自然なプラットフォームを創出すると主張している：

• 希薄なデータからのパラメータ回収（例：血管抵抗）。
• 表面圧力に基づく形状最適化。
• ソルバーと一貫した方法での、乱流クロージャまたはデータ駆動型離散化の学習。

本研究は、現在の制限事項を認めた上で、その範囲を控えめに設定している。具体的には、2次元三角形メッシュに限定されており（3次元への拡張を計画中）、MeshPyジェネレータに依存している（VTKリーダーを計画中）、および現在は乱流モデル（RANSまたはLESクロージャ）を欠いている。本ソルバーは、主にハイブリッド物理/MLモデルおよび逆設計ワークフローの開発を促進するための研究ツールとして提示されている。