<p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/2fb66959-10d1-49e2-bf0f-efac91463000.png" alt="" width="300"> </p>

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上进行原位渲染。

展示

<p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/881be387-22fe-4009-931d-08ce0016a06c.gif" width="800"> </p> <p align="center"> 使用<a href="https://github.com/loliverhennigh/PhantomGaze">PhantomGaze</a>库在GPU上进行原位渲染（无I/O）。使用KBC格子玻尔兹曼模拟，约1000万个单元的NACA翼型流动。 </p> <p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/d47ca839-6b2b-4f38-82c9-e700545464cf.png" alt="" width="500"> </p> <p align="center"> 风洞中的<a href=https://www.epc.ed.tum.de/en/aer/research-groups/automotive/drivaer >DrivAer模型</a>，使用KBC格子玻尔兹曼模拟，约3.17亿个单元 </p> <p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/7d649bdc-cdcf-45a3-9034-853cd5f5acd1.png" alt="" width="700"> </p> <p align="center"> 建筑物内外的气流（约4亿个单元） </p> <p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/4a74faee-3b36-446b-8b6c-899428266d1e.png" alt="" width="900"> </p> <p align="center> 流体密度场从初始状态到在时间步200处通过深度学习优化形成"XLB"模式的各个阶段（详见论文） </p> <br> <p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/a4e0137f-7306-4882-bfc7-ad1b725b80ba.gif" alt="" width="500"> </p> <p align="center"> 雷诺数为100,000的顶盖驱动腔流（约2500万个单元） </p>

功能

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。

• 🌊 **自由表面流：**模拟具有自由表面的流动，如水波和液滴。

• 📡 **电磁波传播：**模拟电磁波的传播。

• 🛩️ **超音速流：**模拟超音速流。

• 🌊🧱 **流固耦合：**模拟流体和固体物体之间的相互作用。

• 🧩 **多相流模拟：**模拟多种不相溶流体的流动。

• 🔥 **燃烧：**模拟燃烧过程和反应流。

• 🪨 **粒子流和离散元方法：**结合基于粒子的方法模拟颗粒流和颗粒物流动。

• 🔧 **更好的几何处理流程：**改进复杂几何体的处理和预处理以用于模拟。