JAX-Fluids

简介

JAX-Fluids 是一个基于 Google JAX 框架的全自动可微 CFD（计算流体力学）求解器，专注于端到端可微的流体仿真。相比传统 CFD 求解器（如 OpenFOAM），它允许对仿真参数（几何、边界条件、物性）做梯度反传，因此天然契合物理信息机器学习（Physics-Informed ML）、优化设计与神经网络耦合仿真。

核心特性

• 全自动可微：对所有仿真参数可求导，无需手动伴随方程
• 多GPU支持：JAX 的 SPMD 抽象天然支持多 GPU/TPU 并行，数据并行零改动
• 高阶精度：WENO-JS 格式，支持 3rd~5th 阶精度
• 多相流：支持两相流（VOF 方法）
• 与神经网络耦合：可作为 Neural ODE / PINN 的物理前向求解器

适用场景

• 形状优化（机翼、管道、钝体）：梯度直接反传至设计参数
• 物理信息神经网络（PINN / PIAN）前向求解器
• 数据驱动湍流建模（closure term 学习）
• 端到端物理仿真-控制联合学习

使用注意

• License 标注 MIT，但原始论文 CPC 部分标注 GPLv3，建议仅作算法参考不直接嵌入闭源商业产品
• 适合学术研究；工业级验证建议与传统求解器（OpenFOAM/ANSYS）交叉验证
• 国内访问 GitHub 需代理，可配合 Gitee 镜像使用