Enable loading of virials for 'qe/cp/traj' if prefix.str file is present

dpdata/plugins/qe.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 import dpdata.md.pbc
 import dpdata.qe.scf
 import dpdata.qe.traj
@@ -26,9 +28,13 @@
             data["coords"],
             data["cells"],
         )
-        data["energies"], data["forces"], es = dpdata.qe.traj.to_system_label(
+        data["energies"], data["forces"], stress, es = dpdata.qe.traj.to_system_label(
             file_name + ".in", file_name, begin=begin, step=step
         )
+        if stress is not None:
+            # Compute virials of all frames by: virial = stress * volume
+            virial = stress * np.linalg.det(data["cells"])[:, np.newaxis, np.newaxis]
+            data["virials"] = virial
         assert cs == es, "the step key between files are not consistent"
         return data
 

dpdata/qe/traj.py

@@ -1,4 +1,5 @@
 #!/usr/bin/python3
+import os
 import warnings
 
 import numpy as np
@@ -16,6 +17,7 @@
 length_convert = LengthConversion("bohr", "angstrom").value()
 energy_convert = EnergyConversion("hartree", "eV").value()
 force_convert = ForceConversion("hartree/bohr", "eV/angstrom").value()
+stress_convert = PressureConversion("GPa", "eV/angstrom^3").value()
 
 
 def load_key(lines, key):
@@ -233,8 +235,15 @@
         step=step,
         convert=force_convert,
     )
-    assert esteps == fsteps, "the step key between files are not consistent "
-    return energy, force, esteps
+    # Load stress from .str file if it exists
+    if os.path.isfile(prefix + ".str"):
+        stress, vsteps = load_data(
+            prefix + ".str", 3, begin=begin, step=step, convert=stress_convert
+        )
+    else:
+        stress, vsteps = None, esteps
+    assert esteps == fsteps == vsteps, "the step key between files are not consistent "
+    return energy, force, stress, esteps
 
 
 if __name__ == "__main__":