# JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modeling

arXiv: https://arxiv.org/abs/2507.04192 · HTML: https://mlsanity.com/paper/2507.04192/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 108 references that are in the map.
- First posted on arXiv 6 Jul 2025.

## Builds on (most cited first, 30 of 33)

- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [On the Spectral Bias of Neural Networks](https://mlsanity.com/paper/1806.08734/) (2018)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Scientific Machine Learning through Physics-Informed Neural Networks: Where we are and What's next](https://mlsanity.com/paper/2201.05624/) (2022)
- [Physics-informed neural networks (PINNs) for fluid mechanics: A review](https://mlsanity.com/paper/2105.09506/) (2021)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [PAC-NeRF: Physics Augmented Continuum Neural Radiance Fields for Geometry-Agnostic System Identification](https://mlsanity.com/paper/2303.05512/) (2023)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Transfer learning enhanced physics informed neural network for phase-field modeling of fracture](https://mlsanity.com/paper/1907.02531/) (2019)
- [ChainQueen: A Real-Time Differentiable Physical Simulator for Soft Robotics](https://mlsanity.com/paper/1810.01054/) (2018)
- [A Differentiable Programming System to Bridge Machine Learning and Scientific Computing](https://mlsanity.com/paper/1907.07587/) (2019)
- [Physics informed deep learning for computational elastodynamics without labeled data](https://mlsanity.com/paper/2006.08472/) (2020)
- [Physics-Informed Neural Networks for Multiphysics Data Assimilation with Application to Subsurface Transport](https://mlsanity.com/paper/1912.02968/) (2019)
- [An End-to-End Differentiable Framework for Contact-Aware Robot Design](https://mlsanity.com/paper/2107.07501/) (2021)
- [Interfacing Finite Elements with Deep Neural Operators for Fast Multiscale Modeling of Mechanics Problems](https://mlsanity.com/paper/2203.00003/) (2022)
- [JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science](https://mlsanity.com/paper/2212.00964/) (2022)
- [DaXBench: Benchmarking Deformable Object Manipulation with Differentiable Physics](https://mlsanity.com/paper/2210.13066/) (2022)
- [Differentiable Programming for Earth System Modeling](https://mlsanity.com/paper/2208.13825/) (2022)
- [Differentiable Programming for Differential Equations: A Review](https://mlsanity.com/paper/2406.09699/) (2024)
- [Graph Neural Network-based surrogate model for granular flows](https://mlsanity.com/paper/2305.05218/) (2023)
- [A Hybrid Deep Neural Operator/Finite Element Method for Ice-Sheet Modeling](https://mlsanity.com/paper/2301.11402/) (2023)
- [Scalable and modular material point method for large-scale simulations](https://mlsanity.com/paper/1909.13380/) (2019)
- [JAX-DIPS: Neural bootstrapping of finite discretization methods and application to elliptic problems with discontinuities](https://mlsanity.com/paper/2210.14312/) (2022)
- [4D topology optimization: Integrated optimization of the structure and self-actuation of soft bodies for dynamic motions](https://mlsanity.com/paper/2302.00905/) (2023)
- [Inverse analysis of granular flows using differentiable graph neural network simulator](https://mlsanity.com/paper/2401.13695/) (2024)
- [A Neural Material Point Method for Particle-based Emulation](https://mlsanity.com/paper/2408.15753/) (2024)

## Built on by (most cited first, 5 of 5)

- [Differentiable Inverse Modeling with Physics-Constrained Latent Diffusion for Heterogeneous Subsurface Parameter Fields](https://mlsanity.com/paper/2512.22421/) (2025)
- [Estimating Parameter Fields in Multi-Physics PDEs from Scarce Measurements](https://mlsanity.com/paper/2509.00203/) (2025)
- [Convex Low-resource Accent-Robust Language Detection in Speech Recognition](https://mlsanity.com/paper/2605.23235/) (2026)
- [Benchmarking Optimizers to Solve Inverse Problems with Differentiable Physics Simulators](https://mlsanity.com/paper/2609.13819/) (2026)
- [CraftSPH: A high-accuracy and composable differentiable SPH solver implemented in PyTorch](https://mlsanity.com/paper/2609.38208/) (2026)

## Abstract (arXiv)

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 architecture. The solver adopts a hybrid Eulerian-Lagrangian framework 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 seconds (single precision) and 98 seconds (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 accommodates 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.