XLB: A differentiable massively parallel lattice Boltzmann library in Python

doi:10.1016/j.cpc.2024.109187

2024-07-01

PRE

Mohammadmehdi Ataei *

Hesam Salehipour

DOI：10.1016/j.cpc.2024.109187

原文链接

原文求助

摘要

En 中文

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, multiGPU, 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. Program summary Program Title: XLB CPC Library link to program files: https://doi .org /10 .17632 /n25h65vw59 .1 Developer's repository link: https://github .com /Autodesk /XLB Licensing provisions: Apache -2.0 Programming language: Python Nature of problem: XLB is a scalable Lattice Boltzmann simulation library based on JAX. It is designed to be accessible, extensible, and high performance. The library can be readily augmented with novel boundary conditions, collision models, or simulation capabilities. Moreover, XLB is differentiable, offering the unique advantage of integration with JAX's extensive machine learning ecosystem, and tackling physics -based machine learning, optimization, and inverse problems. Solution method: XLB utilizes the JAX library for just -in -time (JIT) compilation and distributed computing, enabling it to scale effectively across distributed multi-GPU/multi-TPU systems. XLB's programming model combines an object -oriented interface for easy extensibility. The library can be readily augmented with novel boundary conditions, collision models, or other simulation capabilities.

Keyword:

Open source software

Lattice Boltzmann method

JAX

Machine learning

Differentiable programming

Scientific computing

Computational fluid dynamics

High performance computing