refactor(abacus): link the pp/orb files defined in STRU not jdata (#1127)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: pxlxingliang

dpgen/generator/lib/abacus_scf.py

@@ -342,9 +342,9 @@ def get_additional_from_STRU(geometry_inlines, nele):
             if orb_file_kw == geometry_inlines[iline].split()[0]:
                 orb_file = []
                 for iele in range(nele):
-                    orb_file.append(geometry_inlines[iline + iele + 1].rstrip())
+                    orb_file.append(geometry_inlines[iline + iele + 1].strip())
             if dpks_descriptor_kw == geometry_inlines[iline].split()[0]:
-                dpks_descriptor = geometry_inlines[iline + 1].rstrip()
+                dpks_descriptor = geometry_inlines[iline + 1].strip()
     return orb_file, dpks_descriptor
 
 

dpgen/generator/run.py

@@ -68,7 +68,7 @@
     make_abacus_scf_input,
     make_abacus_scf_kpt,
 )
-from dpgen.generator.lib.abacus_scf import get_abacus_input_parameters
+from dpgen.generator.lib.abacus_scf import get_abacus_input_parameters, get_abacus_STRU
 
 # from dpgen.generator.lib.pwscf import cvt_1frame
 from dpgen.generator.lib.pwmat import make_pwmat_input_dict
@@ -3000,41 +3000,65 @@ def sys_link_fp_vasp_pp(iter_index, jdata):
             os.chdir(cwd)
 
 
-def _link_fp_abacus_orb_descript(iter_index, jdata):
-    # assume lcao orbital files, numerical descrptors and model for dpks are all in fp_pp_path.
-    fp_pp_path = jdata["fp_pp_path"]
-
-    fp_orb_files = jdata["fp_orb_files"]
-    assert os.path.exists(fp_pp_path)
-    fp_dpks_descriptor = None
-    fp_dpks_model = None
-    if "fp_dpks_descriptor" in jdata:
-        fp_dpks_descriptor = jdata["fp_dpks_descriptor"]
-    if "user_fp_params" in jdata:
-        if "deepks_model" in jdata["user_fp_params"]:
-            fp_dpks_model = jdata["user_fp_params"]["deepks_model"]
-
-    fp_pp_path = os.path.abspath(fp_pp_path)
+def _link_fp_abacus_pporb_descript(iter_index, jdata):
+    # assume pp orbital files, numerical descrptors and model for dpks are all in fp_pp_path.
+    fp_pp_path = os.path.abspath(jdata["fp_pp_path"])
 
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
-
     fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     fp_tasks.sort()
     if len(fp_tasks) == 0:
         return
+
     cwd = os.getcwd()
     for ii in fp_tasks:
         os.chdir(ii)
-        for jj in fp_orb_files:
-            orb_file = os.path.join(fp_pp_path, jj)
-            os.symlink(orb_file, jj)
-        if fp_dpks_descriptor is not None:
-            descrptor = os.path.join(fp_pp_path, fp_dpks_descriptor)
-            os.symlink(descrptor, fp_dpks_descriptor)
-        if fp_dpks_model is not None:
-            dpks_model = os.path.join(fp_pp_path, fp_dpks_model)
-            os.symlink(dpks_model, fp_dpks_model)
+
+        # get value of 'deepks_model' from INPUT
+        input_param = get_abacus_input_parameters("INPUT")
+        fp_dpks_model = input_param.get("deepks_model", None)
+        if fp_dpks_model != None:
+            model_file = os.path.join(fp_pp_path, fp_dpks_model)
+            assert os.path.isfile(model_file), (
+                "Can not find the deepks model file %s, which is defined in %s/INPUT"
+                % (model_file, ii)
+            )
+            os.symlink(model_file, fp_dpks_model)
+
+        # get pp, orb, descriptor filenames from STRU
+        stru_param = get_abacus_STRU("STRU")
+        pp_files = stru_param.get("pp_files", [])
+        orb_files = stru_param.get("orb_files", [])
+        descriptor_file = stru_param.get("dpks_descriptor", None)
+        pp_files = [] if pp_files == None else pp_files
+        orb_files = [] if orb_files == None else orb_files
+
+        for jj in pp_files:
+            ifile = os.path.join(fp_pp_path, jj)
+            assert os.path.isfile(ifile), (
+                "Can not find the pseudopotential file %s, which is defined in %s/STRU"
+                % (ifile, ii)
+            )
+            os.symlink(ifile, jj)
+
+        for jj in orb_files:
+            ifile = os.path.join(fp_pp_path, jj)
+            assert os.path.isfile(
+                ifile
+            ), "Can not find the orbital file %s, which is defined in %s/STRU" % (
+                ifile,
+                ii,
+            )
+            os.symlink(ifile, jj)
+
+        if descriptor_file != None:
+            ifile = os.path.join(fp_pp_path, descriptor_file)
+            assert os.path.isfile(ifile), (
+                "Can not find the deepks descriptor file %s, which is defined in %s/STRU"
+                % (ifile, ii)
+            )
+            os.symlink(ifile, descriptor_file)
         os.chdir(cwd)
 
 
@@ -3248,11 +3272,8 @@ def make_fp_abacus_scf(iter_index, jdata):
             fp.write(ret_stru)
 
         os.chdir(cwd)
-    # link pp files
-    _link_fp_vasp_pp(iter_index, jdata)
-    if "basis_type" in fp_params:
-        if fp_params["basis_type"] == "lcao":
-            _link_fp_abacus_orb_descript(iter_index, jdata)
+    # link pp and orbital files
+    _link_fp_abacus_pporb_descript(iter_index, jdata)
 
 
 def make_fp_siesta(iter_index, jdata):