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

IF 3.4 2区物理与天体物理 Q1 COMPUTER SCIENCE, INTERDISCIPLINARY APPLICATIONS

Computer Physics Communications Pub Date : 2025-09-19 DOI:10.1016/j.cpc.2025.109866

Nour-Eddine Toutlini , Azzeddine Soulaïmani , Abdelaziz Beljadid

{"title":"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","authors":"Nour-Eddine Toutlini , Azzeddine Soulaïmani , Abdelaziz Beljadid","doi":"10.1016/j.cpc.2025.109866","DOIUrl":null,"url":null,"abstract":"<div><div>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.</div><div>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 <span><span>https://github.com/NourEddine-Toutlini/JAX-WSPM</span><svg><path></path></svg></span>, offering a flexible, user-friendly, and high-performance tool for simulations in porous media.</div></div>","PeriodicalId":285,"journal":{"name":"Computer Physics Communications","volume":"318 ","pages":"Article 109866"},"PeriodicalIF":3.4000,"publicationDate":"2025-09-19","publicationTypes":"Journal Article","fieldsOfStudy":null,"isOpenAccess":false,"openAccessPdf":"","citationCount":"0","resultStr":null,"platform":"Semanticscholar","paperid":null,"PeriodicalName":"Computer Physics Communications","FirstCategoryId":"101","ListUrlMain":"https://www.sciencedirect.com/science/article/pii/S0010465525003686","RegionNum":2,"RegionCategory":"物理与天体物理","ArticlePicture":[],"TitleCN":null,"AbstractTextCN":null,"PMCID":null,"EPubDate":"","PubModel":"","JCR":"Q1","JCRName":"COMPUTER SCIENCE, INTERDISCIPLINARY APPLICATIONS","Score":null,"Total":0}

引用次数: 0

Abstract

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.

查看原文本刊更多论文

JAX- wspm：基于JAX库的gpu加速并行框架，用于用隐式有限元法模拟非饱和多孔介质中的水流和溶质输运

非饱和多孔介质中水流和溶质输运的精确模拟需要求解复杂的非线性偏微分方程。传统上，隐式有限元方法具有鲁棒性和稳定性。然而，在处理耦合动力学时，它们的计算费用是众所周知的。在这项研究中，我们提出了JAX- wspm，这是一个使用JAX库构建的gpu加速框架，利用即时（JIT）编译和自动区分（AD）来降低计算成本并提高多孔介质中耦合水流和溶质输运系统的可扩展性。我们使用隐式有限元方法来求解模拟非饱和介质中水流的Richards方程和输运方程。JAX-WSPM为计算对溶质输运方程至关重要的水通量实施了两种互补策略：一种基于传统有限元公式，另一种支持自动微分。此外，采用自适应时间步进策略来优化性能。通过基准测试，验证了JAX-WSPM在求解Richards方程和耦合流-溶质输运系统中的准确性、效率和可扩展性。当将gpu加速的JAX-WSPM实现与基于cpu的JAX-WSPM和串行Python实现进行比较时，结果证实了框架的准确性和效率，并展示了显著的速度提升。例如，当在具有103万个自由度的网格上执行模拟时，gpu加速求解器相对于在单个CPU上运行的串行Python实现实现实现的速度提高了大约107倍。JAX-WSPM可在https://github.com/NourEddine-Toutlini/JAX-WSPM上获得，它为多孔介质中的模拟提供了一种灵活的、用户友好的高性能工具。

本文章由计算机程序翻译，如有差异，请以英文原文为准。

求助全文

约1分钟内获得全文求助全文

来源期刊

Computer Physics Communications 物理-计算机：跨学科应用

CiteScore

12.10

自引率

3.20%

发文量

287

审稿时长

5.3 months

期刊介绍： The focus of CPC is on contemporary computational methods and techniques and their implementation, the effectiveness of which will normally be evidenced by the author(s) within the context of a substantive problem in physics. Within this setting CPC publishes two types of paper. Computer Programs in Physics (CPiP) These papers describe significant computer programs to be archived in the CPC Program Library which is held in the Mendeley Data repository. The submitted software must be covered by an approved open source licence. Papers and associated computer programs that address a problem of contemporary interest in physics that cannot be solved by current software are particularly encouraged. Computational Physics Papers (CP) These are research papers in, but are not limited to, the following themes across computational physics and related disciplines. mathematical and numerical methods and algorithms; computational models including those associated with the design, control and analysis of experiments; and algebraic computation. Each will normally include software implementation and performance details. The software implementation should, ideally, be available via GitHub, Zenodo or an institutional repository.In addition, research papers on the impact of advanced computer architecture and special purpose computers on computing in the physical sciences and software topics related to, and of importance in, the physical sciences may be considered.