*Result*: JAX‐LaB: A High‐Performance, Differentiable Lattice Boltzmann Library for Modeling Multiphase Fluid Dynamics in Geosciences and Engineering.
*Further Information*
*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 (Ataei & Salehipour, 2024, https://doi.org/10.1016/j.cpc.2024.109187), and it is built on the JAX framework (Bradbury et al., 2018, http://github.com/jax‐ml/jax). 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 >107 ${ >} 1{0}^{7}$ while maintaining low spurious currents. Fluid wetting is achieved using the "improved" virtual density scheme (Q. Li et al., 2019, https://doi.org/10.1103/physreve.100.053313), 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 CO2 $\mathrm{C}{\mathrm{O}}_{\mathrm{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. Plain Language Summary: Several natural and engineering processes can be classified as multiphase flows involving either a single chemical entity (for instance, water) in multiple phases (typically liquid and gas) or distinct chemical entities such as water and air coexisting in distinct phases. For many such applications, numerical simulations are necessary, as experimental investigations can be challenging. To accurately model multiphase flows, it is essential not only to solve the governing equations that describe fluid dynamics, but also to properly capture phase distributions and account for complex fluid behaviors such as liquid‐vapor interfaces, which continuously form, deform, and break during flow, and wettability, which governs fluid‐solid interactions. In this work, we introduce a Python‐based library to model multiphase systems. The library performs computation using JAX, a high‐performance Python library for scientific computations. We describe the various numerical techniques implemented in the library and then benchmark them against several well‐known problems with analytical solutions, highlighting the applicability of JAX‐LaB for modeling many different multiphase fluid flows. We then showcase several exemplar problems where JAX‐LaB can be used in the geosciences and beyond. Finally, we test the performance of the library on single and multiple Graphics Processing Units (GPUs). Key Points: We present a JAX‐based differentiable multiphase Lattice Boltzmann library, scalable across CPUs, GPUs, and distributed systemsA modified Shan‐Chen pseudopotential incorporating arbitrary equations of state and tunable wettability and surface tension is employedThe library features decoupled compute and write precision, and achieves giga‐scale lattice updates per second on multi‐GPU setups [ABSTRACT FROM AUTHOR]
Copyright of Journal of Advances in Modeling Earth Systems is the property of Wiley-Blackwell and its content may not be copied or emailed to multiple sites without the copyright holder's express written permission. Additionally, content may not be used with any artificial intelligence tools or machine learning technologies. However, users may print, download, or email articles for individual use. This abstract may be abridged. No warranty is given about the accuracy of the copy. Users should refer to the original published version of the material for the full abstract. (Copyright applies to all Abstracts.)*