support virial in qe/traj

Author: Han Wang

dpdata/qe/traj.py

@@ -11,6 +11,8 @@
 if TYPE_CHECKING:
     from dpdata.utils import FileType
 
+import os
+
 from ..unit import (
     EnergyConversion,
     ForceConversion,
@@ -20,6 +22,7 @@
 
 ry2ev = EnergyConversion("rydberg", "eV").value()
 kbar2evperang3 = PressureConversion("kbar", "eV/angstrom^3").value()
+gpa2evperbohr = PressureConversion("GPa", "eV/bohr^3").value()
 
 length_convert = LengthConversion("bohr", "angstrom").value()
 energy_convert = EnergyConversion("hartree", "eV").value()
@@ -228,6 +231,32 @@ def to_system_data(input_name, prefix, begin=0, step=1):
             )
     except FileNotFoundError:
         data["cells"] = np.tile(cell, (data["coords"].shape[0], 1, 1))
+
+    # handle virial
+    stress_fname = prefix + ".str"
+    if os.path.exists(stress_fname):
+        # 1. Read stress tensor (in GPa) for each structure
+        stress, vsteps = load_data(stress_fname, 3, begin=begin, step=step, convert=1.0)
+        if csteps != vsteps:
+            csteps.append(None)
+            vsteps.append(None)
+            for int_id in range(len(csteps)):
+                if csteps[int_id] != vsteps[int_id]:
+                    break
+            step_id = begin + int_id * step
+            raise RuntimeError(
+                f"the step key between files are not consistent. "
+                f"The difference locates at step: {step_id}, "
+                f".pos is {csteps[int_id]}, .str is {vsteps[int_id]}"
+            )
+        # 2. Calculate volume from cell. revert unit to bohr before taking det
+        volumes = np.linalg.det(data["cells"] / length_convert).reshape(-1)
+        # 3. Calculate virials for each structure
+        virials = []
+        for i in range(stress.shape[0]):
+            virials.append(stress[i] * gpa2evperbohr * volumes[i])
+        data["virials"] = np.array(virials)  # shape: [nframe, 3, 3]
+
     return data, csteps
 
 

tests/test_qe_cp_traj.py

@@ -68,5 +68,40 @@ def test_case_null(self):
                 self.assertAlmostEqual(cell[ii][jj], ref[ii][jj])
 
 
+class TestVirial(unittest.TestCase):
+    def test(self):
+        self.system = dpdata.LabeledSystem("qe.traj/si/si", fmt="qe/cp/traj")
+        self.assertEqual(self.system["virials"].shape, (2, 3, 3))
+        np.testing.assert_almost_equal(
+            self.system["virials"][0],
+            np.array(
+                [
+                    [0.31120718, -0.03261485, -0.02537362],
+                    [-0.03261485, 0.3100397, 0.04211053],
+                    [-0.02537362, 0.04211057, 0.30571264],
+                ]
+            ),
+        )
+        np.testing.assert_almost_equal(
+            self.system["virials"][1],
+            np.array(
+                [
+                    [0.31072979, -0.03151186, -0.02302297],
+                    [-0.03151186, 0.30951293, 0.04078447],
+                    [-0.02302297, 0.04078451, 0.30544987],
+                ]
+            ),
+        )
+
+    def test_raise(self):
+        with self.assertRaises(RuntimeError) as c:
+            self.system = dpdata.LabeledSystem(
+                "qe.traj/si.wrongstr/si", fmt="qe/cp/traj"
+            )
+            self.assertTrue(
+                "the step key between files are not consistent." in str(c.exception)
+            )
+
+
 if __name__ == "__main__":
     unittest.main()