diff --git a/tests/test_relaxation.py b/tests/test_relaxation.py index c23b675..f9889aa 100644 --- a/tests/test_relaxation.py +++ b/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,10 @@ 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) == 12, 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.3) no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device")