Top-level SimState

examples/scripts/1_Introduction/1.2_MACE.py

@@ -11,6 +11,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 
 
@@ -19,9 +20,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/1_Introduction/1.3_Batched_MACE.py

@@ -11,17 +11,16 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 
 
 # Set device and data type
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.10_Batched_MACE_FrechetCellFilter_FIRE.py

@@ -14,7 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 from torch_sim.units import UnitConversion
 
 
@@ -23,9 +23,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py

@@ -11,7 +11,7 @@
 
 import torch
 
-from torch_sim.state import SimState
+import torch_sim as ts
 from torch_sim.unbatched.models.soft_sphere import UnbatchedSoftSphereModel
 from torch_sim.unbatched.unbatched_optimizers import fire
 
@@ -75,7 +75,7 @@
 # Cu atomic mass in atomic mass units
 masses = torch.full((positions.shape[0],), 63.546, device=device, dtype=dtype)
 
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,

examples/scripts/2_Structural_optimization/2.3_MACE_FIRE.py

@@ -13,7 +13,8 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-from torch_sim.state import SimState
+import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_optimizers import fire
 
@@ -23,9 +24,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,
@@ -58,7 +58,7 @@
     dtype=dtype,
     enable_cueq=False,
 )
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,

examples/scripts/2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py

@@ -14,6 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_optimizers import unit_cell_fire
 from torch_sim.units import UnitConversion
@@ -24,9 +25,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py

@@ -14,6 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_optimizers import frechet_cell_fire
 from torch_sim.units import UnitConversion
@@ -24,9 +25,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py

@@ -14,7 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 from torch_sim.optimizers import gradient_descent
 
 
@@ -23,9 +23,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.7_Batched_MACE_FIRE.py

@@ -14,7 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 from torch_sim.optimizers import fire
 
 
@@ -23,9 +23,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py

@@ -14,7 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 from torch_sim.optimizers import unit_cell_gradient_descent
 from torch_sim.units import UnitConversion
 
@@ -24,9 +24,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py

@@ -14,7 +14,7 @@
 from mace.calculators.foundations_models import mace_mp
 
 import torch_sim as ts
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 from torch_sim.optimizers import unit_cell_fire
 from torch_sim.units import UnitConversion
 
@@ -24,9 +24,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py

@@ -15,7 +15,7 @@
 
 import torch_sim as ts
 from torch_sim.integrators import MDState, nvt_langevin
-from torch_sim.models.mace import MaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 from torch_sim.monte_carlo import swap_monte_carlo
 from torch_sim.units import MetalUnits
 
@@ -26,9 +26,8 @@
 kT = 1000 * MetalUnits.temperature
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py

@@ -11,8 +11,8 @@
 
 import torch
 
+import torch_sim as ts
 from torch_sim.quantities import calc_kinetic_energy, calc_kT
-from torch_sim.state import SimState
 from torch_sim.unbatched.models.lennard_jones import UnbatchedLennardJonesModel
 from torch_sim.unbatched.unbatched_integrators import npt_langevin
 from torch_sim.units import MetalUnits as Units
@@ -92,7 +92,7 @@
     compute_forces=True,
     compute_stress=True,
 )
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,

examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py

@@ -12,8 +12,9 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
+import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.quantities import calc_kinetic_energy, calc_kT
-from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_integrators import (
     npt_langevin,
@@ -28,9 +29,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,
@@ -62,7 +62,7 @@
     dtype=dtype,
     enable_cueq=False,
 )
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,

examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py

@@ -14,6 +14,7 @@
 from mace.calculators.foundations_models import mace_off
 
 import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.quantities import calc_kinetic_energy
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_integrators import nve
@@ -25,9 +26,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-off/raw/refs/heads/main/mace_off23/MACE-OFF23b_medium.model"
 loaded_model = mace_off(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,

examples/scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py

@@ -11,8 +11,8 @@
 
 import torch
 
+import torch_sim as ts
 from torch_sim.quantities import calc_kinetic_energy
-from torch_sim.state import SimState
 from torch_sim.unbatched.models.lennard_jones import UnbatchedLennardJonesModel
 from torch_sim.unbatched.unbatched_integrators import nve
 from torch_sim.units import MetalUnits as Units
@@ -76,7 +76,7 @@
 # Create the masses tensor (Argon = 39.948 amu)
 masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype)
 
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,

examples/scripts/3_Dynamics/3.2_MACE_NVE.py

@@ -13,8 +13,9 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
+import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.quantities import calc_kinetic_energy
-from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_integrators import nve
 from torch_sim.units import MetalUnits as Units
@@ -25,9 +26,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,
@@ -59,7 +59,7 @@
     enable_cueq=False,
 )
 
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,

examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py

@@ -13,8 +13,9 @@
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
+import torch_sim as ts
+from torch_sim.models.mace import MaceUrls
 from torch_sim.quantities import calc_kinetic_energy
-from torch_sim.state import SimState
 from torch_sim.unbatched.models.mace import UnbatchedMaceModel
 from torch_sim.unbatched.unbatched_integrators import nve
 from torch_sim.units import MetalUnits as Units
@@ -25,9 +26,8 @@
 dtype = torch.float32
 
 # Option 1: Load the raw model from the downloaded model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 loaded_model = mace_mp(
-    model=mace_checkpoint_url,
+    model=MaceUrls.mace_mpa_medium,
     return_raw_model=True,
     default_dtype=dtype,
     device=device,
@@ -58,7 +58,7 @@
     dtype=dtype,
     enable_cueq=torch.cuda.is_available(),
 )
-state = SimState(
+state = ts.SimState(
     positions=positions,
     masses=masses,
     cell=cell,