Test abstract model

tests/test_abstract_model.py

@@ -11,24 +11,51 @@
 def extxyz_file():
     return StringIO(
         """2
-        Properties=species:S:1:pos:R:3:forces:R:3 energy=-100 pbc="F F F"
-        Si  0.0  0.0  0.0  0.46610144  0.6694725  -0.49123398
-        Si  0.0  0.0  0.0  -0.15967495  -0.51654063  -0.67342865
+        Properties=species:S:1:pos:R:3:forces:R:3 energy=-1 pbc="F T F"
+        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
         """
     )
 
 
 def test_from_atoms(extxyz_file):
+    """Test extracting data from ASE Atoms object."""
+    expected_forces = np.array([[0.4, 0.6, -0.4], [-0.1, -0.5, -0.6]])
+    expected_stress = np.array([-1.0, -1.0, -1.0, -2.1, 2.0, 1.8])
+
     atoms = read(extxyz_file, format="extxyz")
-    expected_forces = np.array(
-        [
-            [0.46610144,  0.6694725 , -0.49123398],
-            [-0.15967495, -0.51654063, -0.67342865],
-        ]
-    )
-    # atoms.calc.results["forces"] = forces
+    atoms.calc.results["stress"] = expected_stress
     data = AbstractModel.from_atoms(atoms)
-    assert data is not None
-    assert data["pbc"] == [False, False, False]
-    assert data["energy"] == -100
-    assert data["forces"] == expected_forces
+
+    # Test info
+    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"
+
+    # Test results
+    assert data["energy"] == -1
+    assert data["forces"] == pytest.approx(expected_forces)
+    assert data["stress"] == pytest.approx(expected_stress)
+
+
+def test_from_atoms_no_calc(extxyz_file):
+    """Test extracting data from ASE Atoms object without results."""
+    expected_stress = np.array([-1.0, -1.0, -1.0, -2.1, 2.0, 1.8])
+
+    atoms = read(extxyz_file, format="extxyz")
+    atoms.calc.results["stress"] = expected_stress
+    data = AbstractModel.from_atoms(atoms, store_calc=False)
+
+    # Test info
+    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"
+
+    # Test results
+    assert "energy" not in data
+    assert "forces" not in data
+    assert "stress" not in data