google/jax-cfd

JAX‑CFD: JAXによる微分可能な計算流体力学

概要 – JAX を使用してコア CFD ソルバー（有限体積法および擬似スペクトル法）を実装したオープンソースの研究用ライブラリです。コードは自動微分可能であり、GPU/TPU で実行可能です。また、Haiku で構築された機械学習拡張モデルもいくつか提供されています。

重要性 – JAX の JIT コンパイルと自動微分を活用することで、研究者は以下のことが可能です：

• 流体シミュレーションと相互作用するニューラルネットワークの学習 (e.g., surrogate models, closure models)。
• CFD ソルバー上で直接、勾配ベースの最適化やデータ同化を行う。
• コードを変更せずに CPU, GPU, または TPU で同じコードを実行する。

主な構成要素

• jax_cfd.base – ステガードグリッド (Arakawa‑C/MAC) 上の有限体積/差分法ソルバー。
• jax_cfd.spectral – アンチエイリアシング機能を備えた擬似スペクトル渦度ソルバー。
• jax_cfd.ml – ML 拡張 CFD モデル (Haiku ニューラルネットワーク) およびサンプル プリトレーニング モデル。
• jax_cfd.data – シミュレーションデータのロード、処理、および可視化のためのユーティリティ (Xarray + Pillow)。

数値計算の焦点

• 非定常乱流。
• 空間離散化：ステガードグリッド FV/Diff または擬似スペクトル渦度。
• 時間離散化：1次明示的移流、暗黙的/明示的拡散。
• 圧力求解：共役勾配法または FFT ベースの対角化 (周期境界条件のみ)。
• 移流スキーム：2次 Van Leer。
• 乱流モデル：Smagorinsky 渦粘性モデル。

始め方

1. コアパッケージのインストール：

   pip install jax-cfd
   

2. オプションの追加機能 (ML またはデータユーティリティ) のインストール：

   pip install jax-cfd[ml]   # Haiku + ML helpers
   pip install jax-cfd[data] # Xarray, Pillow, etc.
   

3. notebooks/ フォルダ内の用意された Colab ノートブックを実行して、以下を確認できます：
   • 2-D ステガードグリッド FV デモ
   • 2-D 擬似スペクトル乱流デモ
   • チャンネルフローの例
   • PNAS 2021 論文の結果の再現 (データ分析 & モデル推論)

プロジェクトのステータス – メンテナンス終了 (2024 年 6 月) とされています。README には、JAX‑Fluids, PhiFlow, および Exponax などの、現在アクティブにメンテナンスされている代替案が記載されています。

典型的なユースケース

• 微分可能な CFD の学術研究、例：乱流閉鎖モデルの学習や逆観測演算子の学習。
• 流体ソルバーを介した勾配が必要な ML‑CFD ハイブリッドモデルのプロトタイピング。
• 流体動力学のワークロードにおける JAX の性能ベンチマーク。

引用 – 有限体積法コードまたは ML モデルを使用する場合、PNAS 2021 論文を引用してください。スペクトル法コードを使用する場合、arXiv 2022 pre‑print を引用してください。

関連プロジェクト – README には、他の微分可能な CFD ライブラリ (PhiFlow, Autograd fluid examples) や、より広範な JAX 科学ツール (JAX‑MD, JAX‑DFT, jax‑cosmo, Veros) がリストアップされています。

---

上記の内容は、リポジトリの README から直接引用されています。