XLB: A differentiable massively parallel lattice Boltzmann library in Python

TL;DR

XLB is a Python library for lattice Boltzmann simulations that is highly scalable, differentiable, and compatible with machine learning tools, enabling advanced fluid dynamics modeling and physics-based ML applications.

Contribution

It introduces a differentiable, scalable, and extensible Python lattice Boltzmann library built on JAX, supporting large-scale simulations and integration with ML workflows.

Findings

Scales to billions of cells with giga-scale update rates

Supports multi-GPU and TPU distributed computing

Enables physics-based machine learning and inverse problems

Abstract

The lattice Boltzmann method (LBM) has emerged as a prominent technique for solving fluid dynamics problems due to its algorithmic potential for computational scalability. We introduce XLB library, a Python-based differentiable LBM library based on the JAX platform. The architecture of XLB is predicated upon ensuring accessibility, extensibility, and computational performance, enabling scaling effectively across CPU, TPU, multi-GPU, and distributed multi-GPU or TPU systems. The library can be readily augmented with novel boundary conditions, collision models, or multi-physics simulation capabilities. XLB's differentiability and data structure is compatible with the extensive JAX-based machine learning ecosystem, enabling it to address physics-based machine learning, optimization, and inverse problems. XLB has been successfully scaled to handle simulations with billions of cells, achieving giga-scale lattice updates per second. XLB is released under the permissive Apache-2.0 license and is available on GitHub at https://github.com/Autodesk/XLB.