Test to ASE function

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,13 +35,15 @@ 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]
     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)
@@ -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