release

Author: wgsxm

.gitignore

@@ -0,0 +1,5 @@
+*.egg-info
+__pycache__
+*.ipynb
+output
+build

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Yuchen Lin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,87 @@
+# A PyTorch Implementation of MLS-MPM (Moving Least Squares Material Point Method)
+This repository provides a PyTorch implementation of the MLS-MPM (Moving Least Squares Material Point Method). The algorithm is implemented using **a few lines of tensor operations in PyTorch**, making it naturally differentiable and optimized for GPU acceleration. 
+The code is vectorized without any explicit loops, which makes it efficient for large-scale simulations. 
+
+[Gradient Checkpointing](https://pytorch.org/docs/stable/checkpoint.html) is highly recommended when integrating the MLS-MPM into a trainable deep learning framework. See [OmniPhysGS](https://github.com/wgsxm/omniphysgs) (ICLR 2025) for an example. 
+## Installation
+### From source
+```bash
+git clone https://github.com/wgsxm/MPM-PyTorch.git
+cd MPM-PyTorch
+pip install .
+```
+## Quick Start
+Run the following code to try a simple example of MLS-MPM simulation. The code simulates a 3D elastic jelly falling onto a rigid floor. By default, the code will produce a video of the simulation in the `output` directory.
+```bash
+python simulate.py --config examples/jelly.yaml
+```
+<img src="assets/jelly.gif" width=400> 
+<img src="assets/sand.gif" width=400> 
+
+## Usage
+Refer to `simulate.py` for more details. 
+```python
+from mpm_pytorch import MPMSolver, set_boundary_conditions
+from mpm_pytorch.constitutive_models import *
+particles = ... # Particle positions to simulate
+# Create a MPM solver with default parameters
+mpm_solver = MPMSolver(particles)
+# Set boundary conditions (optional)
+boundary_conditions = ... # Refer to example configs
+set_boundary_conditions(mpm_solver, boundary_conditions)
+# Create constitutive models
+elasticity = CorotatedElasticity(E=2e6)
+plasicity = IdentityPlasticity()
+# Init particle state
+x = particles
+v = torch.zeros_like(x)
+C = torch.zeros((x.shape[0], 3, 3), device=x.device)
+F = torch.eye(3, device=x.device).unsqueeze(0).repeat(x.shape[0], 1, 1)
+# Start simulation for T steps
+for i in range(T):
+    # Update stress
+    stress = elasticity(F)
+    # Particle to grid, grid update, grid to particle
+    grid_update = mpm_solver(x, v, C, F, stress)
+    # Plasticity correction
+    F = plasticity(F)
+```
+
+## Fast Batched SVD
+Batched SVD is a common operation in constitutive models. Original PyTorch SVD is not optimized for batched computation. We adopt a `warp-lang` implementation of differentiable batched SVD in `mpm_pytorch.constitutive_models.warp_svd` from [NCLaw](https://github.com/PingchuanMa/NCLaw/tree/main/nclaw/warp). It can run on CPU or GPU with CUDA. 
+
+The result of the decomposed matrices is not guaranteed to be the same as the original PyTorch SVD. You can also use the original PyTorch SVD or other batched SVD implementations if there is any environment conflict (the version of `warp-lang` requires `numpy<2`). 
+
+We provide a script to benchmark the batched SVD implementation. 
+```bash
+python benchmark_svd.py
+```
+If you encounter error as `ModuleNotFoundError: No module named 'imp'` when using higher version of Python, simply change the import statement from `import imp` to `import importlib as imp`. 
+
+## Citation
+If you find our work helpful, please consider citing:
+```
+@inproceedings{
+  lin2025omniphysgs,
+  title={OmniPhys{GS}: 3D Constitutive Gaussians for General Physics-Based Dynamics Generation},
+  author={Yuchen Lin and Chenguo Lin and Jianjin Xu and Yadong MU},
+  booktitle={The Thirteenth International Conference on Learning Representations},
+  year={2025},
+}
+```
+```
+@article{hu2018moving,
+  title={A moving least squares material point method with displacement discontinuity and two-way rigid body coupling},
+  author={Hu, Yuanming and Fang, Yu and Ge, Ziheng and Qu, Ziyin and Zhu, Yixin and Pradhana, Andre and Jiang, Chenfanfu},
+  journal={ACM Transactions on Graphics (TOG)},
+  year={2018},
+}
+```
+```
+@article{stomakhin2013material,
+  title={A material point method for snow simulation},
+  author={Stomakhin, Alexey and Schroeder, Craig and Chai, Lawrence and Teran, Joseph and Selle, Andrew},
+  journal={ACM Transactions on Graphics (TOG)},
+  year={2013},
+}
+```

assets/jelly.gif

@@ -0,0 +1,111 @@
+import time
+import torch
+from tqdm import tqdm
+
+from mpm_pytorch.constitutive_models.warp_svd import SVD
+
+wp_svd = SVD()
+wp_svd(torch.randn(1, 3, 3))
+def warp_svd(x):
+    return wp_svd(x)
+    
+def torch_svd(x):
+    U, s, Vh = torch.svd(x)
+    return U, s, Vh.transpose(-2, -1)
+
+if __name__ == '__main__':
+    test_time = 10
+    n = 100000
+    atol = 1e-5
+
+    print("\nTest correctness")
+    Is = torch.eye(3).unsqueeze(0).repeat(n, 1, 1)
+    for _ in tqdm(range(test_time), desc="Test orthogonality of U and V"):
+        x = torch.randn(n, 3, 3)
+        U, s, V = warp_svd(x)
+        UU = U @ U.transpose(-2, -1)
+        VV = V @ V.transpose(-2, -1)
+        assert torch.allclose(UU, Is, atol=1e-5)
+        assert torch.allclose(VV, Is, atol=1e-5)
+    bar = tqdm(range(test_time), desc="Test correctness of decomposition")
+    for _ in bar:
+        x = torch.randn(n, 3, 3)
+        U, s, V = warp_svd(x)
+        warp_pred = U @ torch.diag_embed(s) @ V
+        warp_mae = torch.abs(warp_pred - x).mean()
+        U, s, V = torch_svd(x)
+        torch_pred = U @ torch.diag_embed(s) @ V
+        torch_mae = torch.abs(torch_pred - x).mean()
+        bar.set_postfix(warp_mae=warp_mae.item(), torch_mae=torch_mae.item())
+        assert warp_mae < atol
+    
+    print("\nTest differentiability")
+    bar = tqdm(range(test_time), desc="Test differentiability")
+    for _ in bar:
+        x = torch.randn(n, 3, 3, requires_grad=True)
+        U, s, V = warp_svd(x)
+        warp_pred = U @ torch.diag_embed(s) @ V
+        warp_pred.sum().backward()
+        assert x.grad is not None
+        warp_mae = torch.abs(x.grad - torch.ones_like(x)).mean()
+        x.grad = None
+        U, s, V = torch_svd(x)
+        torch_pred = U @ torch.diag_embed(s) @ V
+        torch_pred.sum().backward()
+        torch_mae = torch.abs(x.grad - torch.ones_like(x)).mean()
+        bar.set_postfix(warp_mae=warp_mae.item(), torch_mae=torch_mae.item())
+        assert warp_mae < atol
+    
+    print("\nTest speed")
+    x = torch.randn(n, 3, 3)
+    start = time.time()
+    for _ in tqdm(range(test_time), desc="Test warp_svd forward (CPU)"):
+        U, s, V = warp_svd(x)
+    end = time.time()
+    print("warp_svd: ", (end - start) / test_time, "seconds")
+    start = time.time()
+    for _ in tqdm(range(test_time), desc="Test torch_svd forward (CPU)"):
+        U, s, V = torch_svd(x)
+    end = time.time()
+    print("torch_svd: ", (end - start) / test_time, "seconds")
+
+    print()
+    x = torch.randn(n, 3, 3, requires_grad=True)
+    start = time.time()
+    for _ in tqdm(range(test_time), desc="Test warp_svd forward + backward (CPU)"):
+        U, s, V = warp_svd(x)
+        warp_pred = U @ torch.diag_embed(s) @ V
+        warp_pred.sum().backward()
+    end = time.time()
+    print("warp_svd: ", (end - start) / test_time, "seconds")
+    x.grad = None
+    start = time.time()
+    for _ in tqdm(range(test_time), desc="Test torch_svd forward + backward (CPU)"):
+        U, s, V = torch_svd(x)
+        torch_pred = U @ torch.diag_embed(s) @ V
+        torch_pred.sum().backward()
+    end = time.time()
+    print("torch_svd: ", (end - start) / test_time, "seconds")
+
+    if torch.cuda.is_available():
+        print()
+        x = torch.randn(n, 3, 3).cuda().requires_grad_()
+        start = time.time()
+        for _ in tqdm(range(test_time), desc="Test warp_svd forward + backward (GPU)"):
+            U, s, V = warp_svd(x)
+            warp_pred = U @ torch.diag_embed(s) @ V
+            warp_pred.sum().backward()
+        end = time.time()
+        print("warp_svd: ", (end - start) / test_time, "seconds")
+        x.grad = None
+        start = time.time()
+        for _ in tqdm(range(test_time), desc="Test torch_svd forward + backward (GPU)"):
+            U, s, V = torch_svd(x)
+            torch_pred = U @ torch.diag_embed(s) @ V
+            torch_pred.sum().backward()
+        end = time.time()
+        print("torch_svd: ", (end - start) / test_time, "seconds")
+
+    
+    
+            

assets/sand.gif

@@ -0,0 +1,20 @@
+output_dir: './output'
+tag: 'jelly'
+
+material:
+  elasticity: 'CorotatedElasticity'
+  plasticity: 'IdentityPlasticity'
+  color: 'blue'
+
+sim:
+  num_frames: 150
+  steps_per_frame: 10
+  initial_velocity: [0.0, 0.0, -0.5]
+  boundary_conditions: 
+    - type: 'surface_collider'
+      point: [1.0, 1.0, 0.02]
+      normal: [0.0, 0.0, 1.0]
+      surface: 'sticky'
+      friction: 0.0
+      start_time: 0.0
+      end_time: 1e3

benchmark_svd.py

@@ -0,0 +1,20 @@
+output_dir: './output'
+tag: 'sand'
+
+material:
+  elasticity: 'StVKElasticity'
+  plasticity: 'DruckerPragerPlasticity'
+  color: 'orange'
+
+sim:
+  num_frames: 150
+  steps_per_frame: 10
+  initial_velocity: [0.0, 0.0, -0.5]
+  boundary_conditions: 
+    - type: 'surface_collider'
+      point: [1.0, 1.0, 0.02]
+      normal: [0.0, 0.0, 1.0]
+      surface: 'sticky'
+      friction: 0.0
+      start_time: 0.0
+      end_time: 1e3

examples/jelly.yaml

@@ -0,0 +1,3 @@
+from .mpm_solver import MPMSolver
+from .boundary_condition import set_boundary_conditions
+from .constitutive_models import get_constitutive