JAX-MPM: a learning-augmented differentiable meshfree framework for GPU-accelerated Lagrangian simulation and geophysical inverse modeling.

Saved in:
Bibliographic Details
Title: JAX-MPM: a learning-augmented differentiable meshfree framework for GPU-accelerated Lagrangian simulation and geophysical inverse modeling.
Authors: Du, Honghui1 (AUTHOR), He, QiZhi1 (AUTHOR) qzhe@umn.edu
Source: Engineering with Computers. Jun2026, Vol. 42 Issue 3, p1-25. 25p.
Subjects: Material point method, Meshfree methods, Automatic differentiation, Graphics processing units, Lagrange equations, Inversion (Geophysics), Machine learning, Benchmarking (Management), Deep learning
Abstract: Differentiable programming has emerged as a powerful paradigm in scientific computing, enabling automatic differentiation through simulation pipelines and naturally supporting both forward and inverse modeling. We present JAX-MPM, a general-purpose differentiable meshfree solver based on the material point method (MPM) and implemented in the modern JAX ecosystem. The framework adopts a hybrid Eulerian–Lagrangian formulation to capture large deformations, frictional contact, and inelastic material behavior, with emphasis on geomechanics and geophysical hazard applications. Leveraging GPU acceleration and automatic differentiation, JAX-MPM enables efficient gradient-based optimization directly through its time-stepping solvers and supports joint training of physical models with deep learning to infer unknown system conditions and uncover hidden constitutive parameters. We validate JAX-MPM through a series of 2D and 3D benchmark simulations, including dam-break and granular collapse problems, demonstrating both numerical accuracy and GPU-accelerated performance. Results show that a high-resolution 3D granular cylinder collapse with 2.7 million particles completes 1000 time steps in approximately 22 s (single precision) and 98 s (double precision) on a single GPU. Beyond high-fidelity forward modeling, we demonstrate the framework's inverse modeling capabilities through tasks such as velocity field reconstruction and the estimation of spatially varying friction from sparse data. In particular, JAX-MPM introduces a differentiable observation layer that unifies data assimilation from both Lagrangian (particle-based) and Eulerian (region-based) observations, and can be seamlessly coupled with neural network representations. These results establish JAX-MPM as a unified and scalable differentiable meshfree platform that advances fast physical simulation and data assimilation for complex solid and geophysical systems. [ABSTRACT FROM AUTHOR]
Copyright of Engineering with Computers is the property of Springer Nature 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.)
Database: Engineering Source
Description
Abstract:Differentiable programming has emerged as a powerful paradigm in scientific computing, enabling automatic differentiation through simulation pipelines and naturally supporting both forward and inverse modeling. We present JAX-MPM, a general-purpose differentiable meshfree solver based on the material point method (MPM) and implemented in the modern JAX ecosystem. The framework adopts a hybrid Eulerian–Lagrangian formulation to capture large deformations, frictional contact, and inelastic material behavior, with emphasis on geomechanics and geophysical hazard applications. Leveraging GPU acceleration and automatic differentiation, JAX-MPM enables efficient gradient-based optimization directly through its time-stepping solvers and supports joint training of physical models with deep learning to infer unknown system conditions and uncover hidden constitutive parameters. We validate JAX-MPM through a series of 2D and 3D benchmark simulations, including dam-break and granular collapse problems, demonstrating both numerical accuracy and GPU-accelerated performance. Results show that a high-resolution 3D granular cylinder collapse with 2.7 million particles completes 1000 time steps in approximately 22 s (single precision) and 98 s (double precision) on a single GPU. Beyond high-fidelity forward modeling, we demonstrate the framework's inverse modeling capabilities through tasks such as velocity field reconstruction and the estimation of spatially varying friction from sparse data. In particular, JAX-MPM introduces a differentiable observation layer that unifies data assimilation from both Lagrangian (particle-based) and Eulerian (region-based) observations, and can be seamlessly coupled with neural network representations. These results establish JAX-MPM as a unified and scalable differentiable meshfree platform that advances fast physical simulation and data assimilation for complex solid and geophysical systems. [ABSTRACT FROM AUTHOR]
ISSN:01770667
DOI:10.1007/s00366-026-02320-6