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

• URL: http://arxiv.org/abs/2507.04192v2
• Date: Sat, 27 Sep 2025 19:33:27 GMT
• Title: JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modeling
• Authors: Honghui Du, QiZhi He,
• Abstract summary: We present JAX-MPM, a differentiable meshfree solver based on the material point method (MPM)<n>The solver adopts a hybrid Eulerian-Lagrangian framework to capture large deformations, frictional contact, and in material behavior.<n>JAX-MPM enables efficient gradient-based optimization directly through its time-stepping solvers.
• Score: 1.4287758028119788
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 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.

Related papers

• FastPhysGS: Accelerating Physics-based Dynamic 3DGS Simulation via Interior Completion and Adaptive Optimization [56.17833729527066]
  We propose FastPhysGS, a framework for physics-based dynamic 3DGS simulation.<n>FastPhysGS achieves high-fidelity physical simulation in 1 minute using only 7 GB runtime memory.
  arXiv  Detail & Related papers  (2026-02-02T07:00:42Z)
• Adaptive Mesh-Quantization for Neural PDE Solvers [51.26961483962011]
  Graph Neural Networks can handle the irregular meshes required for complex geometries and boundary conditions, but still apply uniform computational effort across all nodes.<n>We propose Adaptive Mesh Quantization: spatially adaptive quantization across mesh node, edge, and cluster features, dynamically adjusting the bit-width used by a quantized model.<n>We demonstrate our framework's effectiveness by integrating it with two state-of-the-art models, MP-PDE and GraphViT, to evaluate performance across multiple tasks.
  arXiv  Detail & Related papers  (2025-11-23T14:47:24Z)
• Geometric Operator Learning with Optimal Transport [77.16909146519227]
  We propose integrating optimal transport (OT) into operator learning for partial differential equations (PDEs) on complex geometries.<n>For 3D simulations focused on surfaces, our OT-based neural operator embeds the surface geometry into a 2D parameterized latent space.<n> Experiments with Reynolds-averaged Navier-Stokes equations (RANS) on the ShapeNet-Car and DrivAerNet-Car datasets show that our method achieves better accuracy and also reduces computational expenses.
  arXiv  Detail & Related papers  (2025-07-26T21:28:25Z)
• Fast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations [12.00988094580341]
  We present an end-to-end differentiable molecular simulation framework (DIMOS) for molecular dynamics and Monte Carlo simulations.<n>Thanks to its modularity, both classical and machine-learning-based approaches can be easily combined into a hybrid description of the system (ML/MM)<n>The superior performance and the high versatility is probed in different benchmarks and applications, with speed-up factors of up to $170times$.
  arXiv  Detail & Related papers  (2025-03-26T13:39:10Z)
• GausSim: Foreseeing Reality by Gaussian Simulator for Elastic Objects [55.02281855589641]
  GausSim is a novel neural network-based simulator designed to capture the dynamic behaviors of real-world elastic objects represented through Gaussian kernels.<n>We leverage continuum mechanics and treat each kernel as a Center of Mass System (CMS) that represents continuous piece of matter.<n>In addition, GausSim incorporates explicit physics constraints, such as mass and momentum conservation, ensuring interpretable results and robust, physically plausible simulations.
  arXiv  Detail & Related papers  (2024-12-23T18:58:17Z)
• Iterative Sizing Field Prediction for Adaptive Mesh Generation From Expert Demonstrations [49.173541207550485]
  Adaptive Meshing By Expert Reconstruction (AMBER) is an imitation learning problem. AMBER combines a graph neural network with an online data acquisition scheme to predict the projected sizing field of an expert mesh. We experimentally validate AMBER on 2D meshes and 3D meshes provided by a human expert, closely matching the provided demonstrations and outperforming a single-step CNN baseline.
  arXiv  Detail & Related papers  (2024-06-20T10:01:22Z)
• PhyRecon: Physically Plausible Neural Scene Reconstruction [81.73129450090684]
  We introduce PHYRECON, the first approach to leverage both differentiable rendering and differentiable physics simulation to learn implicit surface representations. Central to this design is an efficient transformation between SDF-based implicit representations and explicit surface points. Our results also exhibit superior physical stability in physical simulators, with at least a 40% improvement across all datasets.
  arXiv  Detail & Related papers  (2024-04-25T15:06:58Z)
• A Multi-Grained Symmetric Differential Equation Model for Learning Protein-Ligand Binding Dynamics [73.35846234413611]
  In drug discovery, molecular dynamics (MD) simulation provides a powerful tool for predicting binding affinities, estimating transport properties, and exploring pocket sites. We propose NeuralMD, the first machine learning (ML) surrogate that can facilitate numerical MD and provide accurate simulations in protein-ligand binding dynamics. We demonstrate the efficiency and effectiveness of NeuralMD, achieving over 1K$times$ speedup compared to standard numerical MD simulations.
  arXiv  Detail & Related papers  (2024-01-26T09:35:17Z)
• SoftMAC: Differentiable Soft Body Simulation with Forecast-based Contact Model and Two-way Coupling with Articulated Rigid Bodies and Clothes [7.93869411097782]
  We present SoftMAC, a differentiable simulation framework that couples soft bodies with articulated rigid bodies and clothes.<n>To couple MPM particles with deformable and non-volumetric clothes meshes, we also propose a penetration tracing algorithm.<n>We conducted comprehensive experiments to validate the effectiveness and accuracy of the proposed differentiable pipeline.
  arXiv  Detail & Related papers  (2023-12-06T05:36:55Z)
• LagrangeBench: A Lagrangian Fluid Mechanics Benchmarking Suite [0.0]
  We present LagrangeBench, the first benchmarking suite for Lagrangian particle problems. Our contribution is: (a) seven new fluid mechanics datasets (four in 2D and three in 3D) generated with the Smoothed Particle Hydrodynamics (SPH) method. We also go beyond established position errors and introduce physical metrics like kinetic energy MSE and Sinkhorn distance for the particle distribution.
  arXiv  Detail & Related papers  (2023-09-28T11:03:23Z)
• Learning Controllable Adaptive Simulation for Multi-resolution Physics [86.8993558124143]
  We introduce Learning controllable Adaptive simulation for Multi-resolution Physics (LAMP) as the first full deep learning-based surrogate model. LAMP consists of a Graph Neural Network (GNN) for learning the forward evolution, and a GNN-based actor-critic for learning the policy of spatial refinement and coarsening. We demonstrate that our LAMP outperforms state-of-the-art deep learning surrogate models, and can adaptively trade-off computation to improve long-term prediction error.
  arXiv  Detail & Related papers  (2023-05-01T23:20:27Z)
• DeepPhysics: a physics aware deep learning framework for real-time simulation [0.0]
  We propose a solution to simulate hyper-elastic materials using a data-driven approach. A neural network is trained to learn the non-linear relationship between boundary conditions and the resulting displacement field. The results show that our network architecture trained with a limited amount of data can predict the displacement field in less than a millisecond.
  arXiv  Detail & Related papers  (2021-09-17T12:15:47Z)
• Learning Mesh-Based Simulation with Graph Networks [20.29893312074383]
  We introduce MeshGraphNets, a framework for learning mesh-based simulations using graph neural networks. Our results show it can accurately predict the dynamics of a wide range of physical systems, including aerodynamics, structural mechanics, and cloth.
  arXiv  Detail & Related papers  (2020-10-07T13:34:49Z)
• Learning to Simulate Complex Physics with Graph Networks [68.43901833812448]
  We present a machine learning framework and model implementation that can learn to simulate a wide variety of challenging physical domains. Our framework---which we term "Graph Network-based Simulators" (GNS)--represents the state of a physical system with particles, expressed as nodes in a graph, and computes dynamics via learned message-passing. Our results show that our model can generalize from single-timestep predictions with thousands of particles during training, to different initial conditions, thousands of timesteps, and at least an order of magnitude more particles at test time.
  arXiv  Detail & Related papers  (2020-02-21T16:44:28Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.