JAX-WSPM: A GPU-accelerated parallel framework based on the JAX library for modeling water flow and solute transport in unsaturated porous media using an implicit finite element method

Accurate simulation of water flow and solute transport in unsaturated porous media requires solving complex, nonlinear partial differential equations. Traditionally, implicit finite element methods have been used due to their robustness and stability. However, they are well known for their computational expense when addressing coupled dynamics. In this study, we present JAX-WSPM, a GPU-accelerated framework built with the JAX library that leverages just-in-time (JIT) compilation and automatic differentiation (AD) to reduce computational cost and improve scalability for coupled water flow and solute transport systems in porous media. We use an implicit finite element method to solve the Richards equation, which models water flow in unsaturated media, and the transport equation. JAX-WSPM implements two complementary strategies for computing water fluxes that are critical for the solute transport equation: one based on conventional finite element formulations and another that supports automatic differentiation. In addition, an adaptive time-stepping strategy is employed to optimize performance. Benchmark tests are conducted to examine the accuracy, efficiency, and scalability of JAX-WSPM in solving the Richards equation and the coupled flow-solute transport system. The results confirm the accuracy and efficiency of the framework and demonstrate significant speedups when comparing the GPU-accelerated JAX-WSPM implementation to both the CPU-based JAX-WSPM and serial Python implementations. For example, when performing simulations on a mesh with 1.03 million degrees of freedom, the GPU-accelerated solver achieved a speedup of approximately 107× relative to the serial Python implementation running on a single CPU. JAX-WSPM is available at https://github.com/NourEddine-Toutlini/JAX-WSPM, offering a flexible, user-friendly, and high-performance tool for simulations in porous media.