JAX-Fluids 开源项目使用教程

JAX-Fluids 开源项目使用教程

JAXFLUIDS 项目地址: https://gitcode.com/gh_mirrors/ja/JAXFLUIDS

1. 项目介绍

JAX-Fluids 是一个完全可微分的计算流体动力学（CFD）求解器，专为3D可压缩单相和两相流体设计。该项目旨在促进机器学习（ML）和计算流体动力学（CFD）交叉领域的研究。JAX-Fluids 完全使用 JAX 编写，支持 CPU/GPU/TPU，并能够进行自动微分，从而实现数值模型的端到端优化。

JAX-Fluids 通过 JAX 原语进行并行化，能够在现代高性能计算集群上高效扩展（已在多达512个 NVIDIA A100 GPU 和多达2048个 TPU-v3 核心上测试）。

2. 项目快速启动

安装 JAX-Fluids

在安装 JAX-Fluids 之前，请确保您的 pip 是最新版本：

pip install --upgrade pip

CPU 版本安装

如果您只需要 CPU 支持，可以运行以下命令：

pip install --upgrade "jax[cpu]"
git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install .

GPU 和 CPU 版本安装

如果您需要 GPU 和 CPU 支持，首先需要安装带有 GPU 支持的 JAX。有两种方法可以实现这一点：通过 pip 安装 CUDA 和 cuDNN，或者自行安装 CUDA 和 cuDNN。

推荐使用 pip 安装 CUDA 和 cuDNN：

pip install --upgrade "jax[cuda12]"
git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install -e .

快速启动示例

JAX-Fluids 提供了五个 Jupyter Notebook 示例，帮助您快速上手。以下是一个简单的 1D Sod 激波管模拟示例：

import jaxfluids as jf

# 初始化模拟参数
params = jf.SimulationParameters(
    domain_size=(100,),
    cell_size=(0.1,),
    time_step=0.01,
    end_time=0.2
)

# 初始化物理模型
model = jf.EulerModel()

# 初始化网格
grid = jf.CartesianGrid(params)

# 初始化状态
state = jf.InitialState(grid, model)

# 运行模拟
sim = jf.Simulation(params, model, grid, state)
sim.run()

3. 应用案例和最佳实践

应用案例

1. 航天飞机在马赫2下的模拟：使用浸入式固体边界方法通过水平集进行模拟。
2. 激波与气泡相互作用：使用扩散界面方法进行模拟，大约在 TPUv3-64 上处理 800M 个单元。
3. 水中的空气泡激波诱导坍塌：通过视频展示模拟结果，链接：YouTube 视频。

最佳实践

• 优化并行计算：利用 JAX 的并行化能力，在 GPU 或 TPU 上运行大规模模拟。
• 自动微分：通过 JAX 的自动微分功能，优化数值模型参数。
• 高阶空间重构：使用 WENO-3/5/7 等高阶空间重构方法，提高模拟精度。

4. 典型生态项目

• JAX：JAX-Fluids 的核心依赖，提供自动微分和并行计算能力。
• CUDA & cuDNN：用于 GPU 加速，提升计算性能。
• Jupyter Notebook：提供交互式开发环境，方便快速实验和调试。

通过以上模块，您可以快速了解 JAX-Fluids 的基本功能和使用方法，并开始进行自己的流体动力学模拟。

JAXFLUIDS 项目地址: https://gitcode.com/gh_mirrors/ja/JAXFLUIDS