[](https://opensource.org/licenses/Apache-2.0)
[](https://star-history.com/#Autodesk/XLB)
# XLB：用于基于物理的机器学习的 Python 可微大规模并行 Lattice Boltzmann 库
XLB 是一个完全可微的 2D/3D Lattice Boltzmann Method (LBM) 库，利用了硬件加速。它支持 [JAX](https://github.com/google/jax)、[NVIDIA Warp](https://github.com/NVIDIA/warp) 和 [Neon](https://github.com/Autodesk/Neon) 后端，专门设计用于以计算高效和可微的方式解决流体动力学问题。其独特的功能组合使其成为基于物理的机器学习应用的极其合适的工具。借助 Warp 后端，XLB 提供了最先进的单 GPU 性能，而借助全新的 Neon 后端，它扩展到了多 GPU（单分辨率）。更重要的是，Neon 后端为多分辨率模拟提供了网格细化能力。
## 入门指南
要开始使用 XLB，您可以使用 pip 安装它。根据您的硬件和需求，有不同的安装选项：
### 基础安装（仅限 CPU）
```
pip install xlb
```
### 安装 Warp 支持（单 GPU）
对于 NVIDIA Warp 后端（单 GPU，最先进的性能）：
```
pip install "xlb[warp]"
```
### 安装 CUDA 支持（适用于 NVIDIA GPU）
此安装适用于带 CUDA 支持的 JAX 后端：
```
pip install "xlb[cuda]"
```
### 安装 TPU 支持
此安装适用于带 TPU 支持的 JAX 后端：
```
pip install "xlb[tpu]"
```
### 安装 Neon 支持
Neon 后端实现了多 GPU 稠密和单 GPU 多分辨率表示。
使用以下命令安装带 Neon 支持的 XLB：
```
git clone https://github.com/Autodesk/XLB.git
cd XLB
pip install -r requirements.txt
pip install '.[neon]'
```
**要求：** Neon wheel 在 **Linux x86_64** 和 **Linux ARM** 上支持 **Python 3.11** 到 **Python 3.14**。
**注意：** Neon 使用了 warp 的自定义分支。
### 注意事项：
- 对于 Mac 用户：请使用基础的 CPU 安装命令，因为 MacOS 上不支持 JAX 的 GPU 支持。
- 对于 Warp 后端（单 GPU）使用 `xlb[warp]`，或对于 Neon 后端（多 GPU / 多分辨率）使用 `xlb[neon]`。请勿在同一环境中同时安装两者。
- CUDA 和 TPU 的安装选项仅影响 JAX 后端。
要从源码安装最新的开发版本：
```
pip install git+https://github.com/Autodesk/XLB.git
```
各版本的更新日志可以在[这里](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md)找到。
有关入门示例，请参阅 [examples](https://github.com/Autodesk/XLB/tree/main/examples) 文件夹。
## 配套论文
有关该库的基准测试、验证以及更多详细信息，请参阅[配套论文](https://doi.org/10.1016/j.cpc.2024.109187)。
## 引用 XLB
如果您在研究中使用了 XLB，请引用以下论文：
```
@article{ataei2024xlb,
title={{XLB}: A differentiable massively parallel lattice {Boltzmann} library in {Python}},
author={Ataei, Mohammadmehdi and Salehipour, Hesam},
journal={Computer Physics Communications},
volume={300},
pages={109187},
year={2024},
publisher={Elsevier}
}
```
如果您在工作中使用了网格细化功能，请同时引用：
```
@inproceedings{mahmoud2024optimized,
title={Optimized {GPU} implementation of grid refinement in lattice {Boltzmann} method},
author={Mahmoud, Ahmed H and Salehipour, Hesam and Meneghin, Massimiliano},
booktitle={2024 IEEE International Parallel and Distributed Processing Symposium (IPDPS)},
pages={398--407},
year={2024},
organization={IEEE}
}
@inproceedings{meneghin2022neon,
title={Neon: A Multi-{GPU} Programming Model for Grid-based Computations},
author={Meneghin, Massimiliano and Mahmoud, Ahmed H. and Jayaraman, Pradeep Kumar and Morris, Nigel J. W.},
booktitle={Proceedings of the 36th IEEE International Parallel and Distributed Processing Symposium},
pages={817--827},
year={2022},
month={june},
doi={10.1109/IPDPS53621.2022.00084},
url={https://escholarship.org/uc/item/9fz7k633}
}
```
## 主要特性
- **多后端支持：** XLB 包含对 JAX、NVIDIA Warp 和 Neon 后端的支持，为 Lattice Boltzmann 模拟提供了*最先进的*性能。Warp 后端针对单 GPU 运行，而 Neon 后端则支持多 GPU 单分辨率和单 GPU 多分辨率模拟。
- **多分辨率网格细化：** 采用嵌套长方体网格的网格细化以及多种 kernel-fusion 策略，以在 Neon 后端上实现最佳性能。
- **与 JAX 生态系统集成：** 该库可以轻松与 JAX 强大的机器学习库生态系统集成，如 [Flax](https://github.com/google/flax)、[Haiku](https://github.com/deepmind/dm-haiku)、[Optax](https://github.com/deepmind/optax) 等。
- **可微 LBM Kernels：** XLB 提供了可用于可微物理和深度学习应用的可微 LBM kernels。
- **可扩展性：** XLB 能够使用 JAX 后端或 Neon 后端在分布式多 GPU 系统上进行扩展，从而能够在具有数十亿个 cell 的数百个 GPU 上执行大规模模拟。
- **支持各种 LBM 边界条件和 Kernels：** XLB 支持多种 LBM 边界条件和碰撞 kernels。
- **用户友好的界面：** XLB 完全使用 Python 编写，强调高度可访问的接口，允许用户轻松扩展库并快速设置和运行新的模拟。
- **利用 JAX Array 和 Shardmap：** 该库结合了全新的 JAX array 统一数组类型和 JAX shardmap，为用户提供了类似于 numpy 的接口。这使得用户可以仅专注于语义，而将性能优化留给编译器。
- **平台通用性：** 相同的 XLB 代码可以在多种平台上执行，包括多核 CPU、单 GPU 或多 GPU 系统、TPU，并且还支持在多 GPU 系统或 TPU Pod 切片上进行分布式运行。
- **可视化：** XLB 提供了多种可视化选项，包括使用 [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze) 在 GPU 上进行实时渲染。
## 展示
Simulation of a wind turbine based on the immersed boundary method.
On GPU in-situ rendering using PhantomGaze library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells.
DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells
Airflow into, out of, and within a building (~400 million cells)
The stages of a fluid density field from an initial state to the emergence of the "XLB" pattern through deep learning optimization at timestep 200 (see paper for details)
Lid-driven Cavity flow at Re=100,000 (~25 million cells)
## 功能
### LBM
- BGK 碰撞模型（标准 LBM 碰撞模型）
- KBC 碰撞模型（对于高雷诺数流动无条件稳定）
- 用于湍流建模的 Smagorinsky LES 亚网格模型
### 机器学习
- 轻松集成 JAX 的机器学习库生态系统
- 可微 LBM kernels
- 可微边界条件
### Lattice 模型
- D2Q9
- D3Q19
- D3Q27（必须用于 KBC 模拟运行）
### 计算能力
- Warp 后端的单 GPU 支持，具备最先进的性能
- 使用 Neon 后端和单分辨率网格的多 GPU 支持
- 使用 Neon 后端的单 GPU 网格细化支持
- 使用 JAX 后端的分布式多 GPU 支持
- 混合精度支持（存储对比计算）
- 用于多分辨率模拟的多种 kernel-fusion 性能策略
- 核外支持（即将推出）
### 输出
- 二进制和 ASCII VTK 输出（基于 PyVista 库）
- 用于多分辨率数据的 HDF5/XDMF 输出（带 gzip 压缩）
- 使用 [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze) 库进行实时渲染
- 基于 [Orbax](https://github.com/google/orbax) 的分布式异步 checkpointing
- 图像输出（包括多分辨率切片图像）
- 使用 trimesh 的 3D 网格体素化工具
### 边界条件
- **Equilibrium BC：** 在此边界条件中，假设流体群处于平衡状态。可用于设置预定义的速度或压力。
- **Full-Way Bounceback BC：** 在此边界条件中，流体群的速度被反射回边界的流体侧，导致边界处的流体速度为零。
- **Half-Way Bounceback BC：** 与 Full-Way Bounceback BC 类似，在此边界条件中，流体群的速度被部分反射回边界的流体侧，导致边界处的流体速度不为零。
- **Do Nothing BC：** 在此边界条件中，允许流体群穿过边界而不发生任何反射或修改。
- **Zouhe BC：** 此边界条件用于在边界处施加预定义的速度或压力分布。
- **Regularized BC：** 此边界条件用于在边界处施加预定义的速度或压力分布。此边界条件比 Zouhe BC 更稳定，但计算量更大。
- **Extrapolation Outflow BC：** 一种使用外推法避免强烈波反射的流出边界条件。
- **Interpolated Bounceback BC：** 用于表示弯曲边界的插值反弹边界条件。
- **Hybrid BC：** 将规则化方法和反弹方法结合，并带有可选的壁面距离插值，以提高弯曲几何形状的精度。
## 路线图
### 近期已完成
- ✅ **网格细化：** 通过 Neon 后端，实现具有嵌套长方体网格和多种 kernel-fusion 策略的多分辨率 LBM。
- ✅ **使用 [Neon](https://github.com/Autodesk/Neon) + Warp 进行多 GPU 加速：** 通过 Neon 的数据结构结合基于 Warp 的 kernels，为单分辨率设置提供多 GPU 支持。
### 正在进行的工作 (WIP)
*注意：部分正在进行的工作可以在 XLB 仓库的分支中找到。如需为这些功能做出贡献，请联系我们。*
- 💾 **核外计算：** 实现超过可用 GPU 内存的模拟，适用于 NVIDIA 的 Grace Superchips 等 CPU+GPU 统一内存模型（即将推出）。
- 🗜️ **GPU 加速的无损压缩与解压缩：** 实现高性能的无损压缩和解压缩技术，以支持更大规模的模拟并提高性能。
- 🌡️ **流体热模拟功能：** 将传热和热效应纳入流体模拟中。
- 🎯 **基于伴随的形状和拓扑优化：** 实现基于梯度的优化技术以进行设计优化。
- 🧠 **机器学习加速的模拟：** 利用机器学习加速模拟并提高准确性。
- 📉 **使用机器学习的降阶建模：** 开发数据驱动的降阶模型，以实现高效且准确的模拟。
### 愿望清单
*欢迎为这些功能做出贡献。请为愿望清单中的项目提交 PR。*
- 🌊 **自由表面流动：** 模拟带有自由表面的流动，例如水波和液滴。
- 📡 **电磁波传播：** 模拟电磁波的传播。
- 🛩️ **超音速流动：** 模拟超音速流动。
- 🌊🧱 **流固耦合：** 对流体与固体物体之间的相互作用进行建模。
- 🧩 **多相流模拟：** 模拟具有多种不混溶流体的流动。
- 🔥 **燃烧：** 模拟燃烧过程和反应流动。
- 🪨 **颗粒流和离散单元法：** 结合基于颗粒的方法处理颗粒和微粒流动。
- 🔧 **更好的几何处理流水线：** 改进模拟中复杂几何形状的处理和预处理。