JAX-MPM: Differentiable MPM Solver

• JAX-MPM is a differentiable meshfree solver based on a hybrid Lagrangian-Eulerian framework for large-deformation continuum simulations.
• It leverages reverse-mode automatic differentiation to enable efficient PDE-constrained inverse modeling and joint optimization of physical and neural parameters.
• The solver integrates JAX features like jit, vmap, and scan for GPU acceleration and robust handling of complex geomechanical phenomena.

Searching arXiv for JAX-MPM and closely related JAX-based differentiable simulation frameworks. Searching for the exact JAX-MPM paper and adjacent frameworks. JAX-MPM is a general-purpose differentiable meshfree solver within a hybrid Lagrangian-Eulerian framework, built on the material point method (MPM) and implemented entirely in JAX for large-deformation continuum simulation, GPU acceleration, and end-to-end automatic differentiation. It is tailored to complex continuum mechanics involving large deformations, frictional contact, and inelastic material behavior, with emphasis on geomechanics and geophysical hazard applications such as landslides, debris flows, granular collapse, and dam-break-like free-surface flows. Its central objective is to combine the robustness of MPM for extreme deformation with reverse-mode automatic differentiation through the full time-stepping pipeline, so that forward simulation, PDE-constrained inverse modeling, and joint optimization of physical and neural parameters can be handled within a single solver (Du et al., 6 Jul 2025).

1. Scope, motivation, and position within differentiable simulation

JAX-MPM was developed around two requirements. First, forward simulation in geomechanics and geophysical hazards must remain robust under severe distortion, fragmentation, and evolving contact or topology. Second, inverse modeling in those domains requires differentiability because many relevant quantities are unobservable in situ, including initial velocity fields, internal stress states, and spatially varying friction. Conventional mesh-based methods such as FEM and FVM can fail under severe distortion, while traditional inverse methods often require repeated forward solves, Bayesian sampling, or hand-derived adjoints; JAX-MPM addresses these constraints by making the full MPM pipeline differentiable and GPU-efficient (Du et al., 6 Jul 2025).

The framework differs from conventional MPM solvers written in C++, Fortran, or Julia by being implemented entirely in JAX, using reverse-mode automatic differentiation through the simulation steps, and supporting JIT compilation, vectorization, checkpointing, and GPU acceleration. Compared with PINNs, graph neural simulators, neural operators, Taichi-based differentiable MPM, and adjoint-based PDE solvers, it is presented as a discrete, physics-based solver rather than a surrogate or emulator, and it avoids manual adjoint derivation by differentiating through the unrolled solver directly (Du et al., 6 Jul 2025).

A broader contextual point is that JAX-MPM occupies a distinct position among JAX-based scientific computing systems. JetSCI, for example, is a hybrid JAX-PETSc framework for scalable differentiable finite-element simulation rather than an MPM package; a plausible implication is that the two systems exemplify complementary architectures for differentiable mechanics, with JAX-MPM centered on meshfree particle-grid dynamics and JetSCI centered on finite-element residual assembly and PETSc-based global sparse solves (Cattaneo et al., 23 Apr 2026).

2. Hybrid Lagrangian-Eulerian MPM formulation

JAX-MPM uses the standard hybrid Lagrangian-Eulerian MPM framework. Particles or material points carry mass, volume, velocity, stress, strain, and history variables, while a fixed background grid is used for transient computation at each time step. The balance of linear momentum is written as

$$\rho \ddot{u} = \nabla \cdot \sigma + \rho b, \qquad u = \bar{u} \text{ on } \Gamma_D, \qquad \sigma \cdot n = \bar{t} \text{ on } \Gamma_N.$$

Each MPM step is decomposed into particle-to-grid (P2G), grid-based update, grid-to-particle (G2P), and particle state update. The framework uses a quadratic B-spline basis. In one dimension,

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$

The multidimensional grid shape function is the tensor product

$$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$

The P2G transfer uses the standard particle-weighted interpolation formulas. Grid mass is

$m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$

and nodal velocity is

$v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$

The nodal force contains internal stress forces, body force, and external forces:

$$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$

A lumped mass matrix,

$m_i = \sum_p m_p D_{ip},$

is used to avoid expensive inversion (Du et al., 6 Jul 2025).

The grid update is explicit. Acceleration is

$a_i^n = \frac{f_i^n}{m_i^n},$

and velocity is advanced with explicit Euler:

$v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$

For fixed nodes on $\Gamma_D$, the implementation imposes

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$0

On the G2P side, JAX-MPM supports PIC, FLIP, APIC, and TPIC transfers. The paper explicitly gives FLIP and PIC. FLIP uses the grid velocity increment,

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$1

while PIC replaces the particle velocity by interpolated grid velocity,

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$2

Particle position is updated by

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$3

and the particle velocity gradient is

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$4

PIC is described as more dissipative, whereas FLIP is less dissipative but can be noisier. Stress is updated with an Update Stress Last (USL) strategy after the grid momentum solve and velocity update (Du et al., 6 Jul 2025).

3. Constitutive behavior, stress updates, and contact treatment

The framework is intended for both fluid-like and granular or elastoplastic regimes. For dam-break problems, JAX-MPM uses a weakly compressible Newtonian fluid model with stress update

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$5

together with the weakly compressible equation of state

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$6

The density update is based on deformation or volume change (Du et al., 6 Jul 2025).

For granular collapse, the constitutive law is a non-associated Drucker-Prager model with tension cutoff and a Jaumann objective stress rate,

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$7

The constitutive algorithm includes elastic trial stress evaluation, computation of trial mean stress and deviatoric stress, yield checks, return mapping correction for shear failure, and tensile cutoff correction. The appendix is described as giving the main Drucker-Prager shear and tensile yield functions and the corrected stress update (Du et al., 6 Jul 2025).

Contact and boundary interaction are treated explicitly. In the 3D dam-break benchmark with rigid obstacles, the framework uses a frictionless predictor-corrector contact model. The corrected contact-node velocity removes the normal component of the predicted velocity:

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$8

For inverse friction estimation, a Coulomb friction law is imposed on segmented bottom boundaries, with corrected contact-node velocity

$$w(\xi)= \begin{cases} \frac{3}{4}-\xi^2, & |\xi|<\frac{1}{2},\ \frac{1}{2}\left(\frac{3}{2}-|\xi|\right)^2, & \frac{1}{2}\le |\xi|<\frac{3}{2},\ 0, & \text{otherwise}. \end{cases}$$9

The main text and appendix indicate that this treatment is based on Bardenhagen-style MPM contact handling (Du et al., 6 Jul 2025).

4. Differentiable programming architecture and learning-augmented design

JAX-MPM represents the solver as a differentiable recursive operator,

$$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$0

where $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$1 includes particle state variables such as position $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$2, velocity $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$3, deformation gradient $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$4, density $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$5, and optionally stress and plastic strain. Over a full trajectory, the discretized governing equations are written as

$$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$6

The paper states that JAX autodiff can differentiate through the full time-stepping loop, P2G transfer, grid acceleration and update, boundary and contact corrections, G2P transfer, stress update and constitutive response, and neural parameterizations used for inverse problems. Gradients are obtained with reverse-mode AD directly as

$$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$7

This removes the need for manual adjoint derivation (Du et al., 6 Jul 2025).

The implementation relies on the functional-programming style typical of JAX. The paper emphasizes jax.jit for compilation and kernel fusion, jax.vmap for batch or vector parallelism, jax.scan for time stepping, and jax.checkpoint or jax.remat for memory-efficient reverse-mode AD. To control memory use in long simulations, the implementation also uses segmented simulation, reported to reduce memory complexity from $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$8 to roughly $$D_i(x)=\prod_{a=1}^d w\!\left(\frac{x_a-x_{i,a}}{h}\right).$$9 for $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$0 segments. A representative checkpointed substep is $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$8

A defining feature is its learning-augmented structure. Spatially or temporally varying unknown parameters can be represented by a neural network,

$m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$1

with trainable weights $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$2. This supports inverse recovery of fields rather than only scalar parameters and enables end-to-end optimization through the simulator. The framework is described as supporting scalar physical parameters, fields represented by neural networks, initial conditions, friction coefficients, and potentially constitutive parameters, with seamless integration into JAX ML stacks such as Flax, Haiku, and Equinox (Du et al., 6 Jul 2025).

5. Validation benchmarks, scalability, and runtime characteristics

The validation program spans fluid and granular problems in both two and three dimensions. In the 2D dam-break benchmark, the setup uses a water column with $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$3 m and $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$4 m, grid spacing $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$5 m, 4 particles per cell, sound speed $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$6 m/s, time step $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$7 s, and slip boundaries. JAX-MPM is reported to match shallow-water analytical front propagation reasonably well, and FLIP, TPIC, and APIC all give close agreement, with small early-time deviations due to finite-depth effects and vertical accelerations (Du et al., 6 Jul 2025).

The quasi-2D granular collapse benchmark models an aluminum-bar assemblage in plane strain on a $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$8 grid with $m_i^n = \sum_p m_p D_{ip}, \qquad D_{ip}=D_i(x_p),$9 kg/m$v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$0, $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$1 MPa, $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$2, friction angle $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$3, dilation angle $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$4, cohesion $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$5, FLIP transfer, and a total run time of 0.65 s with $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$6 s. The equivalent plastic strain localizes near the base, showing shear-band formation, and the simulated failure surface and free surface match experiments from Bui et al. (Du et al., 6 Jul 2025).

The 3D dam-break with rigid obstacles uses a domain $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$7 m$v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$8, water column dimensions $v_i^n = \frac{1}{m_i^n}\sum_p m_p D_{ip} v_p^n.$9 m, $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$0 m, $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$1 m, rigid columns with side lengths 0.04 m and 0.06 m, fluid viscosity $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$2 Pa·s, sound speed 35 m/s, about 1.92 million particles, and a $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$3 grid. It captures impact, splashing, recirculation, and separation behind obstacles, and the paper reports 1000 time steps in about 8.2 s in float32 or 35.1 s in float64 on a single GPU (Du et al., 6 Jul 2025).

The 3D granular column collapse benchmark uses 2,408,448 particles, 200,000 steps, $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$4 s, and aspect ratios $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$5. Higher aspect ratio produces more free fall, more mobilized material, larger runout, and a transition from roughly linear to sublinear scaling of normalized runout. The reported normalized runout distances are 2.05, 3.97, 7.76, and 10.65 for $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$6, respectively. In comparison with CB-Geo on CPU, JAX-MPM on GPU is reported to exceed a $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$7 speedup at the highest resolution, around 900,000 particles and roughly 9 s per 1000 steps on the GPU, with benchmarks run on an AMD EPYC 7763 CPU with 128 threads and an NVIDIA A100 GPU (Du et al., 6 Jul 2025).

The most demanding benchmark emphasized in the abstract is the 3D granular cylinder collapse. It uses cylinder diameter $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$8 m, height $$f_i^n = -\sum_p V_p^n \, \sigma_p^n \nabla D_{ip} + \sum_p m_p D_{ip} b_p^n + f_i^{\text{ext}}.$$9 m, friction angles $m_i = \sum_p m_p D_{ip},$0, total simulation time 1 s, $m_i = \sum_p m_p D_{ip},$1 s, 2,712,960 particles, and a $m_i = \sum_p m_p D_{ip},$2 grid with roughly 16 million grid nodes on a single NVIDIA A100 GPU. Higher friction yields a steeper final deposit, less spreading, and a more confined pile morphology. Runtime is reported as about 22.3 s per 1000 time steps in float32 and 98.4 s in float64; the abstract summarizes this as a high-resolution 3D granular cylinder collapse with 2.7 million particles completing 1000 steps in approximately 22 seconds in single precision and 98 seconds in double precision (Du et al., 6 Jul 2025).

Benchmark	Setup highlights	Main result
2D dam-break	$m_i = \sum_p m_p D_{ip},$3 m, $m_i = \sum_p m_p D_{ip},$4 m, $m_i = \sum_p m_p D_{ip},$5 m, 4 particles per cell	Matches shallow-water analytical front propagation reasonably well
Quasi-2D granular collapse	$m_i = \sum_p m_p D_{ip},$6 grid, Drucker-Prager material, FLIP	Shear-band localization near the base; failure and free surface match experiments
3D dam-break with obstacles	1.92 million particles, $m_i = \sum_p m_p D_{ip},$7 grid	1000 steps in about 8.2 s (float32) or 35.1 s (float64)
3D granular column collapse	2,408,448 particles, 200,000 steps	GPU JAX-MPM exceeds $m_i = \sum_p m_p D_{ip},$8 speedup over CB-Geo CPU at highest resolution
3D granular cylinder collapse	2,712,960 particles, $m_i = \sum_p m_p D_{ip},$9 grid	1000 steps in about 22.3 s (float32) or 98.4 s (float64)

6. Inverse modeling, observation operators, applications, and limitations

The inverse problem is formulated as

$a_i^n = \frac{f_i^n}{m_i^n},$0

This supports both Lagrangian and Eulerian supervision. For tracked particles,

$a_i^n = \frac{f_i^n}{m_i^n},$1

whereas for fixed spatial regions,

$a_i^n = \frac{f_i^n}{m_i^n},$2

An important design point is that the same differentiable solver can use particle data or field or region data without changing the simulator (Du et al., 6 Jul 2025).

Three inverse demonstrations structure the paper. The first is recovery of a scalar initial velocity parameter in a 2D dam-break, with ground truth

$a_i^n = \frac{f_i^n}{m_i^n},$3

and initial guess $a_i^n = \frac{f_i^n}{m_i^n},$4. Three supervision modes are used: full particle positions with $a_i^n = \frac{f_i^n}{m_i^n},$5, sparse particle positions with $a_i^n = \frac{f_i^n}{m_i^n},$6, and sparse Eulerian velocity monitors with $a_i^n = \frac{f_i^n}{m_i^n},$7. The loss is the trajectory mismatch over the final 100 frames,

$a_i^n = \frac{f_i^n}{m_i^n},$8

All cases recover $a_i^n = \frac{f_i^n}{m_i^n},$9, while the velocity-monitor supervision converges fastest and most stably. The paper highlights that the best observation type is not always the densest one; alignment of observations with the parameter’s influence matters (Du et al., 6 Jul 2025).

The second inverse example reconstructs a spatially varying initial velocity field with a parabolic base and sine perturbation using a feedforward neural network with 3 hidden layers and 30 ReLU units per layer. Sparse Lagrangian particle trajectories are provided at $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$0, 500, and 1000. The reported relative $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$1 errors are 45.1%, 6.4%, and 3.7%, respectively. The interpretation given is that 100 particles are too sparse to identify the field well, whereas 500 and 1000 provide sufficient information; error remains larger near the bottom because of weaker identifiability and contact effects (Du et al., 6 Jul 2025).

The third inverse example estimates segmented bottom-wall Coulomb friction coefficients in a dam-break problem on a domain with $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$2 m and $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$3 m, with initial water column $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$4 m and $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$5 m. The bottom is divided into four segments with ground-truth

$v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$6

Using particle trajectories over 1000 frames from a 10,000-step simulation and sparse settings of $v_i^{n+1} = v_i^n + \Delta t \, a_i^n.$7 and 500, both cases recover the first three upstream segments well. The downstream segment is harder because it interacts later. The estimated last segment is 0.229 for 100 particles, about 14.5% error, and 0.198 for 500 particles, about 1.0% error. This is used to show how observation density and temporal influence region affect inverse recovery quality (Du et al., 6 Jul 2025).

Optimization in these inverse problems uses Adam with a decaying learning rate and an initial learning rate of 0.1. The framework is explicitly targeted at landslides, debris flows, granular collapse and runout, dam-break flooding, soil liquefaction, geomechanical back-analysis, data assimilation from sparse monitoring, learning digital twins for geohazards, and inverse estimation of friction and initial conditions. At the same time, the paper notes several limitations: memory remains demanding for very long or high-resolution differentiable simulations even with checkpointing; nonsmooth constitutive laws such as Drucker-Prager can lead to instability or problematic gradients; implicit time integration could improve stability for longer horizons; and the current neural parameterizations are simple feedforward networks, with more expressive or hybrid schemes identified as future directions (Du et al., 6 Jul 2025).