diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py index 6a98eedd..8fa07d1c 100644 --- a/dpdata/plugins/qe.py +++ b/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 @@ 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..05e401ab 100644 --- a/dpdata/qe/traj.py +++ b/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 @@ 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__":