JAX-LaB：用于地球科学与工程中多相流体动力学建模的高性能可微 Lattice Boltzmann 库

我们介绍了 JAX-LaB，这是一个可微的、基于 Python 的 Lattice Boltzmann 模拟库，旨在模拟水文、地质和人造多孔介质环境中的多相和多物理场流体动力学问题。该库被设计为 XLB 的扩展，并基于 JAX 框架构建。该库提供了高性能的、与硬件无关的实现，可与机器学习库无缝集成，并在 CPU、多 GPU 设置和分布式环境中高效扩展。多相相互作用使用 Shan-Chen 赝势方法进行建模，并结合状态方程 (EOS) 以重现与麦克斯韦构造一致的密度，从而能够在保持低虚假电流的同时，准确模拟密度比 $> 10^7$ 的流动。流体润湿采用“改进型”虚拟密度方案实现，该方案能够精确控制平坦和弯曲表面上的接触角，同时消除了 Shan-Chen 虚拟密度方案中出现的非物理薄膜。该方案直接整合到相互作用力计算中，无需单独处理流体-流体和流体-固体力。我们通过综合分析基准测试验证了该库的准确性和性能，包括 Laplace 定律、平行板中的毛细上升以及通道中的多组分同向流。随后，我们将该代码用于涉及多组分和多相流的几个应用中，包括渗透率估计、超临界 $CO_2$ 注入饱和水的 Fontainebleau 砂岩，以及获取球形颗粒几何形状的特征曲线。最后，在单节点和分布式系统上评估了该代码的单 GPU 性能和多 GPU 扩展性。该库在 Apache 许可下开源，可从 https://github.com/piyush-ppradhan/JAX-LaB 获取。