TorchSim · janosh · Apr 17, 2025 · Apr 8, 2025 · Apr 8, 2025 · Apr 16, 2025
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -18,8 +18,6 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.workflow_run.head_branch }}
 
       - name: Install pandoc
         run: sudo apt-get install pandoc

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,12 +1,14 @@
 ## v0.2.0
 
 ### Bug Fixes 🐛
+
 * Fix integrate reporting kwarg to arg error, https://github.com/Radical-AI/torch-sim/issues/113 (raised by @hn-yu)
 * Allow runners to take large initial batches, https://github.com/Radical-AI/torch-sim/issues/128 (raised by @YutackPark)
 * Add Fairchem model support for PBC, https://github.com/Radical-AI/torch-sim/issues/111 (raised by @ryanliu30)
 
 ### Enhancements 🛠
-* **[breaking]** Rename `HotSwappingAutobatcher` to `InFlightAutobatcher` and `ChunkingAutoBatcher` to `BinningAutoBatcher`, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher
+
+* **breaking** Rename `HotSwappingAutobatcher` to `InFlightAutobatcher` and `ChunkingAutoBatcher` to `BinningAutoBatcher`, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher
 * Support for Orbv3, https://github.com/Radical-AI/torch-sim/pull/140, @AdeeshKolluru
 * Support metatensor models, https://github.com/Radical-AI/torch-sim/pull/141, @frostedoyter @Luthaf
 * Support for graph-pes models, https://github.com/Radical-AI/torch-sim/pull/118 @jla-gardner
@@ -18,10 +20,12 @@
 * New correlation function module, https://github.com/Radical-AI/torch-sim/pull/115 @stefanbringuier
 
 ### Documentation 📖
+
 * Imoved model documentation, https://github.com/Radical-AI/torch-sim/pull/121 @orionarcher
 * Plot of TorchSim module graph in docs, https://github.com/Radical-AI/torch-sim/pull/132 @janosh
 
 ### House-Keeping 🧹
+
 * Only install HF for fairchem tests, https://github.com/Radical-AI/torch-sim/pull/134 @CompRhys
 * Don't download MBD in CI, https://github.com/Radical-AI/torch-sim/pull/135 @orionarcher
 * Tighten graph-pes test bounds, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher

diff --git a/examples/scripts/1_Introduction/1.1_Lennard_Jones.py b/examples/scripts/1_Introduction/1.1_Lennard_Jones.py
@@ -88,7 +88,7 @@
 results = model(state)
 
 # Print the results
-print(f"Energy: {results['energy']}")
+print(f"Energy: {float(results['energy']):.4f}")
 print(f"Forces: {results['forces']}")
 print(f"Stress: {results['stress']}")
 print(f"Energies: {results['energies']}")

diff --git a/examples/scripts/1_Introduction/1.2_MACE.py b/examples/scripts/1_Introduction/1.2_MACE.py
@@ -19,7 +19,7 @@
 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"
+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,
     return_raw_model=True,
@@ -51,6 +51,6 @@
 results = model(state)
 
 # Print results
-print(f"Energy: {results['energy']}")
+print(f"Energy: {float(results['energy']):.4f}")
 print(f"Forces: {results['forces']}")
 print(f"Stress: {results['stress']}")
diff --git a/examples/scripts/1_Introduction/1.3_Batched_MACE.py b/examples/scripts/1_Introduction/1.3_Batched_MACE.py
@@ -19,7 +19,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/2_Structural_optimization/2.10_Batched_MACE_FrechetCellFilter_FIRE.py b/examples/scripts/2_Structural_optimization/2.10_Batched_MACE_FrechetCellFilter_FIRE.py
@@ -22,7 +22,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/2_Structural_optimization/2.3_MACE_FIRE.py b/examples/scripts/2_Structural_optimization/2.3_MACE_FIRE.py
@@ -23,7 +23,7 @@
 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"
+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,
     return_raw_model=True,

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
@@ -24,7 +24,7 @@
 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"
+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,
     return_raw_model=True,

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
@@ -24,7 +24,7 @@
 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"
+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,
     return_raw_model=True,

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
@@ -23,7 +23,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/...les/scripts/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py b/...les/scripts/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py
@@ -23,7 +23,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py b/examples/scripts/2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py
@@ -23,7 +23,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py b/examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py
@@ -26,7 +26,7 @@
 kT = 1000 * MetalUnits.temperature
 
 # 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py b/examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py
@@ -28,7 +28,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/3_Dynamics/3.2_MACE_NVE.py b/examples/scripts/3_Dynamics/3.2_MACE_NVE.py
@@ -25,7 +25,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py b/examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py
@@ -25,7 +25,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py b/examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py
@@ -24,7 +24,7 @@
 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"
+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,
     return_raw_model=True,

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
@@ -27,7 +27,7 @@
 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"
+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,
     return_raw_model=True,

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
@@ -81,7 +81,7 @@ def get_kT(
 
 # Model configuration
 # Option 1: Load from URL (uncomment to use)
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.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,
     return_raw_model=True,

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
@@ -29,7 +29,7 @@
 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"
+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,
     return_raw_model=True,

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
@@ -27,7 +27,7 @@
 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"
+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,
     return_raw_model=True,

diff --git a/examples/scripts/5_Workflow/5.1_a2c_silicon.py b/examples/scripts/5_Workflow/5.1_a2c_silicon.py
@@ -55,7 +55,7 @@
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float32
 
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
+mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 raw_model = mace_mp(model=mace_checkpoint_url, return_raw_model=True)
 
 # Define system and model

diff --git a/examples/scripts/5_Workflow/5.2_a2c_silicon_batched.py b/examples/scripts/5_Workflow/5.2_a2c_silicon_batched.py
@@ -55,7 +55,7 @@
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float32
 
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
+mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 raw_model = mace_mp(model=mace_checkpoint_url, return_raw_model=True)
 
 # Define system and model

diff --git a/examples/scripts/5_Workflow/5.3_In_Flight_WBM.py b/examples/scripts/5_Workflow/5.3_In_Flight_WBM.py
@@ -25,7 +25,7 @@
 
 # --- Model Initialization ---
 print("Loading MACE model...")
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
+mace_checkpoint_url = "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model"
 mace = mace_mp(model=mace_checkpoint_url, return_raw_model=True)
 mace_model = ts.models.MaceModel(
     model=mace,

diff --git a/examples/scripts/5_Workflow/5.4_Elastic.py b/examples/scripts/5_Workflow/5.4_Elastic.py
@@ -19,7 +19,7 @@
 unit_conv = ts.units.UnitConversion
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float64
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.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,
     enable_cueq=False,

diff --git a/examples/scripts/6_Phonons/6.1_Phonons_MACE.py b/examples/scripts/6_Phonons/6.1_Phonons_MACE.py
@@ -89,7 +89,7 @@ def get_labels_qpts(ph: Phonopy, n_points: int = 101) -> tuple[list[str], list[b
 dtype = torch.float32
 
 # Load the raw model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.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,
     return_raw_model=True,

diff --git a/examples/scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py b/examples/scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py
@@ -212,7 +212,7 @@ def get_qha_phonons(
 autobatcher = False
 
 # Load the raw model
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.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,
     return_raw_model=True,

diff --git a/examples/scripts/6_Phonons/6.3_Conductivity_MACE.py b/examples/scripts/6_Phonons/6.3_Conductivity_MACE.py
@@ -50,7 +50,7 @@ def print_relax_info(trajectory_file: str, device: torch.device) -> None:
 dtype = torch.float64
 
 # Load the raw model from URL
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.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, return_raw_model=True, default_dtype=dtype, device=device
 )
@@ -66,11 +66,8 @@ def print_relax_info(trajectory_file: str, device: torch.device) -> None:
 # Structure and input parameters
 struct = bulk("Si", "diamond", a=5.431, cubic=True)  # ASE structure
 mesh = [8, 8, 8]  # Phonon mesh
-supercell_matrix = [
-    1,
-    1,
-    1,
-]  # supercell matrix for phonon calculation (use larger supercell for better accuracy)
+# supercell matrix for phonon calculation (use larger cell for better accuracy)
+supercell_matrix = [1, 1, 1]
 supercell_matrix_fc2 = [2, 2, 2]  # supercell matrix for FC2 calculation
 Nrelax = 300  # number of relaxation steps
 fmax = 1e-3  # force convergence

diff --git a/examples/tutorials/low_level_tutorial.py b/examples/tutorials/low_level_tutorial.py
@@ -68,7 +68,7 @@
 from torch_sim.models import MaceModel
 
 # load mace_mp using the mace package
-mace_checkpoint_url = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.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,
     return_raw_model=True,

diff --git a/tests/test_elastic.py b/tests/test_elastic.py
@@ -2,13 +2,13 @@
 import torch
 
 from torch_sim.elastic import (
-    BravaisType,
     calculate_elastic_moduli,
     calculate_elastic_tensor,
     get_bravais_type,
 )
 from torch_sim.optimizers import frechet_cell_fire
 from torch_sim.state import SimState
+from torch_sim.typing import BravaisType
 from torch_sim.units import UnitConversion
 
 

diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py
@@ -33,7 +33,7 @@ def ase_to_torch_batch(
             Defaults to torch.float32.
 
     Returns:
-        Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
             A tuple containing:
             - pos: Tensor of atomic positions.
             - cell: Tensor of unit cell vectors.

diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py
@@ -23,12 +23,13 @@
 import logging
 from collections.abc import Callable, Iterator
 from itertools import chain
-from typing import Any, Literal
+from typing import Any, get_args
 
 import torch
 
 from torch_sim.models.interface import ModelInterface
 from torch_sim.state import SimState, concatenate_states
+from torch_sim.typing import MemoryScaling
 
 
 def to_constant_volume_bins(  # noqa: C901, PLR0915
@@ -262,7 +263,7 @@ def determine_max_batch_size(
     error or reaches the specified maximum atom count.
 
     Args:
-        state (SimState): SimState to replicate for testing.
+        state (SimState): State to replicate for testing.
         model (ModelInterface): Model to test with.
         max_atoms (int): Upper limit on number of atoms to try (for safety).
             Defaults to 500,000.
@@ -309,7 +310,7 @@ def determine_max_batch_size(
 
 def calculate_memory_scaler(
     state: SimState,
-    memory_scales_with: Literal["n_atoms_x_density", "n_atoms"] = "n_atoms_x_density",
+    memory_scales_with: MemoryScaling = "n_atoms_x_density",
 ) -> float:
     """Calculate a metric that estimates memory requirements for a state.
 
@@ -322,7 +323,7 @@ def calculate_memory_scaler(
     Args:
         state (SimState): State to calculate metric for, with shape information
             specific to the SimState instance.
-        memory_scales_with ("n_atoms_x_density" |s "n_atoms"): Type of metric
+        memory_scales_with ("n_atoms_x_density" | "n_atoms"): Type of metric
             to use. "n_atoms" uses only atom count and is suitable for models that
             have a fixed number of neighbors. "n_atoms_x_density" uses atom count
             multiplied by number density and is better for models with radial cutoffs
@@ -351,7 +352,9 @@ def calculate_memory_scaler(
         volume = torch.abs(torch.linalg.det(state.cell[0])) / 1000
         number_density = state.n_atoms / volume.item()
         return state.n_atoms * number_density
-    raise ValueError(f"Invalid metric: {memory_scales_with}")
+    raise ValueError(
+        f"Invalid metric: {memory_scales_with}, must be one of {get_args(MemoryScaling)}"
+    )
 
 
 def estimate_max_memory_scaler(
@@ -458,7 +461,7 @@ def __init__(
         self,
         model: ModelInterface,
         *,
-        memory_scales_with: Literal["n_atoms", "n_atoms_x_density"] = "n_atoms_x_density",
+        memory_scales_with: MemoryScaling = "n_atoms_x_density",
         max_memory_scaler: float | None = None,
         return_indices: bool = False,
         max_atoms_to_try: int = 500_000,
@@ -752,7 +755,7 @@ def __init__(
         self,
         model: ModelInterface,
         *,
-        memory_scales_with: Literal["n_atoms", "n_atoms_x_density"] = "n_atoms_x_density",
+        memory_scales_with: MemoryScaling = "n_atoms_x_density",
         max_memory_scaler: float | None = None,
         max_atoms_to_try: int = 500_000,
         memory_scaling_factor: float = 1.6,