TorchSim · orionarcher · Apr 8, 2025 · Apr 5, 2025 · Apr 5, 2025 · Apr 5, 2025
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -60,11 +60,12 @@ jobs:
           - { python: "3.11", resolution: highest }
           - { python: "3.12", resolution: lowest-direct }
         model:
+          - { name: fairchem, test_path: "tests/models/test_fairchem.py" }
           - { name: mace, test_path: "tests/models/test_mace.py" }
           - { name: mace, test_path: "tests/test_elastic.py" }
+          - { name: mattersim, test_path: "tests/models/test_mattersim.py" }
           - { name: orb, test_path: "tests/models/test_orb.py" }
           - { name: sevenn, test_path: "tests/models/test_sevennet.py" }
-          - { name: fairchem, test_path: "tests/models/test_fairchem.py" }
     runs-on: ${{ matrix.os }}
 
     steps:

diff --git a/.gitignore b/.gitignore
@@ -24,6 +24,7 @@ docs/reference/torch_sim.*
 
 # ignore trajectory files
 *.h5
+*.h5md
 *.hdf5
 *.traj
 

diff --git a/README.md b/README.md
@@ -37,7 +37,7 @@ import torch
 import torch_sim as ts
 
 # run natively on gpus
-device = torch.device("cuda")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # easily load the model from mace-mp
 from mace.calculators.foundations_models import mace_mp
@@ -48,7 +48,7 @@ mace_model = MaceModel(model=mace, device=device)
 from ase.build import bulk
 cu_atoms = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2))
 many_cu_atoms = [cu_atoms] * 50
-trajectory_files = [f"Cu_traj_{i}" for i in range(len(many_cu_atoms))]
+trajectory_files = [f"Cu_traj_{i}.h5md" for i in range(len(many_cu_atoms))]
 
 # run them all simultaneously with batching
 final_state = ts.integrate(

diff --git a/docs/conf.py b/docs/conf.py
@@ -62,7 +62,7 @@
     "html_image",
 ]
 
-autodoc_mock_imports = ["mace", "fairchem", "orb", "sevennet"]
+autodoc_mock_imports = ["fairchem", "mace", "mattersim", "orb", "sevennet"]
 
 # use type hints
 autodoc_typehints = "description"

diff --git a/examples/scripts/1_Introduction/1.2_MACE.py b/examples/scripts/1_Introduction/1.2_MACE.py
@@ -38,7 +38,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/1_Introduction/1.3_Batched_MACE.py b/examples/scripts/1_Introduction/1.3_Batched_MACE.py
@@ -41,7 +41,6 @@
     # Or load from compiled model
     # model=compiled_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/2_Structural_optimization/2.3_MACE_FIRE.py b/examples/scripts/2_Structural_optimization/2.3_MACE_FIRE.py
@@ -13,7 +13,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_optimizers import fire
@@ -54,7 +53,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py b/examples/scripts/2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py
@@ -14,7 +14,6 @@
 from mace.calculators.foundations_models import mace_mp
 
 from torch_sim.io import atoms_to_state
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_optimizers import unit_cell_fire
 from torch_sim.units import UnitConversion
@@ -52,7 +51,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py b/examples/scripts/2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py
@@ -14,7 +14,6 @@
 from mace.calculators.foundations_models import mace_mp
 
 from torch_sim.io import atoms_to_state
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_optimizers import frechet_cell_fire
 from torch_sim.units import UnitConversion
@@ -52,7 +51,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py b/examples/scripts/2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py
@@ -15,7 +15,6 @@
 
 from torch_sim.io import atoms_to_state
 from torch_sim.models.mace import MaceModel
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.optimizers import gradient_descent
 
 
@@ -56,7 +55,6 @@
 batched_model = MaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/...les/scripts/2_Structural_optimization/2.7_Batched_MACE_UnitCellFilter_Gradient_Descent.py b/...les/scripts/2_Structural_optimization/2.7_Batched_MACE_UnitCellFilter_Gradient_Descent.py
@@ -66,7 +66,6 @@
 model = ts.models.MaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_FIRE.py b/examples/scripts/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_FIRE.py
@@ -66,7 +66,6 @@
 model = ts.models.MaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/2_Structural_optimization/2.9_Batched_MACE_FrechetCellFilter_FIRE.py b/examples/scripts/2_Structural_optimization/2.9_Batched_MACE_FrechetCellFilter_FIRE.py
@@ -65,7 +65,6 @@
 model = ts.models.MaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py b/examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py
@@ -12,7 +12,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kinetic_energy, calc_kT
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -58,7 +57,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py b/examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py
@@ -14,7 +14,6 @@
 from mace.calculators.foundations_models import mace_off
 
 from torch_sim.io import atoms_to_state
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kinetic_energy
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_integrators import nve
@@ -47,7 +46,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.2_MACE_NVE.py b/examples/scripts/3_Dynamics/3.2_MACE_NVE.py
@@ -13,7 +13,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kinetic_energy
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -54,7 +53,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py b/examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py
@@ -13,7 +13,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kinetic_energy
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -54,7 +53,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py b/examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py
@@ -12,7 +12,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kT
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -53,7 +52,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py b/examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py
@@ -12,7 +12,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kT
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -56,7 +55,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py b/examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py
@@ -18,7 +18,6 @@
 from mace.calculators.foundations_models import mace_mp
 from plotly.subplots import make_subplots
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kT
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -137,7 +136,6 @@ def get_kT(
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py b/examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py
@@ -12,7 +12,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kinetic_energy, calc_kT
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -59,7 +58,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,

diff --git a/examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py b/examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py
@@ -12,7 +12,6 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.neighbors import vesin_nl_ts
 from torch_sim.quantities import calc_kT
 from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
@@ -56,7 +55,6 @@
 model = UnbatchedMaceModel(
     model=loaded_model,
     device=device,
-    neighbor_list_fn=vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,

diff --git a/examples/scripts/5_Workflow/5.1_a2c_silicon.py b/examples/scripts/5_Workflow/5.1_a2c_silicon.py
@@ -74,7 +74,6 @@
 model = UnbatchedMaceModel(
     model=raw_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,  # We don't need stress for MD
     dtype=dtype,

diff --git a/examples/scripts/5_Workflow/5.2_a2c_silicon_batched.py b/examples/scripts/5_Workflow/5.2_a2c_silicon_batched.py
@@ -74,7 +74,6 @@
 model = UnbatchedMaceModel(
     model=raw_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=False,  # We don't need stress for MD
     dtype=dtype,
@@ -205,7 +204,6 @@ def step_fn(
 model = ts.models.MaceModel(
     model=raw_model,
     device=device,
-    neighbor_list_fn=ts.neighbors.vesin_nl_ts,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,
-Original file line number
+Diff line change
@@ Expand Up / @@ -24,6 +24,7 @@ docs/reference/torch_sim.* @@
     # ignore trajectory files
     *.h5
+    *.h5md
     *.hdf5
     *.traj
@@ Expand Down @@