XLB：用于基于物理的机器学习的可微分大规模并行格子玻尔兹曼Python库

XLB是一个完全可微分的2D/3D格子玻尔兹曼方法（LBM）库，利用硬件加速。它基于JAX库构建，专门设计用于以计算高效和可微分的方式解决流体动力学问题。其独特的功能组合使其成为物理机器学习应用的极其合适工具。

引用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}
}

主要特点

与JAX生态系统集成： 该库可以轻松与JAX强大的机器学习库生态系统集成，如Flax、Haiku、Optax等。
可微分LBM核心： XLB提供可微分的LBM核心，可用于可微分物理和深度学习应用。
可扩展性： XLB能够在分布式多GPU系统上扩展，使得在数百个GPU上执行数十亿单元的大规模模拟成为可能。
支持各种LBM边界条件和核心： XLB支持多种LBM边界条件和碰撞核心。
用户友好界面： 完全用Python编写，XLB强调高度可访问的界面，使用户能够轻松扩展库并快速设置和运行新的模拟。
利用JAX Array和Shardmap： 该库incorporates新的JAX统一数组类型和JAX shardmap，为用户提供类似numpy的接口。这使用户可以专注于语义，将性能优化留给编译器。
平台多样性： 相同的XLB代码可以在多核CPU、单个或多个GPU系统、TPU上执行，还支持在多GPU系统或TPU Pod切片上进行分布式运行。
可视化： XLB提供多种可视化选项，包括使用PhantomGaze在GPU上进行原位渲染。

展示

使用PhantomGaze库在GPU上进行原位渲染（无I/O）。使用KBC格子玻尔兹曼模拟，约1000万个单元的NACA翼型流动。

风洞中的DrivAer模型，使用KBC格子玻尔兹曼模拟，约3.17亿个单元

建筑物内外的气流（约4亿个单元）

雷诺数为100,000的顶盖驱动腔流（约2500万个单元）

功能

LBM

BGK碰撞模型（标准LBM碰撞模型）
KBC碰撞模型（高雷诺数流动无条件稳定）

机器学习

易于与JAX的机器学习库生态系统集成
可微分LBM核心
可微分边界条件

格子模型

D2Q9
D3Q19
D3Q27（必须用于KBC模拟运行）

计算能力

分布式多GPU支持
混合精度支持（存储与计算）
核外支持（即将推出）

输出

基于PyVista库的二进制和ASCII VTK输出
使用PhantomGaze库进行原位渲染
基于Orbax的分布式异步检查点保存
图像输出
使用trimesh的3D网格体素化

边界条件

**平衡边界条件：**在此边界条件下，假设流体分布处于平衡状态。可用于设置指定的速度或压力。
**全反弹边界条件：**在此边界条件下，流体分布的速度被完全反射回流体边界侧，导致边界处流体速度为零。
**半反弹边界条件：**类似于全反弹边界条件，在此边界条件下，流体分布的速度被部分反射回流体边界侧，导致边界处流体速度不为零。
**无操作边界条件：**在此边界条件下，流体分布可以不经任何反射或修改直接通过边界。
**Zouhe边界条件：**此边界条件用于在边界处施加指定的速度或压力分布。
**正则化边界条件：**此边界条件用于在边界处施加指定的速度或压力分布。该边界条件比Zouhe边界条件更稳定，但计算成本更高。
**外推出流边界条件：**一种使用外推法避免强烈波反射的出流边界条件。
**插值反弹边界条件：**由Bouzidi提出的用于格子玻尔兹曼方法模拟的插值反弹边界条件。

安装指南

要使用XLB，你首先需要使用以下命令安装JAX和其他依赖项：

请参考https://github.com/google/jax获取最新的安装文档。下表摘自JAX的Github页面。

硬件	安装指令
CPU	`pip install -U "jax[cpu]"`
x86_64上的NVIDIA GPU	`pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`
Google TPU	`pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`
AMD GPU	使用Docker或从源代码构建。
Apple GPU	按照Apple的说明操作。

**注意：**由于Metal后端缺乏对某些操作的支持，我们在Apple GPU上执行XLB时遇到了挑战。我们建议在Mac OS上使用CPU后端。我们将在未来测试XLB在Apple GPU上的表现，并相应更新此部分。

安装依赖项：

pip install pyvista numpy matplotlib Rtree trimesh jmp orbax-checkpoint termcolor

运行示例：

git clone https://github.com/Autodesk/XLB
cd XLB
export PYTHONPATH=.
python3 examples/CFD/cavity2d.py

路线图

正在进行中（WIP）

注意：一些正在进行中的功能可以在XLB仓库的分支中找到。如果要为这些功能做贡献，请联系我们。

🚀 **Warp后端：**通过结合JAX利用Warp框架实现最先进的性能。
🌐 **网格细化：**实现自适应网格细化技术以提高模拟精度。
⚡ **使用Neon + Warp的多GPU加速：**使用Neon的数据结构改善扩展性。
💾 **核外计算：**支持超出可用GPU内存的模拟，适用于CPU+GPU一致性内存模型，如NVIDIA的Grace超级芯片。
🗜️ GPU加速无损压缩和解压缩：实现高性能无损压缩和解压缩技术，用于大规模模拟和性能改进。
🌡️ **流体-热仿真能力：**将热传递和热效应纳入流体模拟。
🎯 **伴随型形状和拓扑优化：**实现基于梯度的优化技术用于设计优化。
🧠 **机器学习加速模拟：**利用机器学习加速模拟并提高精度。
📉 **使用机器学习的降阶模型：**开发数据驱动的降阶模型，实现高效准确的模拟。

愿望清单

欢迎对这些功能做出贡献。请为愿望清单项目提交PR。

🌊 **自由表面流：**模拟具有自由表面的流动，如水波和液滴。
📡 **电磁波传播：**模拟电磁波的传播。
🛩️ **超音速流：**模拟超音速流。
🌊🧱 **流固耦合：**模拟流体和固体物体之间的相互作用。
🧩 **多相流模拟：**模拟多种不相溶流体的流动。
🔥 **燃烧：**模拟燃烧过程和反应流。
🪨 **粒子流和离散元方法：**结合基于粒子的方法模拟颗粒流和颗粒物流动。
🔧 **更好的几何处理流程：**改进复杂几何体的处理和预处理以用于模拟。

XLB