XLB：基于JAX的加速格子玻尔兹曼方法库

XLB：推动计算流体动力学与机器学习的融合

在现代科学计算和工程模拟领域，流体动力学一直是一个充满挑战yet又不可或缺的研究方向。随着人工智能技术的快速发展，将传统的数值方法与新兴的机器学习技术相结合，已成为一个极具前景的研究热点。在这样的背景下，Autodesk公司开发并开源了XLB库，这是一个基于JAX的全可微分格子玻尔兹曼方法(LBM)库，旨在为物理驱动的机器学习应用提供强大而灵活的工具。

XLB的核心特性

XLB库的设计理念是将高性能计算与易用性完美结合。它具有以下几个关键特性：

与JAX生态系统的无缝集成：XLB可以轻松地与JAX的机器学习库（如Flax、Haiku和Optax等）集成，为研究人员提供了丰富的工具集。
可微分的LBM内核：这一特性使XLB能够在可微分物理和深度学习应用中发挥重要作用。
出色的可扩展性：XLB能够在分布式多GPU系统上扩展，支持在数百个GPU上运行包含数十亿个单元的大规模模拟。
多样化的LBM边界条件和碰撞核：XLB支持多种LBM边界条件和碰撞核，为研究人员提供了丰富的选择。
用户友好的接口：XLB完全用Python编写，提供了高度可访问的接口，使用户能够轻松扩展库并快速设置和运行新的模拟。
利用JAX Array和Shardmap：XLB采用了新的JAX统一数组类型和JAX shardmap，为用户提供了类似numpy的接口，使他们可以专注于语义，将性能优化交给编译器。
平台多样性：同一份XLB代码可以在多核CPU、单GPU或多GPU系统、TPU上运行，还支持在多GPU系统或TPU Pod切片上进行分布式运行。
可视化功能：XLB提供了多种可视化选项，包括使用PhantomGaze库进行GPU上的原位渲染。

XLB Logo

XLB的应用展示

XLB库的强大功能可以通过其多样化的应用场景得到充分展现。以下是一些引人注目的例子：

NACA翼型周围的流动：使用KBC格子玻尔兹曼模拟方法，XLB能够模拟包含约1000万个单元的NACA翼型周围流动。这个模拟利用PhantomGaze库在GPU上进行原位渲染，无需I/O操作，展示了XLB的高效计算和可视化能力。
风洞中的DrivAer模型：XLB成功模拟了风洞中约3.17亿个单元的DrivAer模型。这个大规模模拟充分展示了XLB在处理复杂几何形状和大规模计算方面的能力。

DrivAer模型模拟

建筑物内外的气流模拟：XLB能够模拟建筑物内外约4亿个单元的气流，这对于建筑设计和城市规划具有重要意义。
深度学习优化的流体密度场：XLB展示了如何通过深度学习优化，将初始状态的流体密度场逐步演化成"XLB"图案。这个例子生动地展示了XLB在物理驱动机器学习应用中的潜力。
高雷诺数lid-driven cavity流动：XLB成功模拟了雷诺数为100,000的lid-driven cavity流动，包含约2500万个单元。这个例子展示了XLB在处理高雷诺数流动时的稳定性和精确性。

Lid-driven Cavity流动

XLB的技术能力

XLB库提供了丰富的技术能力，以满足各种复杂的流体动力学模拟需求：

LBM模型：
- BGK碰撞模型（标准LBM碰撞模型）
- KBC碰撞模型（适用于高雷诺数流动的无条件稳定模型）
机器学习集成：
- 与JAX的机器学习库生态系统轻松集成
- 可微分的LBM内核
- 可微分的边界条件
格子模型：
- D2Q9
- D3Q19
- D3Q27（必须用于KBC模拟运行）
计算能力：
- 分布式多GPU支持
- 混合精度支持（存储vs计算）
- 核外支持（即将推出）
输出功能：
- 基于PyVista库的二进制和ASCII VTK输出
- 使用PhantomGaze库进行原位渲染
- 基于Orbax的分布式异步检查点
- 图像输出
- 使用trimesh的3D网格体素化
边界条件：
- 平衡BC
- 全程回弹BC
- 半程回弹BC
- 无操作BC
- Zouhe BC
- 正则化BC
- 外推出流BC
- 插值回弹BC

安装和使用XLB

要使用XLB，用户首先需要安装JAX和其他依赖项。安装过程因硬件平台而异，XLB提供了详细的安装指南。以下是一个简单的安装和运行示例：

# 安装依赖
pip install pyvista numpy matplotlib Rtree trimesh jmp orbax-checkpoint termcolor

# 克隆XLB仓库并运行示例
git clone https://github.com/Autodesk/XLB
cd XLB
export PYTHONPATH=.
python3 examples/CFD/cavity2d.py

XLB的未来发展

XLB团队有着雄心勃勃的发展计划，包括：

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

此外，XLB团队还列出了一系列愿望清单，包括自由表面流动、电磁波传播、超音速流动、流固耦合、多相流模拟、燃烧模拟等多个具有挑战性的研究方向。

结语

XLB作为一个开源项目，正在为计算流体动力学和物理驱动的机器学习领域带来革命性的变化。它不仅提供了高性能的模拟能力，还通过其可微分性为深度学习应用铺平了道路。随着项目的不断发展和社区的积极参与，我们可以期待XLB在未来为更多前沿研究和工程应用提供强大支持。

对于有志于探索计算流体动力学与机器学习交叉领域的研究者和工程师来说，XLB无疑是一个值得关注和使用的强大工具。我们鼓励感兴趣的读者访问XLB的GitHub仓库，深入了解这个令人兴奋的项目，并考虑为其未来发展做出贡献。让我们一起期待XLB在推动科学计算和人工智能融合方面继续发挥重要作用！