diff --git a/orb_models/forcefield/forcefield_adapter.py b/orb_models/forcefield/forcefield_adapter.py index 44d8563..45859b6 100644 --- a/orb_models/forcefield/forcefield_adapter.py +++ b/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, ( diff --git a/pyproject.toml b/pyproject.toml index 092a6ab..5fbcb6e 100644 --- a/pyproject.toml +++ b/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 = ["."] diff --git a/tests/forcefield/test_orb_torchsim.py b/tests/forcefield/test_orb_torchsim.py index ac37c7d..54a7ce9 100644 --- a/tests/forcefield/test_orb_torchsim.py +++ b/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: