Improve function pert_scaled (#1563)

My suggestion:
- Separate VASP and ABACUS into separated functions (and may be
separated files). This will helpful to someone who want to implement a
new DFT engine and also convenient for maintain.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Enhanced the `pert_scaled` function to support different
initialization styles ("VASP" or "ABACUS").
  
- **Refactor**
- Improved loop logic and file handling for better performance and
reliability in perturbation processes.
- Streamlined the `make_vasp_relax` and `make_scale` functions for
improved readability and clarity.
- Simplified the construction of perturbation commands in the
`pert_scaled` function.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

dpgen/data/gen.py

@@ -728,8 +728,7 @@ def make_scale_ABACUS(jdata):
 
 
 def pert_scaled(jdata):
-    if "init_fp_style" not in jdata:
-        jdata["init_fp_style"] = "VASP"
+    ### Extract data from jdata
     out_dir = jdata["out_dir"]
     scale = jdata["scale"]
     pert_box = jdata["pert_box"]
@@ -746,6 +745,7 @@ def pert_scaled(jdata):
     if "from_poscar" in jdata:
         from_poscar = jdata["from_poscar"]
 
+    ### Get the current working directory and the system path
     cwd = os.getcwd()
     path_sp = os.path.join(out_dir, global_dirname_03)
     assert os.path.isdir(path_sp)
@@ -754,35 +754,35 @@ def pert_scaled(jdata):
     sys_pe.sort()
     os.chdir(cwd)
 
-    pert_cmd = os.path.dirname(__file__)
-    pert_cmd = os.path.join(pert_cmd, "tools")
-    pert_cmd = os.path.join(pert_cmd, "create_random_disturb.py")
-    fp_style = "vasp"
-    poscar_name = "POSCAR"
-    if jdata["init_fp_style"] == "ABACUS":
+    ### Construct the perturbation command
+    init_fp_style = jdata.get("init_fp_style", "VASP")
+    if init_fp_style == "VASP":
+        fp_style = "vasp"
+        poscar_name = "POSCAR"
+    elif init_fp_style == "ABACUS":
         fp_style = "abacus"
         poscar_name = "STRU"
-    pert_cmd = (
-        sys.executable
-        + " "
-        + pert_cmd
-        + " -etmax %f -ofmt %s %s %d %f > /dev/null"
-        % (pert_box, fp_style, poscar_name, pert_numb, pert_atom)
+
+    python_exec = os.path.join(
+        os.path.dirname(__file__), "tools", "create_random_disturb.py"
     )
+    pert_cmd = f"{sys.executable} {python_exec} -etmax {pert_box} -ofmt {fp_style} {poscar_name} {pert_numb} {pert_atom} > /dev/null"
+
+    ### Loop over each system and scale
     for ii in sys_pe:
         for jj in scale:
-            path_work = path_sp
-            path_work = os.path.join(path_work, ii)
-            path_work = os.path.join(path_work, f"scale-{jj:.3f}")
+            path_work = os.path.join(path_sp, ii, f"scale-{jj:.3f}")
             assert os.path.isdir(path_work)
             os.chdir(path_work)
             sp.check_call(pert_cmd, shell=True)
+
+            ### Loop over each perturbation
             for kk in range(pert_numb):
                 if fp_style == "vasp":
-                    pos_in = "POSCAR%d.vasp" % (kk + 1)
+                    pos_in = f"POSCAR{kk+1}.vasp"
                 elif fp_style == "abacus":
-                    pos_in = "STRU%d.abacus" % (kk + 1)
-                dir_out = "%06d" % (kk + 1)
+                    pos_in = f"STRU{kk+1}.abacus"
+                dir_out = f"{kk+1:06d}"
                 create_path(dir_out)
                 if fp_style == "vasp":
                     pos_out = os.path.join(dir_out, "POSCAR")
@@ -807,12 +807,14 @@ def pert_scaled(jdata):
                 else:
                     shutil.copy2(pos_in, pos_out)
                 os.remove(pos_in)
+
+            ### Handle special case (unperturbed ?)
             kk = -1
             if fp_style == "vasp":
                 pos_in = "POSCAR"
             elif fp_style == "abacus":
                 pos_in = "STRU"
-            dir_out = "%06d" % (kk + 1)
+            dir_out = f"{kk+1:06d}"
             create_path(dir_out)
             if fp_style == "vasp":
                 pos_out = os.path.join(dir_out, "POSCAR")
@@ -836,6 +838,7 @@ def pert_scaled(jdata):
                         )
             else:
                 shutil.copy2(pos_in, pos_out)
+
             os.chdir(cwd)
 
 

dpgen/data/surf.py

@@ -18,7 +18,7 @@
 import dpgen.data.tools.fcc as fcc
 import dpgen.data.tools.hcp as hcp
 import dpgen.data.tools.sc as sc
-from dpgen import ROOT_PATH, dlog
+from dpgen import dlog
 from dpgen.dispatcher.Dispatcher import make_submission_compat
 from dpgen.generator.lib.utils import symlink_user_forward_files
 from dpgen.remote.decide_machine import convert_mdata
@@ -354,15 +354,16 @@ def make_vasp_relax(jdata):
     out_dir = jdata["out_dir"]
     potcars = jdata["potcars"]
     cwd = os.getcwd()
-
     work_dir = os.path.join(out_dir, global_dirname_02)
     assert os.path.isdir(work_dir)
     work_dir = os.path.abspath(work_dir)
+
     if os.path.isfile(os.path.join(work_dir, "INCAR")):
         os.remove(os.path.join(work_dir, "INCAR"))
     if os.path.isfile(os.path.join(work_dir, "POTCAR")):
         os.remove(os.path.join(work_dir, "POTCAR"))
     shutil.copy2(jdata["relax_incar"], os.path.join(work_dir, "INCAR"))
+
     out_potcar = os.path.join(work_dir, "POTCAR")
     with open(out_potcar, "w") as outfile:
         for fname in potcars:
@@ -442,15 +443,12 @@ def make_scale(jdata):
         for jj in scale:
             if skip_relax:
                 pos_src = os.path.join(os.path.join(init_path, ii), "POSCAR")
-                assert os.path.isfile(pos_src)
             else:
-                try:
-                    pos_src = os.path.join(os.path.join(init_path, ii), "CONTCAR")
-                    assert os.path.isfile(pos_src)
-                except Exception:
-                    raise RuntimeError(
-                        "not file %s, vasp relaxation should be run before scale poscar"
-                    )
+                pos_src = os.path.join(os.path.join(init_path, ii), "CONTCAR")
+            if not os.path.isfile(pos_src):
+                raise RuntimeError(
+                    f"file {pos_src} not found, vasp relaxation should be run before scale poscar"
+                )
             scale_path = os.path.join(work_path, ii)
             scale_path = os.path.join(scale_path, f"scale-{jj:.3f}")
             create_path(scale_path)
@@ -503,46 +501,45 @@ def pert_scaled(jdata):
     sys_pe.sort()
     os.chdir(cwd)
 
-    pert_cmd = (
-        sys.executable
-        + " "
-        + os.path.join(ROOT_PATH, "data/tools/create_random_disturb.py")
-    )
-    pert_cmd += " -etmax %f -ofmt vasp POSCAR %d %f > /dev/null" % (
-        pert_box,
-        pert_numb,
-        pert_atom,
+    ### Construct the perturbation command
+    python_exec = os.path.join(
+        os.path.dirname(__file__), "tools", "create_random_disturb.py"
     )
+    pert_cmd = f"{sys.executable} {python_exec} -etmax {pert_box} -ofmt vasp POSCAR {pert_numb} {pert_atom} > /dev/null"
+
+    ### Loop over each system and scale
     for ii in sys_pe:
         for jj in scale:
-            path_scale = path_sp
-            path_scale = os.path.join(path_scale, ii)
-            path_scale = os.path.join(path_scale, f"scale-{jj:.3f}")
+            path_scale = os.path.join(path_sp, ii, f"scale-{jj:.3f}")
             assert os.path.isdir(path_scale)
             os.chdir(path_scale)
             dlog.info(os.getcwd())
             poscar_in = os.path.join(path_scale, "POSCAR")
             assert os.path.isfile(poscar_in)
+
+            ### Loop over each perturbation
             for ll in elongs:
-                path_elong = path_scale
-                path_elong = os.path.join(path_elong, f"elong-{ll:3.3f}")
+                path_elong = os.path.join(path_scale, f"elong-{ll:3.3f}")
                 create_path(path_elong)
                 os.chdir(path_elong)
                 poscar_elong(poscar_in, "POSCAR", ll)
                 sp.check_call(pert_cmd, shell=True)
                 for kk in range(pert_numb):
-                    pos_in = "POSCAR%d.vasp" % (kk + 1)
-                    dir_out = "%06d" % (kk + 1)
+                    pos_in = f"POSCAR{kk+1}.vasp"
+                    dir_out = f"{kk+1:06d}"
                     create_path(dir_out)
                     pos_out = os.path.join(dir_out, "POSCAR")
                     poscar_shuffle(pos_in, pos_out)
                     os.remove(pos_in)
+
+                ### Handle special case (unperturbed ?)
                 kk = -1
                 pos_in = "POSCAR"
-                dir_out = "%06d" % (kk + 1)
+                dir_out = f"{kk+1:06d}"
                 create_path(dir_out)
                 pos_out = os.path.join(dir_out, "POSCAR")
                 poscar_shuffle(pos_in, pos_out)
+
                 os.chdir(cwd)