XLB: A differentiable massively parallel lattice Boltzmann library in Python

Mohammadmehdi Ataei,Hesam Salehipour

DOI: https://doi.org/10.1016/j.cpc.2024.109187

2024-04-02

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 <a class="link-external link-https" href="https://github.com/Autodesk/XLB" rel="external noopener nofollow">this https URL</a>.

Computational Physics,Computational Engineering, Finance, and Science,Machine Learning

What problem does this paper attempt to address?

The problem this paper attempts to address is the development of an efficient, scalable, and user-friendly differentiable large-scale parallel lattice Boltzmann method (LBM) library—XLB. Specifically, the goals of the paper include:

1. **Improving Accessibility**: XLB is written in Python and provides an interface similar to Numpy, ensuring that users can easily get started and quickly adopt it.
2. **Enhancing Scalability**: XLB employs an object-oriented programming model, allowing users to easily extend the library's functionality to meet different research needs.
3. **Maintaining High Performance**: Despite its user-friendly design, XLB still focuses on performance and scalability, making it suitable for a wide range of scenarios from entry-level to advanced resource-intensive applications.

XLB is based on the JAX platform, supporting automatic differentiation and distributed computing, which gives it broad application potential in fields such as physics-based machine learning, optimization, and inverse problems. The paper demonstrates the advantages of XLB in handling large-scale fluid simulation tasks by introducing its architecture, implementation details, and performance evaluation.