fix(abacus): put pp/orb/dpks files to folder pporb in 02fp.

dpgen/generator/lib/abacus_scf.py

@@ -25,8 +25,9 @@
     return ret
 
 
-def make_abacus_scf_input(fp_params):
+def make_abacus_scf_input(fp_params, extra_file_path=""):
     # Make INPUT file for abacus pw scf calculation.
+    # put extra files (such as: deepks_model) to extra_file_path folder
     ret = "INPUT_PARAMETERS\n"
     ret += "calculation scf\n"
     for key in fp_params:
@@ -181,7 +182,9 @@
             ), "'deepks_scf' should be either 0 or 1."
             ret += "deepks_scf %d\n" % fp_params["deepks_scf"]
         elif key == "deepks_model":
-            ret += "deepks_model %s\n" % fp_params["deepks_model"]
+            ret += "deepks_model %s\n" % os.path.join(
+                extra_file_path, os.path.split(fp_params["deepks_model"])[1]
+            )
         elif key[0] == "_":
             pass
         elif key == "calculation":
@@ -198,6 +201,7 @@
     fp_dpks_descriptor=None,
     fp_params=None,
     type_map=None,
+    pporb="",  # pull all pp orb dpks files to pporb folder
 ):
     atom_names = sys_data["atom_names"]
     atom_numbs = sys_data["atom_numbs"]
@@ -217,12 +221,17 @@
         )
         idx = type_map.index(atom_names[iatom])
         if "atom_masses" not in sys_data:
-            ret += atom_names[iatom] + " 1.00 " + fp_pp_files[idx] + "\n"
+            ret += (
+                atom_names[iatom]
+                + " 1.00 "
+                + os.path.join(pporb, fp_pp_files[idx])
+                + "\n"
+            )
         else:
             ret += (
                 atom_names[iatom]
                 + " %.3f " % sys_data["atom_masses"][iatom]
-                + fp_pp_files[idx]
+                + os.path.join(pporb, fp_pp_files[idx])
                 + "\n"
             )
 
@@ -267,11 +276,11 @@
         assert len(fp_orb_files) == len(type_map)
         for iatom in range(len(atom_names)):
             idx = type_map.index(atom_names[iatom])
-            ret += fp_orb_files[idx] + "\n"
+            ret += os.path.join(pporb, fp_orb_files[idx]) + "\n"
 
     if fp_dpks_descriptor is not None:
         ret += "\nNUMERICAL_DESCRIPTOR\n"
-        ret += "%s\n" % fp_dpks_descriptor
+        ret += os.path.join(pporb, fp_dpks_descriptor) + "\n"
 
     return ret
 

dpgen/generator/run.py

@@ -25,6 +25,7 @@
 import warnings
 from collections import Counter
 from collections.abc import Iterable
+from pprint import pp
 from typing import List
 
 import dpdata
@@ -3031,6 +3032,7 @@
 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"])
+    type_map = jdata["type_map"]
 
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
@@ -3047,46 +3049,56 @@
         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)
+            model_file = os.path.join(
+                fp_pp_path, os.path.split(fp_dpks_model)[1]
+            )  # only the filename
             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)
+            os.symlink(model_file, fp_dpks_model)  # link to the model file
 
         # 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)
+        atom_names = stru_param["atom_names"]
+        pp_files_stru = stru_param.get("pp_files", None)
+        orb_files_stru = stru_param.get("orb_files", None)
+        descriptor_file_stru = stru_param.get("dpks_descriptor", None)
+
+        if pp_files_stru:
+            assert "fp_pp_files" in jdata, "need to define fp_pp_files in jdata"
+        if orb_files_stru:
+            assert "fp_orb_files" in jdata, "need to define fp_orb_files in jdata"
+        if descriptor_file_stru:
+            assert (
+                "fp_dpks_descriptor" in jdata
+            ), "need to define fp_dpks_descriptor in jdata"
 
-        for jj in orb_files:
-            ifile = os.path.join(fp_pp_path, jj)
+        for idx, iatom in enumerate(atom_names):
+            type_map_idx = type_map.index(iatom)
+            if iatom not in type_map:
+                raise RuntimeError(
+                    "atom name %s in STRU is not defined in type_map" % (iatom)
+                )
+            if pp_files_stru:
+                src_file = os.path.join(fp_pp_path, jdata["fp_pp_files"][type_map_idx])
+                assert os.path.isfile(
+                    src_file
+                ), f"Can not find the pseudopotential file {src_file}"
+                os.symlink(src_file, pp_files_stru[idx])
+            if orb_files_stru:
+                src_file = os.path.join(fp_pp_path, jdata["fp_orb_files"][type_map_idx])
+                assert os.path.isfile(
+                    src_file
+                ), f"Can not find the orbital file {src_file}"
+                os.symlink(src_file, orb_files_stru[idx])
+        if descriptor_file_stru:
+            src_file = os.path.join(fp_pp_path, jdata["fp_dpks_descriptor"])
             assert os.path.isfile(
-                ifile
-            ), "Can not find the orbital file %s, which is defined in %s/STRU" % (
-                ifile,
-                ii,
-            )
-            os.symlink(ifile, jj)
+                src_file
+            ), f"Can not find the descriptor file {src_file}"
+            os.symlink(src_file, descriptor_file_stru)
 
-        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)
 
 
@@ -3201,6 +3213,7 @@
 
 def make_fp_abacus_scf(iter_index, jdata):
     # make config
+    pporb_path = "pporb"
     fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
     if len(fp_tasks) == 0:
         return
@@ -3223,7 +3236,7 @@
         raise RuntimeError(
             "Set 'user_fp_params' or 'fp_incar' in json file to make INPUT of ABACUS"
         )
-    ret_input = make_abacus_scf_input(fp_params)
+    ret_input = make_abacus_scf_input(fp_params, extra_file_path=pporb_path)
 
     # Get orbital and deepks setting
     if "basis_type" in fp_params:
@@ -3295,10 +3308,14 @@
             fp_dpks_descriptor,
             fp_params,
             type_map=jdata["type_map"],
+            pporb=pporb_path,
         )
         with open("STRU", "w") as fp:
             fp.write(ret_stru)
 
+        if not os.path.isdir(pporb_path):
+            os.makedirs(pporb_path)
+
         os.chdir(cwd)
     # link pp and orbital files
     _link_fp_abacus_pporb_descript(iter_index, jdata)
@@ -3756,17 +3773,9 @@
             assert os.path.exists(fp_input_path)
             fp_input_path = os.path.abspath(fp_input_path)
             fp_params = get_abacus_input_parameters(fp_input_path)
-        forward_files = ["INPUT", "STRU"]
+        forward_files = ["INPUT", "STRU", "pporb"]
         if "kspacing" not in fp_params.keys():
-            forward_files = ["INPUT", "STRU", "KPT"]
-        forward_files += fp_pp_files
-        if "fp_orb_files" in jdata:
-            forward_files += jdata["fp_orb_files"]
-        if "fp_dpks_descriptor" in jdata:
-            forward_files.append(jdata["fp_dpks_descriptor"])
-        if "user_fp_params" in jdata:
-            if "deepks_model" in jdata["user_fp_params"]:
-                forward_files.append(jdata["user_fp_params"]["deepks_model"])
+            forward_files.append("KPT")
         backward_files = ["output", "OUT.ABACUS"]
         run_fp_inner(
             iter_index,

tests/generator/test_make_fp.py

@@ -179,7 +179,7 @@
 deepks_out_labels 0\n\
 deepks_descriptor_lmax 0\n\
 deepks_scf 0\n\
-deepks_model model.ptg\n"
+deepks_model pporb/model.ptg\n"
 
 abacus_kpt_ref = "K_POINTS\n\
 0\n\
@@ -690,6 +690,20 @@ def test_make_fp_pwscf_old(self):
 
 
 class TestMakeFPABACUS(unittest.TestCase):
+    def _check_pp(self, idx, fp_pp_path, fp_pp_files):
+        nfile = len(fp_pp_files)
+        fp_path = os.path.join("iter.%06d" % idx, "02.fp")
+        for ii in range(nfile):
+            self.assertTrue(os.path.isfile(os.path.join(fp_pp_path, fp_pp_files[ii])))
+        tasks = glob.glob(os.path.join(fp_path, "task.*"))
+        for ii in tasks:
+            for jj in range(nfile):
+                my_file_cmp(
+                    self,
+                    os.path.join(fp_pp_path, fp_pp_files[jj]),
+                    os.path.join(ii, "pporb", fp_pp_files[jj]),
+                )
+
     def test_make_fp_abacus(self):
         setUpModule()
         if os.path.isdir("iter.000000"):
@@ -723,7 +737,7 @@ def test_make_fp_abacus(self):
         _check_poscars(self, 0, jdata["fp_task_max"], jdata["type_map"])
         _check_abacus_input(self, 0)
         _check_abacus_kpt(self, 0)
-        _check_potcar(self, 0, jdata["fp_pp_path"], jdata["fp_pp_files"])
+        self._check_pp(0, jdata["fp_pp_path"], jdata["fp_pp_files"])
         shutil.rmtree("iter.000000")
 
     def test_make_fp_abacus_kspacing(self):
@@ -803,7 +817,7 @@ def test_make_fp_abacus_from_input(self):
         _check_poscars(self, 0, jdata["fp_task_max"], jdata["type_map"])
         _check_abacus_input(self, 0)
         _check_abacus_kpt(self, 0)
-        _check_potcar(self, 0, jdata["fp_pp_path"], jdata["fp_pp_files"])
+        self._check_pp(0, jdata["fp_pp_path"], jdata["fp_pp_files"])
         shutil.rmtree("iter.000000")