From cd50791d5efcede52e94e1cd5afa403382bf6c9b Mon Sep 17 00:00:00 2001
From: zhiyili1230 <zhiyi@orbitalmaterials.com>
Date: Fri, 4 Oct 2024 12:55:26 +0000
Subject: [PATCH] tests ok

---
 orb_models/__init__.py                     |  2 +
 tests/test_model_backward_compatibility.py | 57 +++++++++++++++-------
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/orb_models/__init__.py b/orb_models/__init__.py
index e69de29..4f536c2 100644
--- a/orb_models/__init__.py
+++ b/orb_models/__init__.py
@@ -0,0 +1,2 @@
+# orb_models/__init__.py
+__version__ = "0.5.0"
diff --git a/tests/test_model_backward_compatibility.py b/tests/test_model_backward_compatibility.py
index 746b4e3..529d1a4 100644
--- a/tests/test_model_backward_compatibility.py
+++ b/tests/test_model_backward_compatibility.py
@@ -1,38 +1,61 @@
+import pytest
 import numpy as np
-from ase.optimize import BFGS
 from ase.build import bulk
+from ase.optimize import BFGS
 from orb_models.forcefield import atomic_system, pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
 
-def test_energy_and_forces_prediction():
-    orbff = pretrained.orb_v1(device="cpu")
+@pytest.mark.parametrize("model_fn", [pretrained.orb_v1, pretrained.orb_v2])
+def test_energy_forces_stress_prediction(model_fn):
+    """Tests model compatibility on energy, forces and stress."""
+    orbff = model_fn(device="cpu")
     atoms = bulk("Cu", "fcc", a=3.58, cubic=True)
     graph = atomic_system.ase_atoms_to_atom_graphs(atoms, device="cpu")
     result = orbff.predict(graph)
-    energy_gold = np.array(-16.3437)
+
     energy = result["graph_pred"][0]
-    print(energy)
-    forces_gold = np.array([9.3569e-05, -4.0413e-05, 6.0380e-05])
-    forces = result["node_pred"][0][0]
-    print(forces)
-    stress_gold = np.array(
-        [-3.2064e-02, -3.2115e-02, -3.1964e-02, -6.2145e-07, 2.5603e-07, 1.6000e-06]
-    )
+    forces = result["node_pred"][0]
     stress = result["stress_pred"][0]
-    print(stress)
 
+    if model_fn == pretrained.orb_v1:
+        energy_gold = np.array(-16.3437)
+        forces_gold = np.array([9.3569e-05, -4.0413e-05, 6.0380e-05])
+        stress_gold = np.array(
+            [-3.2064e-02, -3.2115e-02, -3.1964e-02, -6.2145e-07, 2.5603e-07, 1.6000e-06]
+        )
+    elif model_fn == pretrained.orb_v2:
+        energy_gold = np.array(-16.3459)
+        forces_gold = np.array([-1.0881e-06, 7.6937e-08, -1.7336e-06])
+        stress_gold = np.array(
+            [-3.2431e-02, -3.2185e-02, -3.1791e-02, 1.5402e-06, 1.9556e-06, 6.9951e-07]
+        )
 
-def test_optimization():
-    # Rattle the atoms to get them out of the minimum energy configuration
+    assert np.isclose(energy, energy_gold, atol=1e-4)
+    assert np.allclose(forces, forces_gold, atol=1e-6)
+    assert np.allclose(stress, stress_gold, atol=1e-6)
+
+
+@pytest.mark.parametrize("model_fn", [pretrained.orb_v1, pretrained.orb_v2])
+def test_optimization(model_fn):
+    """Tests model compatibility on optimization."""
     atoms = bulk("Cu", "fcc", a=3.58, cubic=True)
-    orbff = pretrained.orb_v1(device="cpu")
+    orbff = model_fn(device="cpu")
     calc = ORBCalculator(orbff, device="cpu")
     atoms.set_calculator(calc)
     atoms.rattle(0.5)
     rattled_energy = atoms.get_potential_energy()
-    print(rattled_energy)
     dyn = BFGS(atoms)
     dyn.run(fmax=0.01)
     optimized_energy = atoms.get_potential_energy()
-    print(optimized_energy)
+    if model_fn == pretrained.orb_v1:
+        gold_rattled_energy = -11.943148
+        gold_optimized_energy = -16.345758
+    elif model_fn == pretrained.orb_v2:
+        print(rattled_energy)
+        print(optimized_energy)
+        gold_rattled_energy = -12.034759
+        gold_optimized_energy = -16.348310
+
+    assert np.isclose(rattled_energy, gold_rattled_energy, atol=1e-6)
+    assert np.isclose(optimized_energy, gold_optimized_energy, atol=1e-6)