split make_fp and share make_fp_calculation between run and simplify

dpgen/generator/run.py

@@ -3127,7 +3127,25 @@
         os.chdir(cwd)
 
 
-def _make_fp_vasp_configs(iter_index, jdata):
+def _make_fp_vasp_configs(iter_index: int, jdata: dict):
+    """Read the model deviation from model_devi step, and then generate the candidated structures
+    in 02.fp directory.
+
+    Currently, the formats of generated structures are decided by model_devi_eigne.
+
+    Parameters
+    ----------
+    iter_index : int
+        The index of iteration.
+    jdata : dict
+        The json data.
+
+    Returns
+    -------
+    int
+        The number of the candidated structures.
+    """
+    # TODO: we need to unify different data formats
     fp_task_max = jdata["fp_task_max"]
     model_devi_skip = jdata["model_devi_skip"]
     type_map = jdata["type_map"]
@@ -3179,10 +3197,6 @@
 
 
 def make_fp_vasp(iter_index, jdata):
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
     # abs path for fp_incar if it exists
     if "fp_incar" in jdata:
         jdata["fp_incar"] = os.path.abspath(jdata["fp_incar"])
@@ -3203,10 +3217,8 @@
 
 
 def make_fp_pwscf(iter_index, jdata):
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     # make pwscf input
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
@@ -3237,11 +3249,9 @@
 
 
 def make_fp_abacus_scf(iter_index, jdata):
-    # make config
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     pporb_path = "pporb"
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
     # make abacus/pw/scf input
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
@@ -3347,10 +3357,8 @@
 
 
 def make_fp_siesta(iter_index, jdata):
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     # make siesta input
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
@@ -3374,10 +3382,8 @@
 
 
 def make_fp_gaussian(iter_index, jdata):
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     # make gaussian gjf file
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
@@ -3403,10 +3409,8 @@
 
 
 def make_fp_cp2k(iter_index, jdata):
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     # make cp2k input
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, fp_name)
@@ -3445,10 +3449,6 @@
 
 
 def make_fp_pwmat(iter_index, jdata):
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
-    if len(fp_tasks) == 0:
-        return
     # abs path for fp_incar if it exists
     if "fp_incar" in jdata:
         jdata["fp_incar"] = os.path.abspath(jdata["fp_incar"])
@@ -3504,8 +3504,9 @@
        Jinzhe Zeng, Timothy J. Giese, Şölen Ekesan, and Darrin M. York, Journal of Chemical
        Theory and Computation 2021 17 (11), 6993-7009
     """
-    # make config
-    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
+    assert jdata["model_devi_engine"] == "amber"
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
     # make amber input
     cwd = os.getcwd()
     # link two mdin files and param7
@@ -3558,8 +3559,36 @@
 
 
 def make_fp(iter_index, jdata, mdata):
-    fp_style = jdata["fp_style"]
+    """Select the candidate strutures and make the input file of FP calculation.
 
+    Parameters
+    ----------
+    iter_index : int
+        iter index
+    jdata : dict
+        Run parameters.
+    mdata : dict
+        Machine parameters.
+    """
+    fp_tasks = _make_fp_vasp_configs(iter_index, jdata)
+    if len(fp_tasks) == 0:
+        return
+    make_fp_calculation(iter_index, jdata, mdata)
+
+
+def make_fp_calculation(iter_index, jdata, mdata):
+    """Make the input file of FP calculation.
+
+    Parameters
+    ----------
+    iter_index : int
+        iter index
+    jdata : dict
+        Run parameters.
+    mdata : dict
+        Machine parameters.
+    """
+    fp_style = jdata["fp_style"]
     if fp_style == "vasp":
         make_fp_vasp(iter_index, jdata)
     elif fp_style == "pwscf":

dpgen/simplify/arginfo.py

@@ -6,7 +6,10 @@
 from dpgen.generator.arginfo import (
     basic_args,
     data_args,
+    fp_style_abacus_args,
+    fp_style_cp2k_args,
     fp_style_gaussian_args,
+    fp_style_siesta_args,
     fp_style_vasp_args,
     training_args,
 )
@@ -59,7 +62,7 @@ def fp_style_variant_type_args() -> Variant:
     Variant
         variant for fp style
     """
-    doc_fp_style = "Software for First Principles, if `labeled` is false. Options include “vasp”, “gaussian” up to now."
+    doc_fp_style = "Software for First Principles, if `labeled` is false."
     doc_fp_style_none = "No fp."
     doc_fp_style_vasp = "VASP."
     doc_fp_style_gaussian = "Gaussian. The command should be set as `g16 < input`."
@@ -73,6 +76,15 @@ def fp_style_variant_type_args() -> Variant:
             Argument(
                 "gaussian", dict, fp_style_gaussian_args(), doc=doc_fp_style_gaussian
             ),
+            Argument("siesta", dict, fp_style_siesta_args()),
+            Argument("cp2k", dict, fp_style_cp2k_args()),
+            Argument("abacus", dict, fp_style_abacus_args()),
+            # TODO: not supported yet, as it requires model_devi_engine to be amber
+            # Argument(
+            #     "amber/diff", dict, fp_style_amber_diff_args(), doc=doc_amber_diff
+            # ),
+            Argument("pwmat", dict, [], doc="TODO: add doc"),
+            Argument("pwscf", dict, [], doc="TODO: add doc"),
         ],
         optional=True,
         default_tag="none",

dpgen/simplify/simplify.py

@@ -22,7 +22,6 @@
 
 from dpgen import dlog
 from dpgen.dispatcher.Dispatcher import make_submission
-from dpgen.generator.lib.gaussian import make_gaussian_input
 
 # TODO: maybe the following functions can be moved to dpgen.util
 from dpgen.generator.lib.utils import (
@@ -36,6 +35,7 @@
     data_system_fmt,
     fp_name,
     fp_task_fmt,
+    make_fp_calculation,
     make_fp_vasp_cp_cvasp,
     make_fp_vasp_incar,
     make_fp_vasp_kp,
@@ -414,65 +414,14 @@
         ii += 1
 
 
-def make_fp_gaussian(iter_index, jdata):
-    work_path = os.path.join(make_iter_name(iter_index), fp_name)
-    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
-    cwd = os.getcwd()
-    if "user_fp_params" in jdata.keys():
-        fp_params = jdata["user_fp_params"]
-    else:
-        fp_params = jdata["fp_params"]
-    cwd = os.getcwd()
-    for ii in fp_tasks:
-        os.chdir(ii)
-        sys_data = dpdata.System("POSCAR").data
-        ret = make_gaussian_input(sys_data, fp_params)
-        with open("input", "w") as fp:
-            fp.write(ret)
-        os.chdir(cwd)
-
-
-def make_fp_vasp(iter_index, jdata):
-    # abs path for fp_incar if it exists
-    if "fp_incar" in jdata:
-        jdata["fp_incar"] = os.path.abspath(jdata["fp_incar"])
-    # get nbands esti if it exists
-    if "fp_nbands_esti_data" in jdata:
-        nbe = NBandsEsti(jdata["fp_nbands_esti_data"])
-    else:
-        nbe = None
-    # order is critical!
-    # 1, create potcar
-    sys_link_fp_vasp_pp(iter_index, jdata)
-    # 2, create incar
-    make_fp_vasp_incar(iter_index, jdata, nbands_esti=nbe)
-    # 3, create kpoints
-    make_fp_vasp_kp(iter_index, jdata)
-    # 4, copy cvasp
-    make_fp_vasp_cp_cvasp(iter_index, jdata)
-
-
-def make_fp_calculation(iter_index, jdata):
-    fp_style = jdata["fp_style"]
-    if fp_style == "vasp":
-        make_fp_vasp(iter_index, jdata)
-    elif fp_style == "gaussian":
-        make_fp_gaussian(iter_index, jdata)
-    else:
-        raise RuntimeError("unsupported fp_style " + fp_style)
-
-
 def make_fp(iter_index, jdata, mdata):
     labeled = jdata.get("labeled", False)
     if labeled:
         make_fp_labeled(iter_index, jdata)
     else:
         make_fp_configs(iter_index, jdata)
-        make_fp_calculation(iter_index, jdata)
-        # Copy user defined forward_files
-        iter_name = make_iter_name(iter_index)
-        work_path = os.path.join(iter_name, fp_name)
-        symlink_user_forward_files(mdata=mdata, task_type="fp", work_path=work_path)
+        jdata["model_devi_engine"] = "lammps"
+        make_fp_calculation(iter_index, jdata, mdata)
 
 
 def run_iter(param_file, machine_file):

tests/generator/test_gromacs_engine.py

@@ -12,7 +12,7 @@
 __package__ = "generator"
 dirname = os.path.join(os.path.abspath(os.path.dirname(__file__)), "gromacs")
 
-from .context import make_fp_gaussian, make_model_devi
+from .context import make_fp, make_model_devi
 
 
 def _make_fake_graphs(train_path):
@@ -131,7 +131,7 @@ def test_make_model_devi_gromacs_with_openbabel(self):
         self._copy_outputs(
             os.path.join(self.dirname, "outputs"), self.model_devi_task_path
         )
-        make_fp_gaussian(iter_index=0, jdata=self.jdata)
+        make_fp(iter_index=0, jdata=self.jdata, mdata={})
         candi = np.loadtxt(
             os.path.join(self.fp_path, "candidate.shuffled.000.out"), dtype=str
         )