Autodesk/XLB
XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML
XLB – Differentiable Lattice‑Boltzmann Library (Python)
What it is – XLB is a Python library that implements the Lattice Boltzmann Method (LBM) for fluid‑dynamics simulation and makes every step differentiable. It can run on CPUs, single‑GPU, multi‑GPU, and TPU hardware via three back‑ends:
- JAX – leverages JAX’s XLA compiler and works on CPUs, GPUs, and TPUs.
- NVIDIA Warp – a custom CUDA‑kernel framework for state‑of‑the‑art single‑GPU performance.
- Neon – Autodesk’s multi‑GPU programming model that adds grid‑refinement (multi‑resolution) and multi‑GPU scaling.
Because the kernels are differentiable, XLB can be embedded in machine‑learning pipelines (e.g., gradient‑based design optimisation, physics‑informed neural networks, or adjoint‑based shape optimisation).
Quick start (from the README)
# CPU‑only (no special hardware)
pip install xlb
# Single‑GPU with Warp (fastest single‑GPU path)
pip install "xlb[warp]"
# JAX on CUDA‑enabled GPU
pip install "xlb[cuda]"
# JAX on TPU
pip install "xlb[tpu]"
# Multi‑GPU / multi‑resolution with Neon (requires Linux & Python 3.11‑3.14)
git clone https://github.com/Autodesk/XLB.git
cd XLB
pip install -r requirements.txt
pip install '.[neon]'
Do not mix Warp and Neon in the same environment.
Core capabilities (as listed in the README)
| Category | Features |
|---|---|
| Back‑ends | JAX, NVIDIA Warp, Neon (multi‑GPU, grid refinement) |
| LBM models | BGK, KBC (stable at high Reynolds), Smagorinsky LES |
| Lattice sets | D2Q9, D3Q19, D3Q27 (required for KBC) |
| Boundary conditions | Equilibrium, full/half‑way bounce‑back, Zou/He, regularized, extrapolation outflow, interpolated bounce‑back, hybrid, etc. |
| Differentiability | All collision and streaming kernels are JAX‑compatible, enabling gradients w.r.t. initial fields, geometry, or material parameters |
| Scalability | Single‑GPU (Warp), multi‑GPU (Neon), distributed multi‑GPU/TPU via JAX shard‑map |
| Precision | Mixed‑precision (separate store vs compute) |
| I/O & visualisation | VTK (binary/ASCII) via PyVista, HDF5/XDMF with gzip, in‑situ GPU rendering with PhantomGaze, checkpointing with Orbax, image slices, voxelisation via trimesh |
| Machine‑learning integration | Works with Flax, Haiku, Optax, etc.; provides differentiable boundary conditions |
| Roadmap highlights | Out‑of‑core simulation, lossless GPU compression, fluid‑thermal coupling, adjoint shape optimisation, ML‑accelerated solvers, reduced‑order models |
Typical use‑cases
- Physics‑based ML research – embed a fluid simulator in a neural‑network loss and back‑propagate through the simulation to learn control policies or inverse designs.
- Gradient‑based design optimisation – optimise shapes (airfoils, car bodies, building layouts) by differentiating the drag/lift objectives w.r.t. geometry.
- High‑performance CFD – run large‑scale LBM simulations (hundreds of millions of cells) on a single GPU (Warp) or across many GPUs (Neon/JAX).
- Multi‑resolution studies – use Neon’s grid‑refinement to concentrate resolution where needed (e.g., near walls) while keeping the overall memory footprint manageable.
- Educational / research prototyping – pure‑Python API with NumPy‑like syntax makes it easy to experiment without writing CUDA kernels.
Where to learn more
- Examples – the repository’s
examples/folder contains ready‑to‑run scripts for a wind‑turbine, NACA airfoil, DrivAer car model, building airflow, lid‑driven cavity, and a differentiable optimisation demo. - Paper – XLB: A differentiable massively parallel lattice Boltzmann library in Python (Computer Physics Communications, 2024) provides benchmark numbers and validation details.
- Citations – If you publish work that uses XLB, cite the main paper and, when using grid‑refinement, also cite the Neon‑related conference papers listed in the README.
License
Apache 2.0 – free for academic and commercial use.
Bottom line – XLB is a genuine, production‑grade library that bridges high‑performance fluid simulation with modern differentiable programming. It is well‑suited for researchers and engineers who need fast, scalable CFD that can be differentiated for machine‑learning‑driven design or optimisation tasks.
Related
- Project
- Project
- Project
- Project