From 4b69200cd3d861e77df3dfa94a7fa353e4db4b15 Mon Sep 17 00:00:00 2001
From: Peng Xingliang <91927439+pxlxingliang@users.noreply.github.com>
Date: Tue, 31 Oct 2023 09:07:06 +0800
Subject: [PATCH] abacus: fix bug of finding the final relax STRU (#1344)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 dpgen/auto_test/lib/abacus.py           | 31 +++++++++++++------------
 dpgen/data/gen.py                       |  9 ++++---
 tests/auto_test/test_abacus_property.py |  1 +
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/dpgen/auto_test/lib/abacus.py b/dpgen/auto_test/lib/abacus.py
index 6a90fea8f..34b53af62 100644
--- a/dpgen/auto_test/lib/abacus.py
+++ b/dpgen/auto_test/lib/abacus.py
@@ -1,4 +1,5 @@
 #!/usr/bin/python3
+import glob
 import os
 
 import dpdata
@@ -309,23 +310,23 @@ def final_stru(abacus_path):
             out_stru = bool(line.split()[1])
     logf = os.path.join(abacus_path, f"OUT.{suffix}/running_{calculation}.log")
     if calculation in ["relax", "cell-relax"]:
-        if not out_stru:
+        if os.path.isfile(os.path.join(abacus_path, "OUT.%s/STRU_ION_D" % suffix)):
             return "OUT.%s/STRU_ION_D" % suffix
         else:
-            with open(logf) as f1:
-                lines = f1.readlines()
-            for i in range(1, len(lines)):
-                max_step = ""
-                if "ALGORITHM --------------- ION=" in lines[-i]:
-                    index_ben = lines[-i].index("ION=") + 4
-                    index_end = lines[-i].index("ELEC")
-                    max_step = int(lines[-i][index_ben:index_end])
-                    if max_step < 2:
-                        max_step = ""
-                    else:
-                        max_step -= 2
-                    break
-            return f"OUT.{suffix}/STRU_ION{str(max_step)}_D"
+            # find the final name by STRU_ION*_D,
+            # for abacus version < v3.2.2, there has no STRU_ION_D file but has STRU_ION0_D STRU_ION1_D ... STRU_ION10_D ...
+            # so we need to find the last STRU_ION*_D file
+            stru_ions = glob.glob(
+                os.path.join(abacus_path, f"OUT.{suffix}/STRU_ION*_D")
+            )
+            if len(stru_ions) > 0:
+                # sort the file name by the number in the file name
+                stru_ions.sort(key=lambda x: int(x.split("_")[-2][3:]))
+                final_stru_ion = os.path.basename(stru_ions[-1])
+                return f"OUT.{suffix}/{final_stru_ion}"
+            else:
+                # if there has no STRU_ION_D, return the input STRU
+                return "STRU"
     elif calculation == "md":
         with open(logf) as f1:
             lines = f1.readlines()
diff --git a/dpgen/data/gen.py b/dpgen/data/gen.py
index 4a2527f6d..27134ef64 100644
--- a/dpgen/data/gen.py
+++ b/dpgen/data/gen.py
@@ -710,13 +710,16 @@ def make_scale_ABACUS(jdata):
                 assert os.path.isfile(pos_src)
             else:
                 try:
-                    pos_src = os.path.join(
-                        os.path.join(init_path, ii), "OUT.ABACUS/STRU_ION_D"
+                    from dpgen.auto_test.lib.abacus import (
+                        final_stru as abacus_final_stru,
                     )
+
+                    pos_src = abacus_final_stru(os.path.join(init_path, ii))
+                    pos_src = os.path.join(init_path, ii, pos_src)
                     assert os.path.isfile(pos_src)
                 except Exception:
                     raise RuntimeError(
-                        "not file %s, vasp relaxation should be run before scale poscar"
+                        "Can not find STRU_ION_D in OUT.ABACUS!!!\nABACUS relaxation should be run before scale poscar"
                     )
             scale_path = os.path.join(work_path, ii)
             scale_path = os.path.join(scale_path, "scale-%.3f" % jj)
diff --git a/tests/auto_test/test_abacus_property.py b/tests/auto_test/test_abacus_property.py
index edcc3985c..764092258 100644
--- a/tests/auto_test/test_abacus_property.py
+++ b/tests/auto_test/test_abacus_property.py
@@ -174,6 +174,7 @@ def test_make_property_elastic(self):
         os.remove(
             os.path.realpath(os.path.join(self.equi_path, "OUT.ABACUS", "STRU_ION_D"))
         )
+        os.remove(os.path.realpath(os.path.join(self.equi_path, "STRU")))
         with self.assertRaises(RuntimeError):
             elastic.make_confs(work_path, self.equi_path, refine=False)