XLB: A differentiable massively parallel lattice Boltzmann library in Python (2311.16080v3)

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.

• The paper showcases differentiable programming in CFD by integrating JAX for automatic differentiation to support physics-informed optimization.
• The paper demonstrates XLB's scalability by achieving giga-scale lattice updates per second across various hardware platforms including CPUs, GPUs, and TPUs.
• The paper validates its approach using benchmarks like the Taylor-Green Vortex and 3D lid-driven cavity, paving the way for integrated machine learning applications in fluid dynamics.

XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python

This paper introduces XLB, a Python-based library leveraging the Lattice Boltzmann Method (LBM) for fluid dynamics simulations. XLB is developed on the JAX platform, designed for high-level numerical computing, which allows seamless integration with machine learning frameworks while maintaining computational efficiency. The library's architecture supports scalability across various hardware, including CPU, TPU, multi-GPU, and distributed systems, making it adaptable for broad research applications, from physics-based machine learning to large-scale simulations.

Key Contributions

The XLB library contributes several advancements in computational fluid dynamics (CFD):

1. Differentiable Programming: Built on JAX, XLB supports automatic differentiation, beneficial for optimization tasks and machine learning models that integrate fluid dynamics simulations. This compatibility allows exploring physics-informed neural networks and other innovative applications in scientific computing.
2. Scalability and Performance: XLB demonstrates scalability across multiple hardware configurations. It achieves giga-scale lattice updates per second, proving its suitability for large-scale simulations, evidenced by successful execution on massive grid sizes with billions of cells.
3. Accessibility and Extensibility: By employing Python and resembling Numpy interfaces, XLB ensures accessibility for users familiar with scientific computing. Its object-oriented design allows for extensive customization, enabling users to introduce new models or boundary conditions with ease.
4. Benchmark Validation: The library was validated against well-known benchmarks like the Taylor-Green Vortex, 3D lid-driven cavity, and flow over a cylinder, covering a wide range of fluid flow conditions. These validations confirm the library's accuracy and reliability in both steady and unsteady flow simulations.
5. Machine Learning Integration: XLB facilitates advanced machine learning applications in CFD, such as surrogate modeling and inverse problem solving. The examples provided in the paper, including flow reconstruction using neural networks, underline its potential in integrating learning-based methodologies with traditional CFD approaches.

Numerical Results and Analysis

XLB's differentiated features are supported by robust numerical results. The paper discusses extensive benchmarking tests, showing the effectiveness of various collision models and boundary conditions. It highlights the significant reduction in error when deploying machine learning correctors to improve coarse-grained simulation accuracy. Moreover, derivative-based methods using neural networks in flow optimization tasks, such as forming specific patterns within certain simulation steps, are made feasible through XLB's integration of JAX.

Theoretical and Practical Implications

The development of XLB adds significant value both theoretically and practically to CFD. Theoretically, the integration with differentiable programming and machine learning paradigms opens new research avenues. Practical implications include the ability for researchers to experiment with flow optimization and real-time simulation adjustments in a user-friendly yet high-performance environment.

Future Directions

While the paper elaborates on current features, future enhancements to XLB are anticipated. These include adjoint-based optimization, real-time data assimilation, and expanded computational backends to enhance performance further. Out-of-core computing features to handle simulations of unprecedented size also highlight the library's growth trajectory.

Overall, XLB represents a powerful tool for modern CFD tasks, bridging the gap between traditional physics-based simulations and emerging data-driven techniques. Its release under the Apache-2.0 license encourages broad adoption and community collaboration, fostering rapid innovation in computational fluid dynamics.