google/jax-cfd

JAX‑CFD: JAX 中的可微分计算流体动力学

是什么 – 一个开源研究库，使用 JAX 实现核心 CFD 求解器（有限体积法和伪谱法），因此代码可以自动求导，并能在 GPU/TPU 上运行。它还附带了一些使用 Haiku 构建的机器学习增强模型。

为什么重要 – 通过利用 JAX 的即时编译 (JIT) 和自动微分功能，研究人员可以：

• 训练与流体模拟交互的神经网络 (e.g., surrogate models, closure models)。
• 直接在 CFD 求解器上进行基于梯度的优化或数据同化。
• 在不更改代码的情况下，在 CPU, GPU, 或 TPU 上运行相同的代码。

核心组件

• jax_cfd.base – 基于交错网格 (Arakawa‑C/MAC) 的有限体积/差分求解器。
• jax_cfd.spectral – 具有抗混叠功能的伪谱涡度求解器。
• jax_cfd.ml – ML 增强的 CFD 模型 (Haiku 神经网络) 和示例预训练模型。
• jax_cfd.data – 用于加载、处理和可视化数据模拟的工具 (Xarray + Pillow)。

数值计算重点

• 非定常湍流。
• 空间离散化：交错网格 FV/Diff 或伪谱涡度。
• 时间离散化：一阶显式平流，隐式/显式扩散。
• 压力求解：共轭梯度法或基于 FFT 的对角化 (仅限周期性 BCs)。
• 平流方案：二阶 Van Leer。
• 湍流闭合：Smagorinsky 涡粘性模型。

如何开始

1. 安装核心包：

   pip install jax-cfd
   

2. 添加可选的额外组件以使用 ML 或数据工具：

   pip install jax-cfd[ml]   # Haiku + ML helpers
   pip install jax-cfd[data] # Xarray, Pillow, etc.
   

3. 运行 notebooks/ 文件夹中的现成 Colab notebooks，以查看：
   • 2-D 交错网格 FV 演示
   • cd 2-D 伪谱湍流演示
   • 通道流示例
   • 重现 PNAS 2021 论文的结果 (数据分析 & 模型推理)

项目状态 – 已标记为不再维护 (2024 年 6 月)。README 指向了积极维护的替代方案，例如 JAX‑Fluids, PhiFlow, 和 Exponax。

典型使用场景

• 学术研究上可微分 CFD，例如：学习湍流闭合模型或逆向观察算符。
• 构建需要通过流体求解器进行梯度的 ML‑CFD 混合模型原型。
• 在流体动力学工作负载上基准测试 JAX 性能。

引用

• 如果你使用有限体积法代码或 ML 模型，请引用 PNAS 2021 论文。
• 对于谱法代码，请引用 arXiv 2022 pre‑print。

相关项目 – README 列出了其他可微分 CFD 库 (PhiFlow, Autograd fluid examples) 和更广泛的 JAX 科学工具 (JAX‑MD, JAX‑DFT, jax‑cosmo, Veros)。

---

以上所有细节均直接取自仓库的 README。