JAX-CFD：基于 JAX 的可微计算流体力学研究工具

JAX-CFD 是 Google 用 JAX 实现的实验性 CFD 研究项目，面向非定常湍流，结合自动微分、机器学习与 GPU/TPU 加速。

这个工具是什么

JAX-CFD 是 Google 的 JAX 生态下的实验性研究项目，用于探索机器学习、自动微分与硬件加速器（GPU/TPU）在计算流体力学中的潜力，代码以 JAX 实现，相关方法发表于 PNAS 2021 论文《Machine learning accelerated computational fluid dynamics》。项目按子模块组织：jax_cfd.base 提供基于 JAX 的有限体积/有限差分核心方法，jax_cfd.spectral 提供伪谱方法，jax_cfd.ml 基于 JAX 与 Haiku 提供机器学习增强模型，jax_cfd.data 基于 Xarray 与 Pillow 提供数据处理工具。README 明确说明该项目已不再维护，并建议考虑 JAX-Fluids、PhiFlow 或 Exponax 等替代方案。

核心能力

• 提供交错网格（Arakawa C/MAC）上的有限体积与有限差分方法，以及用于涡量的伪谱方法
• 支持自动微分与 GPU/TPU 加速，便于将 CFD 求解器与机器学习模型结合
• 包含 jax_cfd.ml 子模块，提供基于 Haiku 的机器学习增强模型与预训练模型推理示例
• 提供 jax_cfd.data 数据处理工具，用于准备、评估与后处理 JAX-CFD 生成的数据

科研中的典型用法

• 在 notebooks 目录的 Colab 示例中运行二维交错网格有限体积模拟、伪谱强迫湍流模拟与槽道流模拟
• 复现 PNAS 论文结果：运行预训练机器学习模型，并进行数据分析与评估
• 借助自动微分与可微求解器，将 CFD 模拟嵌入机器学习训练或反问题研究流程

快速上手

基础安装执行 pip install jax-cfd，仅依赖 NumPy、SciPy 与 JAX；如需机器学习或数据处理子模块，使用 pip install jax-cfd[ml]、pip install jax-cfd[data] 或 pip install jax-cfd[complete]。入门可从 notebooks 目录的演示笔记本开始，覆盖有限体积、伪谱与槽道流等模拟设置。注意项目已不再维护。

生态与相近工具

同领域可考虑 README 推荐的 JAX-Fluids、PhiFlow 与 Exponax，它们同为可与深度学习结合的可微 CFD 代码；README 还列出使用 JAX-CFD 的项目（如变分数据同化的学习逆观测算子）。由于 JAX-CFD 已停止维护且仅支持周期性边界条件与一阶时间离散，新项目宜优先评估上述替代工具。

项目信息

• 仓库：google/jax-cfd
• 星标：964　·　派生：142
• 主要语言：Jupyter Notebook
• 许可证：Apache-2.0
• 最近提交：2026-09-09

本文属于《学科研究工具知识库 · 工程与材料》第 9 篇（按 GitHub 星标排序）。收录标准：该学科公开可获取的开源研究工具，星标与活跃度为主要参考，不代表对其性能或适用性的背书；选型前请结合自身场景评估。