From 51cc5c02528952a53fda0bcdb1a35331414b9f8f Mon Sep 17 00:00:00 2001 From: ElliottKasoar Date: Thu, 28 Nov 2024 17:10:31 +0000 Subject: [PATCH] Test to ASE function --- tests/test_abstract_model.py | 50 ++++++++++++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/tests/test_abstract_model.py b/tests/test_abstract_model.py index 78964ab..5d4d7c6 100644 --- a/tests/test_abstract_model.py +++ b/tests/test_abstract_model.py @@ -11,7 +11,7 @@ def extxyz_file(): return StringIO( """2 - Properties=species:S:1:pos:R:3:forces:R:3 energy=-1 pbc="F T F" + Properties=species:S:1:pos:R:3:forces:R:3 energy=-1 pbc="F T F" info="test" Si 0.0 0.0 0.0 0.4 0.6 -0.4 Si 0.0 0.0 0.0 -0.1 -0.5 -0.6 """ @@ -35,6 +35,7 @@ def test_from_atoms(extxyz_file): "formula", "calculator_name", "calculator_parameters", + "info", } assert info_keys == set(data.info_keys) assert data["pbc"] == [False, True, False] @@ -42,6 +43,7 @@ def test_from_atoms(extxyz_file): assert len(data["cell"]) == 3 assert all(arr == [0.0, 0.0, 0.0] for arr in data["cell"]) assert data["formula"] == "Si2" + assert data["info"] == "test" # Test arrays assert {"numbers", "positions"} == set(data.arrays_keys) @@ -74,12 +76,13 @@ def test_from_atoms_no_calc(extxyz_file): data = AbstractModel.from_atoms(atoms, store_calc=False) # Test info - assert {"pbc", "n_atoms", "cell", "formula"} == set(data.info_keys) + assert {"pbc", "n_atoms", "cell", "formula", "info"} == set(data.info_keys) assert data["pbc"] == [False, True, False] assert data["n_atoms"] == 2 assert len(data["cell"]) == 3 assert all(arr == [0.0, 0.0, 0.0] for arr in data["cell"]) assert data["formula"] == "Si2" + assert data["info"] == "test" # Test arrays assert {"numbers", "positions"} == set(data.arrays_keys) @@ -105,3 +108,46 @@ def test_from_atoms_no_calc(extxyz_file): "hash_structure", } assert derived_keys == set(data.derived_keys) + + +def test_to_ase(extxyz_file): + """Test returning data to ASE Atoms object with results.""" + atoms = read(extxyz_file, format="extxyz") + data = AbstractModel.from_atoms(atoms, store_calc=True) + + new_atoms = data.to_ase() + + # Test info set + assert new_atoms.cell == pytest.approx(atoms.cell) + assert new_atoms.pbc == pytest.approx(atoms.pbc) + assert new_atoms.positions == pytest.approx(atoms.positions) + assert new_atoms.numbers == pytest.approx(atoms.numbers) + + assert new_atoms.info["n_atoms"] == len(atoms) + assert new_atoms.info["formula"] == atoms.get_chemical_formula() + + assert new_atoms.calc.results["energy"] == pytest.approx( + atoms.calc.results["energy"] + ) + assert new_atoms.calc.results["forces"] == pytest.approx( + atoms.calc.results["forces"] + ) + + +def test_to_ase_no_results(extxyz_file): + """Test returning data to ASE Atoms object without results.""" + atoms = read(extxyz_file, format="extxyz") + data = AbstractModel.from_atoms(atoms, store_calc=False) + + new_atoms = data.to_ase() + + # Test info set + assert new_atoms.cell == pytest.approx(atoms.cell) + assert new_atoms.pbc == pytest.approx(atoms.pbc) + assert new_atoms.positions == pytest.approx(atoms.positions) + assert new_atoms.numbers == pytest.approx(atoms.numbers) + + assert new_atoms.info["n_atoms"] == len(atoms) + assert new_atoms.info["formula"] == atoms.get_chemical_formula() + + assert new_atoms.calc is None