XLB: A differentiable massively parallel lattice Boltzmann library in Python

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.

References in corpus (6)

• Enhancing Computational Fluid Dynamics with Machine Learning
• Surveying the reach and maturity of machine learning and artificial intelligence in astronomy
• JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows
• JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science
• Learned Turbulence Modelling with Differentiable Fluid Solvers: Physics-based Loss-functions and Optimisation Horizons
• On the accuracy and performance of the lattice Boltzmann method with 64-bit, 32-bit and novel 16-bit number formats

Cited by in corpus (4)

• Reduced-order modeling of unsteady fluid flow using neural network ensembles
• Optimized Flow Control based on Automatic Differentiation in Compressible Turbulent Channel Flows
• Differentiable Cosmological Hydrodynamics for Field-Level Inference and High Dimensional Parameter Constraints
• A Lattice Boltzmann Method for Non-Newtonian Blood Flow in Coiled Intracranial Aneurysms