作者：德米特里·科奇科夫、杰米·A·史密斯、彼得·诺加德、吉迪恩·德累斯顿、阿亚·阿利瓦、斯蒂芬·霍耶
JAX-CFD 是一个实验研究项目，旨在探索机器学习、自动微分和硬件加速器 (GPU/TPU) 在计算流体动力学方面的潜力。它是在 JAX中实现的。
要了解有关我们一般方法的更多信息，请阅读我们的论文机器学习加速计算流体动力学 (PNAS 2021)。
“notebooks”目录包含几个使用 JAX-CFD 代码的演示。
-
不同模拟设置的演示：
-
重现我们的 PNAS 论文的结果：
JAX-CFD 围绕子模块进行组织：
jax_cfd.base：CFD 的核心有限体积/差分方法，用 JAX 编写。jax_cfd.spectral：CFD 的核心伪谱方法，用 JAX 编写。jax_cfd.ml：CFD 机器学习增强模型，用 JAX 和Haiku编写。jax_cfd.data：数据处理实用程序，用于准备、评估和后处理使用 JAX-CFD 创建的数据，用 Xarray和 Pillow编写。
基本安装pip install jax-cfd仅需要 NumPy、SciPy 和 JAX。要安装其他子模块的依赖项，请使用pip install jax-cfd[ml],
pip install jax-cfd[data]或pip install jax-cfd[complete]。
JAX-CFD 目前专注于非定常湍流：
- 空间离散化：
- 交错网格（“Arakawa C”或“MAC”网格）上的有限体积/差分方法，每个单元中心的压力和相应面上定义的速度分量。
- 涡度的伪谱方法使用非线性项的抗混叠滤波技术来保持稳定性。
- 时间离散化：目前只有一阶时间离散化，使用显式时间步进进行平流，使用隐式或显式时间步进进行扩散。
- 压力求解：CG 或实值 FFT 的快速对角化（适用于周期性边界条件）。
- 边界条件：目前仅支持周期性边界条件。
- 平流：我们实施二阶精确的“Van Leer”方案。
- 闭包：我们目前实施 Smagorinsky 涡流粘度模型。
TODO：添加一个笔记本，更深入地解释我们的数值模型。
从长远来看，我们有兴趣扩展 JAX-CFD 以实现与相关研究相关的方法，例如，
- 共置网格
- 替代边界条件（例如，非周期边界和浸入边界方法）
- 高阶时间步长
- 几何多重网格
- 稳态模拟（例如 RANS）
- 跨多个 TPU/GPU 的分布式模拟
我们欢迎在这些方面进行合作！在开始重要工作之前，请联系（通过 GitHub 或通过电子邮件）进行协调。
其他与深度学习兼容的可微分 CFD 代码：
- PhiFlow支持 TensorFlow、PyTorch 和 JAX
- Autograd 中的流体模拟
JAX 科学：
我们错过了什么吗？请告诉我们！
如果您使用我们的有限体积法 (FVM) 或 ML 模型，请引用：
@article{Kochkov2021-ML-CFD,
author = {Kochkov, Dmitrii and Smith, Jamie A. and Alieva, Ayya and Wang, Qing and Brenner, Michael P. and Hoyer, Stephan},
title = {Machine learning{\textendash}accelerated computational fluid dynamics},
volume = {118},
number = {21},
elocation-id = {e2101784118},
year = {2021},
doi = {10.1073/pnas.2101784118},
publisher = {National Academy of Sciences},
issn = {0027-8424},
URL = {https://www.pnas.org/content/118/21/e2101784118},
eprint = {https://www.pnas.org/content/118/21/e2101784118.full.pdf},
journal = {Proceedings of the National Academy of Sciences}
}
如果您使用我们的光谱代码，请引用：
@article{Dresdner2022-Spectral-ML,
doi = {10.48550/ARXIV.2207.00556},
url = {https://arxiv.org/abs/2207.00556},
author = {Dresdner, Gideon and Kochkov, Dmitrii and Norgaard, Peter and Zepeda-Núñez, Leonardo and Smith, Jamie A. and Brenner, Michael P. and Hoyer, Stephan},
title = {Learning to correct spectral methods for simulating turbulent flows},
publisher = {arXiv},
year = {2022},
copyright = {arXiv.org perpetual, non-exclusive license}
}
本地安装以进行开发：
git clone https://github.com/google/jax-cfd.git
cd jax-cfd
pip install jaxlib
pip install -e ".[complete]"
然后手动运行测试套件：
pytest -n auto jax_cfd --dist=loadfile --ignore=jax_cfd/base/validation_test.py