私は、逆設計や学習済みクロージャーといったMLループの中での微分可能Navier Stokesシミュレーションのために、JAXベースのCFDフレームワークを構築してきました。
最適化や学習パイプラインの内部に組み込めるように、完全なソルバースタックを微分可能に保つことが目標です。
設計上の選択:
- 外部依存なしの完全なJAXネイティブ
- CPUを最初にしたベクトル化実装
- 速度・圧力・渦度の各場を通じたエンドツーエンドの微分可能性
- Navier Stokes（投影法）およびLBM（D2Q9）サポート
- 幾何形状の取り扱いのための滑らかなマスクを用いたBrinkmanスタイルの強制項
現在:
- 投影と圧力補正を用いた2D非圧縮Navier Stokesソルバー
- LBMソルバーを同一フレームワークに統合
- 性能はCPU制約で、グリッドに依存
- 128x128で約560 FPS
- 512x96で約300 FPS
- パイプライン全体を通した微分可能な流れ場
- ソルバーのループ内部でニューラル演算子および学習済み補正を行うためのフック
本当の価値はここにあります:
- 逆設計：幾何形状が流れにマッピングされ、勾配が幾何形状へと逆伝播する
- 乱流または残差クロージャーを、直接ソルバー内で学習する
- CFDを、MLシステム向けの微分可能なデータ生成器として使う
- 勾配の流れを壊さずに、物理と学習モデルをハイブリッド化する
多くのCFDおよびMLパイプラインでは、ソルバーを依然としてブラックボックスとして扱っています。そのため、勾配に基づく設計が難しい、あるいは不可能になります。
AeroJAXは、物理の構造を保ったまま、パイプライン全体を微分可能にする試みです。
[リンク] [コメント]
