Merge remote-tracking branch 'origin/main' into upstream-fairchem

Author: CompRhys

.github/workflows/test.yml

@@ -69,6 +69,7 @@ jobs:
           - { name: mattersim, test_path: "tests/models/test_mattersim.py" }
           - { name: metatomic, test_path: "tests/models/test_metatomic.py" }
           - { name: nequip, test_path: "tests/models/test_nequip_framework.py" }
+          - { name: nequix, test_path: "tests/models/test_nequix.py" }
           - { name: orb, test_path: "tests/models/test_orb.py" }
           - { name: sevenn, test_path: "tests/models/test_sevennet.py" }
         exclude:
@@ -154,7 +155,7 @@ jobs:
       - name: Find example scripts
         id: set-matrix
         run: |
-          EXAMPLES=$(find examples -name "*.py" | jq -R -s -c 'split("\n")[:-1]')
+          EXAMPLES=$(find examples -name "*.py" -not -path "examples/benchmarking/*" | jq -R -s -c 'split("\n")[:-1]')
           echo "examples=$EXAMPLES" >> $GITHUB_OUTPUT
 
   test-examples:

.gitignore

@@ -28,6 +28,9 @@ docs/reference/torch_sim.*
 *.hdf5
 *.traj
 
+# ignore torch.save outputs
+*.pt
+
 # coverage
 coverage.xml
 .coverage*

README.md

@@ -14,7 +14,7 @@ era. By rewriting the core primitives of atomistic simulation in Pytorch, it all
 orders of magnitude acceleration of popular machine learning potentials.
 
 * Automatic batching and GPU memory management allowing significant simulation speedup
-* Support for MACE, Fairchem, SevenNet, ORB, MatterSim, graph-pes, and metatomic MLIP models
+* Support for MACE, Fairchem, SevenNet, ORB, MatterSim, graph-pes, metatomic, and Nequix MLIP models
 * Support for classical lennard jones, morse, and soft-sphere potentials
 * Molecular dynamics integration schemes like NVE, NVT Langevin, and NPT Langevin
 * Relaxation of atomic positions and cell with gradient descent and FIRE

docs/user/reproducibility.md

@@ -79,17 +79,34 @@ Because TorchSim runs batched simulations, all systems in a batch share a single
 
 If strict reproducibility is required, keep your batching setup fixed.
 
-### Serialising the RNG state
+### Serialising state for reproducible restarts
 
-If you wish to be able to resume a session and ensure determinism you need to persist and reload the `torch.Generator` state. This can be done using `torch.save()` and `torch.Generator().set_state()`:
+To resume a simulation and ensure determinism you need to persist and reload the complete state, including the `torch.Generator` RNG state. The simplest approach is to save the full state dict with `torch.save()`:
 
 ```python
+from dataclasses import asdict
+from torch_sim.integrators import MDState
+
 # save
+torch.save(asdict(state), "checkpoint.pt")
+
+# restore (weights_only=False needed for torch.Generator in PyTorch 2.6+)
+restored = MDState(**torch.load("checkpoint.pt", weights_only=False))
+```
+
+This captures positions, momenta, forces, energy, cell, and the `torch.Generator` in a single file. Since `torch.save()` uses pickle, the generator is serialised automatically.
+
+> **Pickle caveat**: The `torch.Generator` object in the dict requires `weights_only=False` and may not unpickle across PyTorch versions. For portable checkpoints, save the tensors normally and extract the RNG state as a plain `uint8` tensor via `get_state()` — this loads with `weights_only=True` and is version-safe:
+
+```python
+# save RNG state as a plain uint8 tensor (no pickle needed)
 rng_state = state.rng.get_state()
 torch.save(rng_state, "rng_state.pt")
 
 # restore
 gen = torch.Generator(device=state.device)
-gen.set_state(torch.load("rng_state.pt"))
+gen.set_state(torch.load("rng_state.pt", weights_only=True))
 state.rng = gen
 ```
+
+See the [reproducible restart tutorial](../../examples/tutorials/reproducible_restart_tutorial.py) for a complete worked example.