abacus: update the read of stress and force for ABACUS v3.4.1 (#560)

Signed-off-by: Peng Xingliang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

dpdata/abacus/relax.py

@@ -2,7 +2,15 @@
 
 import numpy as np
 
-from .scf import bohr2ang, get_cell, get_coords, get_geometry_in, kbar2evperang3
+from .scf import (
+    bohr2ang,
+    collect_force,
+    collect_stress,
+    get_cell,
+    get_coords,
+    get_geometry_in,
+    kbar2evperang3,
+)
 
 # Read in geometries from an ABACUS RELAX(CELL-RELAX) trajectory in OUT.XXXX/runnning_relax/cell-relax.log.
 
@@ -40,8 +48,6 @@ def get_coords_from_log(loglines, natoms):
     energy = []
     cells = []
     coords = []
-    force = []
-    stress = []
     coord_direct = []  # if the coordinate is direct type or not
 
     for i in range(len(loglines)):
@@ -91,18 +97,8 @@ def get_coords_from_log(loglines, natoms):
             # get the energy of current structure
             energy.append(float(line.split()[-2]))
 
-        elif line[4:15] == "TOTAL-FORCE":
-            force.append([])
-            for j in range(5, 5 + natoms):
-                force[-1].append(
-                    list(map(lambda x: float(x), loglines[i + j].split()[1:4]))
-                )
-        elif line[1:13] == "TOTAL-STRESS":
-            stress.append([])
-            for j in range(4, 7):
-                stress[-1].append(
-                    list(map(lambda x: float(x), loglines[i + j].split()[0:3]))
-                )
+    force = collect_force(loglines)
+    stress = collect_stress(loglines)
 
     # delete last structures which has no energy
     while len(energy) < len(coords):

dpdata/abacus/scf.py

@@ -1,5 +1,6 @@
 import os
 import re
+import warnings
 
 import numpy as np
 
@@ -123,33 +124,79 @@ def get_energy(outlines):
     return Etot, False
 
 
-def get_force(outlines, natoms):
+def collect_force(outlines):
     force = []
-    force_inlines = get_block(
-        outlines, "TOTAL-FORCE (eV/Angstrom)", skip=4, nlines=np.sum(natoms)
-    )
-    if force_inlines is None:
-        print(
-            "TOTAL-FORCE (eV/Angstrom) is not found in OUT.XXX/running_scf.log. May be you haven't set 'cal_force 1' in the INPUT."
-        )
+    for i, line in enumerate(outlines):
+        if "TOTAL-FORCE (eV/Angstrom)" in line:
+            value_pattern = re.compile(
+                r"^\s*[A-Z][a-z]?[1-9][0-9]*\s+[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?\s+[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?\s+[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?\s*$"
+            )
+            j = i
+            # find the first line of force
+            noforce = False
+            while not value_pattern.match(outlines[j]):
+                j += 1
+                if (
+                    j >= i + 10
+                ):  # if can not find the first line of force in 10 lines, then stop
+                    warnings.warn("Warning: can not find the first line of force")
+                    noforce = True
+                    break
+            if noforce:
+                break
+
+            force.append([])
+            while value_pattern.match(outlines[j]):
+                force[-1].append([float(ii) for ii in outlines[j].split()[1:4]])
+                j += 1
+    return force  # only return the last force
+
+
+def get_force(outlines, natoms):
+    force = collect_force(outlines)
+    if len(force) == 0:
         return [[]]
-    for line in force_inlines:
-        force.append([float(f) for f in line.split()[1:4]])
-    force = np.array(force)
-    return force
+    else:
+        return np.array(force[-1])  # only return the last force
 
 
-def get_stress(outlines):
+def collect_stress(outlines):
     stress = []
-    stress_inlines = get_block(outlines, "TOTAL-STRESS (KBAR)", skip=3, nlines=3)
-    if stress_inlines is None:
-        return None
-    for line in stress_inlines:
-        stress.append([float(f) for f in line.split()])
-    stress = np.array(stress) * kbar2evperang3
+    for i, line in enumerate(outlines):
+        if "TOTAL-STRESS (KBAR)" in line:
+            value_pattern = re.compile(
+                r"^\s*[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?\s+[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?\s+[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?\s*$"
+            )
+            j = i
+            nostress = False
+            while not value_pattern.match(outlines[j]):
+                j += 1
+                if (
+                    j >= i + 10
+                ):  # if can not find the first line of stress in 10 lines, then stop
+                    warnings.warn("Warning: can not find the first line of stress")
+                    nostress = True
+                    break
+            if nostress:
+                break
+
+            stress.append([])
+            while value_pattern.match(outlines[j]):
+                stress[-1].append(
+                    list(map(lambda x: float(x), outlines[j].split()[0:3]))
+                )
+                j += 1
     return stress
 
 
+def get_stress(outlines):
+    stress = collect_stress(outlines)
+    if len(stress) == 0:
+        return None
+    else:
+        return np.array(stress[-1]) * kbar2evperang3  # only return the last stress
+
+
 def get_frame(fname):
     data = {
         "atom_names": [],