From 6ed8779596d4b24e1e59d31f2e286d7e130bf5ea Mon Sep 17 00:00:00 2001
From: Rashid Rafeek <rashidrafeek@gmail.com>
Date: Thu, 30 Nov 2023 23:58:35 +0530
Subject: [PATCH 1/2] Enable loading of virials for 'qe/cp/traj' if prefix.str
 file is present

---
 dpdata/plugins/qe.py |  7 ++++++-
 dpdata/qe/traj.py    | 13 +++++++++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py
index 6a98eedd..30ac2278 100644
--- a/dpdata/plugins/qe.py
+++ b/dpdata/plugins/qe.py
@@ -2,6 +2,7 @@
 import dpdata.qe.scf
 import dpdata.qe.traj
 from dpdata.format import Format
+import numpy as np
 
 
 @Format.register("qe/cp/traj")
@@ -26,9 +27,13 @@ def from_labeled_system(self, file_name, begin=0, step=1, **kwargs):
             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
 
diff --git a/dpdata/qe/traj.py b/dpdata/qe/traj.py
index e27990cb..1e2cb11b 100644
--- a/dpdata/qe/traj.py
+++ b/dpdata/qe/traj.py
@@ -1,5 +1,6 @@
 #!/usr/bin/python3
 import warnings
+import os
 
 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 @@ def to_system_label(input_name, prefix, begin=0, step=1):
         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__":

From dbe6699ea70ecdcb7edbd7a45a131ed3c6a6ef50 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 30 Nov 2023 18:34:29 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/plugins/qe.py | 3 ++-
 dpdata/qe/traj.py    | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py
index 30ac2278..8fa07d1c 100644
--- a/dpdata/plugins/qe.py
+++ b/dpdata/plugins/qe.py
@@ -1,8 +1,9 @@
+import numpy as np
+
 import dpdata.md.pbc
 import dpdata.qe.scf
 import dpdata.qe.traj
 from dpdata.format import Format
-import numpy as np
 
 
 @Format.register("qe/cp/traj")
diff --git a/dpdata/qe/traj.py b/dpdata/qe/traj.py
index 1e2cb11b..05e401ab 100644
--- a/dpdata/qe/traj.py
+++ b/dpdata/qe/traj.py
@@ -1,6 +1,6 @@
 #!/usr/bin/python3
-import warnings
 import os
+import warnings
 
 import numpy as np