JAX-LaB: A High-Performance, Differentiable, Lattice Boltzmann Library for Modeling Multiphase Fluid Dynamics in Geosciences and Engineering

2025-06-21

AI Summary

To address the low computational efficiency and non-differentiability of multiphase, multi-physics flow simulations in porous media, this paper introduces XLB—a high-performance, differentiable lattice Boltzmann method (LBM) framework built on JAX. Methodologically, XLB integrates an enhanced Shan-Chen pseudopotential model with a Maxwell-consistent equation of state to enable robust high-density-ratio two-phase flows; incorporates a virtual-density wetting model for precise contact-angle control and elimination of unphysical liquid films; and leverages automatic differentiation alongside optimized external-force formulations to support end-to-end machine learning integration and seamless CPU/GPU cross-platform parallelism. Validated against Laplace’s law, capillary rise, and multicomponent co-flow benchmarks, XLB significantly suppresses spurious currents and demonstrates excellent strong scalability on single- and multi-GPU systems. The open-source implementation is released under the Apache 2.0 license.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Low Level & Physics-based Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Data management and stream processing for Web, mobile and wireless applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems

Abstract

We present JAX-LaB, a differentiable, Python-based Lattice Boltzmann library for simulating multiphase and multiphysics flows in hydrologic, geologic, and engineered porous media. Built as an extension of the XLB library, JAX-LaB utilizes JAX for computations and offers a performant, hardware-agnostic implementation that integrates seamlessly with machine learning workflows and scales efficiently across CPUs, GPUs, and distributed systems. Multiphase interactions are modeled using the Shan-Chen pseudopotential method, which is coupled with an equation of state and an improved forcing scheme to obtain liquid-vapor densities that are consistent with Maxwell's construction, enabling simulations of systems with very large density ratios while maintaining minimal spurious currents. Wetting is handled using the "improved" virtual density scheme, which allows precise control of contact angles and eliminates non-physical films seen in other Shan-Chen wetting methods. We validate the library through several analytical benchmarks, such as Laplace's law, capillary rise, and cocurrent multicomponent flow, and demonstrate some exemplary use cases for the library. We also report single- and multi-GPU performance scaling of the library. The library is open-source under the Apache license and available at https://github.com/piyush-ppradhan/JAX-LaB.

Problem

• Simulate multiphase flows in geosciences and engineering
• Model large density ratios with minimal spurious currents
• Control wetting behavior and eliminate non-physical films

Innovation

• Differentiable Lattice Boltzmann library for multiphase flows
• Uses JAX for hardware-agnostic, scalable computations
• Improved Shan-Chen method for large density ratios