update to new extra properties SimState interface

orb_models/forcefield/forcefield_adapter.py

@@ -497,6 +497,7 @@ def _get_charge_and_spin(atoms: ase.Atoms | ts.SimState) -> dict[str, torch.Tens
     elif (
         _TORCH_SIM_AVAILABLE
         and isinstance(atoms, ts.SimState)
+        and (atoms.has_extras("charge") or atoms.has_extras("spin"))
         and (atoms.charge is not None or atoms.spin is not None)
     ):
         assert atoms.charge is not None and atoms.spin is not None, (

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-torchsim = ["torch-sim-atomistic>=0.5.1"]
+torchsim = ["torch-sim-atomistic>=0.6.0"]
 
 [tool.setuptools.packages.find]
 where = ["."]

tests/forcefield/test_orb_torchsim.py

@@ -4,72 +4,44 @@
 
 import torch
 import torch_sim as ts
-from torch_sim.elastic import full_3x3_to_voigt_6_stress
+from torch_sim.models.interface import validate_model_outputs
+from torch_sim.testing import (
+    CONSISTENCY_SIMSTATES,
+    SIMSTATE_GENERATORS,
+    assert_model_calculator_consistency,
+)
 
 from orb_models.forcefield.forcefield_adapter import ForcefieldAtomsAdapter
 from orb_models.forcefield.inference.calculator import ORBCalculator
 from orb_models.forcefield.inference.orb_torchsim import OrbTorchSimModel
 
+DEVICE = torch.device("cpu")
+DTYPE = torch.float64
+
 
+# Use the TorchSim consistency check harness for external models
+@pytest.mark.parametrize("sim_state_name", CONSISTENCY_SIMSTATES)
 @pytest.mark.parametrize(
     "edge_method",
     ["knn_scipy", "knn_alchemi"],
 )
-def test_orb_torchsim_interface(edge_method, conservative_regressor, mptraj_10_systems_db):
-    atoms_list = [mptraj_10_systems_db.get_atoms(i) for i in range(1, 5)]
-
+def test_orb_torchsim_consistency(sim_state_name, edge_method, conservative_regressor):
     adapter = ForcefieldAtomsAdapter(6.0, 120)
     calculator = ORBCalculator(
         model=conservative_regressor,
         atoms_adapter=adapter,
         edge_method=edge_method,
     )
-
-    # TorchSim results
-    sim_state = ts.io.atoms_to_state(atoms_list, "cpu", torch.get_default_dtype())
     sim_model = OrbTorchSimModel(conservative_regressor, adapter, edge_method=edge_method)
-    model_results = sim_model(sim_state)
 
-    # ASE calculator results
-    calc_energy_list = []
-    calc_force_list = []
-    calc_stress_list = []
-    for atoms in atoms_list:
-        atoms.calc = calculator
-        calc_energy = atoms.get_potential_energy()
-        calc_forces = torch.tensor(
-            atoms.get_forces(),
-            dtype=model_results["forces"].dtype,
-        )
-        calc_stress = torch.tensor(
-            atoms.get_stress(),
-            dtype=model_results["stress"].dtype,
-        )
-        calc_energy_list.append(calc_energy)
-        calc_force_list.append(calc_forces)
-        calc_stress_list.append(calc_stress)
+    sim_state = SIMSTATE_GENERATORS[sim_state_name](DEVICE, DTYPE)
+    assert_model_calculator_consistency(sim_model, calculator, sim_state)
 
-    # Test consistency with specified tolerances
-    torch.testing.assert_close(
-        model_results["energy"],
-        torch.tensor(calc_energy_list),
-        rtol=1e-5,
-        atol=1e-5,
-    )
-    torch.testing.assert_close(
-        model_results["forces"],
-        torch.cat(calc_force_list),
-        rtol=1e-5,
-        atol=1e-5,
-    )
-    if "stress" in model_results:
-        torch.testing.assert_close(
-            full_3x3_to_voigt_6_stress(model_results["stress"]),
-            torch.stack(calc_stress_list),
-            rtol=1e-5,
-            atol=1e-5,
-            equal_nan=True,
-        )
+
+def test_orb_torchsim_validate_outputs(conservative_regressor):
+    adapter = ForcefieldAtomsAdapter(6.0, 120)
+    sim_model = OrbTorchSimModel(conservative_regressor, adapter, dtype=DTYPE)
+    validate_model_outputs(sim_model, device=DEVICE, dtype=DTYPE)
 
 
 class TestOrbTorchSimStressToggle: