Use first-party implementation of Orb Model Interface.

pyproject.toml

@@ -52,7 +52,7 @@ symmetry = ["moyopy>=0.3"]
 mace = ["mace-torch>=0.3.14"]
 mattersim = ["mattersim>=0.1.2"]
 metatomic = ["metatomic-torch>=0.1.3", "metatrain[pet]>=2025.12"]
-orb = ["orb-models>=0.5.2"]
+orb = ["orb-models[torchsim] @ git+https://github.com/orbital-materials/orb-models.git"]
 sevenn = ["sevenn>=0.11.0"]
 graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"]
 nequip = ["nequip>=0.16.2"]

tests/models/test_orb.py

@@ -1,22 +1,18 @@
 import traceback
 
-import numpy as np
 import pytest
-from ase.geometry.cell import cell_to_cellpar as ase_c2p
 
 from tests.conftest import DEVICE
 from tests.models.conftest import (
     make_model_calculator_consistency_test,
     make_validate_model_outputs_test,
 )
-from torch_sim import SimState
-from torch_sim.models.orb import cell_to_cellpar
 from torch_sim.testing import SIMSTATE_GENERATORS
 
 
 try:
     from orb_models.forcefield import pretrained
-    from orb_models.forcefield.calculator import ORBCalculator
+    from orb_models.forcefield.inference.calculator import ORBCalculator
 
     from torch_sim.models.orb import OrbModel
 except ImportError:
@@ -25,32 +21,36 @@
 
 @pytest.fixture
 def orbv3_conservative_inf_omat_model() -> OrbModel:
-    orb_ff = pretrained.orb_v3_conservative_inf_omat(
+    orb_ff, atoms_adapter = pretrained.orb_v3_conservative_inf_omat(
         device=DEVICE, precision="float32-high"
     )
-    return OrbModel(model=orb_ff, device=DEVICE)
+    return OrbModel(orb_ff, atoms_adapter, device=DEVICE)
 
 
 @pytest.fixture
 def orbv3_direct_20_omat_model() -> OrbModel:
-    orb_ff = pretrained.orb_v3_direct_20_omat(device=DEVICE, precision="float32-high")
-    return OrbModel(model=orb_ff, device=DEVICE)
+    orb_ff, atoms_adapter = pretrained.orb_v3_direct_20_omat(
+        device=DEVICE, precision="float32-high"
+    )
+    return OrbModel(orb_ff, atoms_adapter, device=DEVICE)
 
 
 @pytest.fixture
 def orbv3_conservative_inf_omat_calculator() -> ORBCalculator:
     """Create an ORBCalculator for the pretrained model."""
-    orb_ff = pretrained.orb_v3_conservative_inf_omat(
+    orb_ff, atoms_adapter = pretrained.orb_v3_conservative_inf_omat(
         device=DEVICE, precision="float32-high"
     )
-    return ORBCalculator(model=orb_ff, device=DEVICE)
+    return ORBCalculator(model=orb_ff, atoms_adapter=atoms_adapter, device=DEVICE)
 
 
 @pytest.fixture
 def orbv3_direct_20_omat_calculator() -> ORBCalculator:
     """Create an ORBCalculator for the pretrained model."""
-    orb_ff = pretrained.orb_v3_direct_20_omat(device=DEVICE, precision="float32-high")
-    return ORBCalculator(model=orb_ff, device=DEVICE)
+    orb_ff, atoms_adapter = pretrained.orb_v3_direct_20_omat(
+        device=DEVICE, precision="float32-high"
+    )
+    return ORBCalculator(model=orb_ff, atoms_adapter=atoms_adapter, device=DEVICE)
 
 
 test_orb_conservative_consistency = make_model_calculator_consistency_test(
@@ -78,16 +78,3 @@ def orbv3_direct_20_omat_calculator() -> ORBCalculator:
 test_validate_direct_model_outputs = make_validate_model_outputs_test(
     model_fixture_name="orbv3_direct_20_omat_model",
 )
-
-
-def test_cell_to_cellpar(ti_sim_state: SimState) -> None:
-    assert np.allclose(
-        ase_c2p(ti_sim_state.row_vector_cell.squeeze()),
-        cell_to_cellpar(ti_sim_state.row_vector_cell.squeeze(0)).cpu().numpy(),
-    )
-    assert np.allclose(
-        ase_c2p(ti_sim_state.row_vector_cell.squeeze(), radians=True),
-        cell_to_cellpar(ti_sim_state.row_vector_cell.squeeze(0), radians=True)
-        .cpu()
-        .numpy(),
-    )