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.

翻译：本文介绍了JAX-LaB，一个基于Python的可微分格子玻尔兹曼模拟库，专为水文、地质及工程多孔介质环境中的多相与多物理场流体动力学问题建模而设计。该库作为XLB的扩展构建于JAX框架之上，提供了高性能且硬件无关的实现，能够无缝集成机器学习库，并在CPU、多GPU配置及分布式环境中高效扩展。多相相互作用采用Shan-Chen赝势方法建模，并结合状态方程（EOS）复现符合Maxwell构造的密度，从而在密度比$>10^7$的流动模拟中保持高精度并抑制虚假流。流体润湿通过“改进型”虚拟密度方案实现，该方案能精确控制平面与曲面上的接触角，同时消除Shan-Chen虚拟密度方案中出现的非物理薄膜。此方案直接集成于相互作用力计算中，无需分别处理流体-流体与流体-固体作用力。我们通过全面的解析基准测试验证了该库的准确性与性能，包括拉普拉斯定律、平行板毛细上升以及通道内多组分并流。随后，我们将该代码应用于多组分与多相流动的多个场景，包括渗透率估算、水饱和枫丹白露砂岩中超临界$CO_2$注入以及球堆积几何结构的特征曲线获取。最后，我们在单节点与分布式系统上评估了代码的单GPU性能与多GPU扩展能力。该库基于Apache许可证开源，可通过https://github.com/piyush-ppradhan/JAX-LaB获取。