tweak test_relaxation

chgnet/model/dynamics.py

@@ -170,6 +170,7 @@ def calculate(
         self.results |= {
             long_key: model_prediction[key] * factor
             for key, (long_key, factor) in key_map.items()
+            if key in model_prediction
         }
         self.results["free_energy"] = self.results["energy"]
         self.results["crystal_fea"] = model_prediction["crystal_fea"]

tests/test_relaxation.py

@@ -32,7 +32,11 @@ def test_relaxation(
     chgnet.graph_converter = converter
     relaxer = StructOptimizer(model=chgnet)
     result = relaxer.relax(
-        li_mn_o2, verbose=True, ase_filter=ase_filter, assign_magmoms=assign_magmoms
+        li_mn_o2,
+        verbose=True,
+        ase_filter=ase_filter,
+        assign_magmoms=assign_magmoms,
+        steps=10,  # few steps for faster testing
     )
     assert list(result) == ["final_structure", "trajectory"]
     final_struct, traj = result["final_structure"], result["trajectory"]
@@ -50,10 +54,12 @@ def test_relaxation(
     assert {*traj.__dict__} == {
         *"atoms energies forces stresses magmoms atom_positions cells".split()
     }
-    assert len(traj) == 2 if algorithm == "legacy" else 4
+    assert len(traj) == (
+        2 if algorithm == "legacy" else 4
+    ), f"{len(traj)=}, {algorithm=}"
 
     # make sure final structure is more relaxed than initial one
-    assert traj.energies[-1] == pytest.approx(-58.94209, rel=1e-4)
+    assert traj.energies[-1] == pytest.approx(-58.6675376, abs=0.2)
 
 
 no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device")