Improve Typing of ModelInterface (#215)

Signed-off-by: Rhys Goodall <rhys.goodall@outlook.com>
Co-authored-by: Rhys Goodall <rhys.goodall@outlook.com>

Author: curtischong

.github/PULL_REQUEST_TEMPLATE.md

@@ -8,7 +8,8 @@
 Before a pull request can be merged, the following items must be checked:
 
 * [ ] Doc strings have been added in the [Google docstring format](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html#example-google).
-  Run [ruff](https://beta.ruff.rs/docs/rules/#pydocstyle-d) on your code.
+* [ ] Run [ruff](https://beta.ruff.rs/docs/rules/#pydocstyle-d) on your code.
+* [ ] Run `uvx ty check` on the repo.
 * [ ] Tests have been added for any new functionality or bug fixes.
 
 We highly recommended installing the pre-commit hooks running in CI locally to speedup the development process. Simply run `pip install pre-commit && pre-commit install` to install the hooks which will check your code before each commit.

examples/scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py

@@ -89,8 +89,8 @@
     positions=positions,
     masses=masses,
     cell=cell.unsqueeze(0),
-    pbc=True,
     atomic_numbers=atomic_numbers,
+    pbc=True,
 )
 
 # Run initial simulation and get results

examples/scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py

@@ -80,8 +80,8 @@
     positions=positions,
     masses=masses,
     cell=cell.unsqueeze(0),
-    pbc=True,
     atomic_numbers=atomic_numbers,
+    pbc=True,
 )
 
 # Initialize the Soft Sphere model

examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py

@@ -97,8 +97,8 @@
     positions=positions,
     masses=masses,
     cell=cell.unsqueeze(0),
-    pbc=True,
     atomic_numbers=atomic_numbers,
+    pbc=True,
 )
 # Run initial simulation and get results
 results = model(state)
@@ -148,11 +148,11 @@
 
 
 stress = model(state)["stress"]
-calc_kinetic_energy = calc_kinetic_energy(
+kinetic_energy = calc_kinetic_energy(
     masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
 )
 volume = torch.linalg.det(state.cell)
-pressure = get_pressure(stress, calc_kinetic_energy, volume)
+pressure = get_pressure(stress, kinetic_energy, volume)
 pressure = pressure.item() / Units.pressure
 print(f"Final {pressure=:.4f}")
 print(stress * UnitConversion.eV_per_Ang3_to_GPa)

examples/scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py

@@ -78,11 +78,7 @@
 masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype)
 
 state = ts.SimState(
-    positions=positions,
-    masses=masses,
-    cell=cell,
-    pbc=True,
-    atomic_numbers=atomic_numbers,
+    positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True
 )
 # Initialize the Lennard-Jones model
 # Parameters:

examples/scripts/3_Dynamics/3.2_MACE_NVE.py

@@ -60,11 +60,7 @@
 )
 
 state = ts.SimState(
-    positions=positions,
-    masses=masses,
-    cell=cell,
-    pbc=True,
-    atomic_numbers=atomic_numbers,
+    positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True
 )
 
 # Run initial inference

examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py

@@ -59,11 +59,7 @@
 )
 
 state = ts.SimState(
-    positions=positions,
-    masses=masses,
-    cell=cell,
-    pbc=True,
-    atomic_numbers=atomic_numbers,
+    positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True
 )
 
 dt = 0.002 * Units.time  # Timestep (ps)

examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py

@@ -96,8 +96,8 @@
     positions=positions,
     masses=masses,
     cell=cell.unsqueeze(0),
-    pbc=True,
     atomic_numbers=atomic_numbers,
+    pbc=True,
 )
 # Run initial simulation and get results
 results = model(state)

examples/scripts/4_High_level_api/4.1_high_level_api.py

@@ -54,7 +54,9 @@
 prop_calculators = {
     10: {"potential_energy": lambda state: state.energy},
     20: {
-        "kinetic_energy": lambda state: calc_kinetic_energy(state.momenta, state.masses)
+        "kinetic_energy": lambda state: calc_kinetic_energy(
+            momenta=state.momenta, masses=state.masses
+        )
     },
 }
 

examples/scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py

@@ -24,12 +24,13 @@
 from phonopy.structure.atoms import PhonopyAtoms
 
 import torch_sim as ts
+from torch_sim.models.interface import ModelInterface
 from torch_sim.models.mace import MaceModel, MaceUrls
 
 
 def get_relaxed_structure(
     struct: Atoms,
-    model: torch.nn.Module | None,
+    model: ModelInterface,
     Nrelax: int = 300,
     fmax: float = 1e-3,
     *,
@@ -80,7 +81,7 @@ def get_relaxed_structure(
 def get_qha_structures(
     state: ts.state.SimState,
     length_factors: np.ndarray,
-    model: torch.nn.Module | None,
+    model: ModelInterface,
     Nmax: int = 300,
     fmax: float = 1e-3,
     *,
@@ -129,7 +130,7 @@ def get_qha_structures(
 
 def get_qha_phonons(
     scaled_structures: list[PhonopyAtoms],
-    model: torch.nn.Module | None,
+    model: ModelInterface,
     supercell_matrix: np.ndarray | None,
     displ: float = 0.05,
     *,