improve: ase try to get virials from different sources

deepmodeling · May 17, 2024 · aa8d22e · aa8d22e
1 parent 626e692
commit aa8d22e
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/dpdata/plugins/ase.py b/dpdata/plugins/ase.py
@@ -93,13 +93,21 @@ def from_labeled_system(self, atoms: ase.Atoms, **kwargs) -> dict:
             "energies": np.array([energies]),
             "forces": np.array([forces]),
         }
-        try:
-            stress = atoms.get_stress(False)
-        except PropertyNotImplementedError:
-            pass
-        else:
-            virials = np.array([-atoms.get_volume() * stress])
+
+        # try to get virials from different sources
+        virials = atoms.info.get("virial")
+        if virials is None:
+            virials = atoms.info.get("virials")
+        if virials is None:
+            try:
+                stress = atoms.get_stress(False)
+            except PropertyNotImplementedError:
+                pass
+            else:
+                virials = np.array([-atoms.get_volume() * stress])
+        if virials is not None:
             info_dict["virials"] = virials
+
         return info_dict
 
     def from_multi_systems(