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

• URL: http://arxiv.org/abs/2507.04192v1
• Date: Sun, 06 Jul 2025 00:00:11 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 within a hybrid Lagrangian-Eulerian framework.<n>Built on the material point method (MPM) and implemented using the JAX computing framework, JAX-MPM is fully differentiable and GPU-accelerated.<n>We validate JAX-MPM on several 2D and 3D benchmarks, including dam-breaks and granular collapses, demonstrating its accuracy and performance.
• Score: 1.450261153230204
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Differentiable programming that enables automatic differentiation through simulation pipelines has emerged as a powerful paradigm in scientific computing, supporting both forward and inverse modeling and facilitating integration with deep learning frameworks. We present JAX-MPM, a general-purpose differentiable meshfree solver within a hybrid Lagrangian-Eulerian framework, tailored for simulating complex continuum mechanics involving large deformations, frictional contact, and inelastic material behavior, with emphasis on geomechanics and geophysical hazard applications. Built on the material point method (MPM) and implemented using the JAX computing framework, JAX-MPM is fully differentiable and GPU-accelerated, enabling efficient gradient-based optimization directly through time-stepping solvers. It supports joint training of physical models and neural networks, allowing the learning of embedded closures and neural constitutive models. We validate JAX-MPM on several 2D and 3D benchmarks, including dam-breaks and granular collapses, demonstrating its accuracy and performance. A high-resolution 3D granular cylinder collapse with 2.7 million particles completes 1000 steps in ~22 seconds (single precision) and ~98 seconds (double precision) on a single GPU. Beyond forward modeling, we demonstrate inverse modeling capabilities such as velocity field reconstruction and spatially varying friction estimation. These results establish JAX-MPM as a unified, scalable platform for differentiable meshfree simulation and data-driven geomechanical inference.

Related papers

• 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.