tumaer/JAXFLUIDS

Differentiable Fluid Dynamics Package

What it solves

JAX-Fluidsは、機械学習と計算流体力学 (CFD) のギャップを埋めるために設計されました。3D圧縮、単相および二相流に対して完全微分可能なソルバーを提供し、研究者が自動微分を用いて数値モデルの端から端までの最適化を行うことを可能にします。

How it works

JAXで完全に構築されたこのパッケージは、デカルト座標系グリッド上で有限体積法を用いてNavier-Stokes方程式を解きます。WENOやTENOなどの様々な高階自適応空間再構成法をサポートし、複数のRiemannソルバー、および陽的な時間ステップ法をサポートしています。複雑な流体については、level-set法およびdiffuse-interface法を用いて、二相流シミュレーションおよび没入型固体境界 (immersed solid boundaries) 制御をサポートしています。

Who it's for

MLとCFDの交差点における研究者向けです。CPU, GPU, および TPU hardware へのスケーリング- 規模のHPC- 集団 (HPC clusters) 大規模- 規模の- n/a (註：原文- 在 CPU, GPU, 与 TPU 上运行。)

Highlights

  • Fully Differentiable: モデル最適化のための自動微分を可能にします。

  • High Performance: HPCクラスター上で効率的にスケーリングします。最大 512 NVIDIA A100 GPUs と 2048 TPU-v3 cores にてテスト済みです。

  • Versatile Flow Models: 3D圧縮流体、二相流シミュレーション、および没入型固体境界をサポートします。

  • Hardware Agnostic: CPU, GPU, および TPU 上で動作します。