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 方法。

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: 在 CPU, GPU, 和 TPU 上运行。