tumaer/JAXFLUIDS

Differentiable Fluid Dynamics Package

What it solves

JAX-Fluids 是為了彌合機器學習與計算流體力學 (CFD) 之間的差距而設計的。它為 3D 可壓縮單相與兩相流提供了一個完全可微的解算器，允許研究人員使用自動微分對數值模型進行端到端優化。

How it works

該套件完全基於 JAX 建構，在笛卡兒網格上使用有限體積法求解 Navier-Stokes 方程式。它支援多種高階自適應空間重建 (例如 WENO 和 TENO)，多個 Riemann 解算器，以及顯式時間步進。對於複雜流體，它包含用於兩相模擬與浸入式固體邊界 (immersed solid boundaries) 的 level-set 方法與 diffuse-interface 方法。

Who it's for

它適用於在 ML 與 CFD 交界處工作的研究人員，這些研究人員需要一個高性能、可微的解算器，且能跨 CPU、GPU 與 TPU 硬體進行擴展，包括大規模 HPC 集群。

Highlights

  • Fully Differentiable: 實現了用於模型優化的自動微分。

  • High Performance: 在 HPC 集群上高效擴展，已在多達 512 NVIDIA A100 GPUs 與 2048 TPU-v3 cores 上進行過測試。

  • Versatile Flow Models: 支援 3D 可壓縮流體、兩相模擬與浸入式固體邊界。

  • Hardware Agnostic: n/a (註：原文- 在 CPU, GPU, 與 TPU 上運行。)

  • Hardware Agnostic: 在 CPU, GPU, 與 TPU 上運行。