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

Abstract

We introduce JAX-LaB, a differentiable, Python-based Lattice Boltzmann simulation library designed for modeling multiphase and multiphysics fluid dynamics problems in hydrologic, geologic, and engineered porous media settings. The library is designed as an extension to XLB, and it is built on the JAX framework. The library delivers a performant, hardware-agnostic implementation that seamlessly integrates with machine learning libraries and scales efficiently across CPUs, multi-GPU setups, and distributed environments. Multiphase interactions are modeled using the Shan-Chen pseudopotential method, coupled with an equation of state (EOS) to reproduce densities consistent with Maxwell's construction, enabling accurate simulation of flows with density ratios $> 10^7$ while maintaining low spurious currents. Fluid wetting is achieved using the "improved" virtual density scheme, which enables precise control of contact angle on flat and curved surfaces, while eliminating non-physical films seen in the Shan-Chen virtual density scheme. This scheme integrates directly into the interaction force calculations, removing the need to handle fluid-fluid and fluid-solid forces separately. We validate the library's accuracy and performance through comprehensive analytical benchmarks, including Laplace's law, capillary rise in parallel plates, and multi-component cocurrent flow in a channel. We then use the code for several applications involving multicomponent and multiphase flows, including permeability estimation, injection of supercritical $CO_2$ in a water-saturated Fontainebleau sandstone, and obtaining the characteristic curves for a sphere pack geometry. Finally, the single-GPU performance and multi-GPU scaling of the code are evaluated on both single-node and distributed systems. The library is open-source under the Apache license and available at https://github.com/piyush-ppradhan/JAX-LaB.