XLB
A differentiable, massively parallel Lattice Boltzmann library in Python for physics-based machine learning and fluid dynamics simulations.
What is XLB?
XLB is a differentiable, massively parallel Lattice Boltzmann Method library in Python for simulating fluid dynamics. It solves complex flow problems efficiently while providing differentiable kernels that enable seamless integration with machine learning frameworks for physics-based AI applications.
Researchers and engineers in computational fluid dynamics, physics-based machine learning, and scientific computing who need scalable, differentiable simulations for optimization and AI-driven design.
Developers choose XLB for its unique combination of full differentiability, massive parallel scalability across multiple hardware backends, and user-friendly Python interface that bridges traditional CFD with modern ML workflows.
Overview
XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML
Use Cases
Best For
- Physics-informed neural networks requiring differentiable fluid simulations
- Large-scale computational fluid dynamics research on multi-GPU clusters
- Gradient-based optimization in aerodynamic design and shape optimization
- Integrating traditional CFD simulations with JAX-based machine learning pipelines
- Educational purposes for learning Lattice Boltzmann methods with modern Python tools
- Rapid prototyping of fluid dynamics experiments with in-situ GPU visualization
Not Ideal For
- Projects requiring multi-GPU acceleration with the Warp backend, as it currently only supports single GPU
- Simulations involving free surface or multiphase flows, as these features are on the wishlist and not yet implemented
- Teams developing on macOS who need GPU-accelerated simulations, since JAX backend lacks GPU support on MacOS
- Applications needing real-time, interactive fluid simulations with low latency, due to the computational intensity of LBM methods
Pros & Cons
Pros
Supports JAX and Warp backends for optimal performance on CPUs, GPUs, and TPUs, allowing the same code to run on diverse hardware without modification.
Provides fully differentiable LBM kernels and boundary conditions, enabling seamless integration with JAX-based machine learning libraries like Flax and Optax for gradient-based optimization.
Scales to billions of cells on distributed multi-GPU systems using the JAX backend, making it suitable for large-scale research and industrial simulations.
Written entirely in Python with a numpy-like interface, simplifying setup and extension of simulations without requiring deep C++ or CUDA expertise.
Offers in-situ GPU rendering via PhantomGaze, reducing I/O overhead and enabling real-time visualization during simulations without saving to disk.
Cons
The Warp backend currently only supports single GPU, which restricts multi-GPU acceleration for users preferring this backend over JAX, as noted in the README.
Advanced features like free surface flows, multiphase simulations, and combustion are not yet implemented and are listed in the wishlist, limiting scope for complex applications.
On macOS, GPU acceleration is unavailable with the JAX backend, forcing users to rely on CPU-only simulations or switch platforms for better performance.
While user-friendly, XLB requires familiarity with Lattice Boltzmann methods and CFD concepts, which can be a steep learning curve for newcomers without prior experience.
Frequently Asked Questions
Related Projects
ALXGoogle Research
OryxProbabilistic reasoning and statistical analysis in TensorFlow
BRAXMassively parallel rigidbody physics simulation on accelerator hardware.
MctxMonte Carlo tree search in JAX
Found a gem we're missing?
Open-Awesome is built by the community, for the community. Submit a project, suggest an awesome list, or help improve the catalog on GitHub.