fix: qe input file not seperated by empty lines

wanghan-iapcm · Sep 19, 2024 · 345aedd · 345aedd
1 parent 480242e
commit 345aedd
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 3 deletions.
diff --git a/dpdata/qe/scf.py b/dpdata/qe/scf.py
@@ -11,15 +11,32 @@
 bohr2ang = 0.52917721067
 kbar2evperang3 = 1e3 / 1.602176621e6
 
+_QE_BLOCK_KEYWORDS = [
+    "ATOMIC_SPECIES",
+    "ATOMIC_POSITIONS",
+    "K_POINTS",
+    "ADDITIONAL_K_POINTS",
+    "CELL_PARAMETERS",
+    "CONSTRAINTS",
+    "OCCUPATIONS",
+    "ATOMIC_VELOCITIES",
+    "ATOMIC_FORCES",
+    "SOLVENTS",
+    "HUBBARD",
+]
+
 
 def get_block(lines, keyword, skip=0):
     ret = []
     for idx, ii in enumerate(lines):
         if keyword in ii:
             blk_idx = idx + 1 + skip
-            while len(lines[blk_idx]) == 0:
+            while len(lines[blk_idx].split()) == 0:
                 blk_idx += 1
-            while len(lines[blk_idx]) != 0 and blk_idx != len(lines):
+            while (
+                len(lines[blk_idx].split()) != 0
+                and (lines[blk_idx].split()[0] not in _QE_BLOCK_KEYWORDS)
+            ) and blk_idx != len(lines):
                 ret.append(lines[blk_idx])
                 blk_idx += 1
             break
@@ -111,6 +128,8 @@ def get_energy(lines):
 
 def get_force(lines, natoms):
     blk = get_block(lines, "Forces acting on atoms", skip=1)
+    if len(blk) == 0:
+        raise RuntimeError("no force found in the output file.")
     ret = []
     blk = blk[0 : sum(natoms)]
     for ii in blk:
@@ -123,7 +142,7 @@ def get_force(lines, natoms):
 def get_stress(lines):
     blk = get_block(lines, "total   stress")
     if len(blk) == 0:
-        return
+        return None
     ret = []
     for ii in blk:
         ret.append([float(jj) for jj in ii.split()[3:6]])

diff --git a/tests/test_qe_pw_scf.py b/tests/test_qe_pw_scf.py
@@ -188,5 +188,50 @@ def setUp(self):
         )
 
 
+class TestNa(unittest.TestCase):
+    def test(self):
+        ss = dpdata.LabeledSystem("qe.scf/na.out", fmt="qe/pw/scf")
+        b2a = dpdata.unit.LengthConversion("bohr", "angstrom").value()
+        print(ss.data.keys())
+        self.assertEqual(ss["atom_numbs"], [3])
+        self.assertEqual(ss["atom_names"], ["Na"])
+        self.assertEqual(ss.get_nframes(), 1)
+        np.testing.assert_array_equal(ss["atom_types"], [0, 0, 0])
+        np.testing.assert_allclose(
+            ss["coords"][0],
+            np.array(
+                [
+                    0.940587444301534,
+                    0.397635863676890,
+                    0.059472381901808,
+                    0.059413515648061,
+                    0.602364552456546,
+                    0.559472465518034,
+                    0.602364619812068,
+                    0.059413062489401,
+                    0.059472381901808,
+                ]
+            ).reshape(3, 3)
+            @ ss["cells"][0],
+        )
+        np.testing.assert_allclose(
+            ss["cells"][0],
+            np.array(
+                [
+                    7.171683039200000,
+                    0.000000000000000,
+                    0.000000000000000,
+                    -4.270578118300000,
+                    5.761527588200000,
+                    0.000000000000000,
+                    -0.000000045600000,
+                    0.000000023000000,
+                    12.826457854999999,
+                ]
+            ).reshape(3, 3),
+        )
+        np.testing.assert_allclose(ss["forces"][0], np.zeros([3, 3]))
+
+
 if __name__ == "__main__":
     unittest.main()