support NvNMD-train and NvNMD-explore

dpgen2/constants.py

@@ -1,8 +1,11 @@
 train_index_pattern = "%04d"
 train_task_pattern = "task." + train_index_pattern
 train_script_name = "input.json"
+train_cnn_script_name = "input_cnn.json"
+train_qnn_script_name = "input_qnn.json"
 train_log_name = "train.log"
 model_name_pattern = "model.%03d.pb"
+nvnmd_model_name_pattern = "nvnmd_model.%03d"
 pytorch_model_name_pattern = "model.%03d.pth"
 model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth)"
 lmp_index_pattern = "%06d"

dpgen2/entrypoint/args.py

@@ -127,12 +127,59 @@ def dp_train_args():
     ]
 
 
+def nvnmd_train_args():
+    doc_numb_models = "Number of models trained for evaluating the model deviation"
+    doc_config = "Configuration of training"
+    doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. "
+    doc_init_models_paths = "the paths to initial models"
+    doc_init_models_uri = "The URI of initial models"
+    doc_optional_files = "Optional files for training"
+
+    return [
+        Argument(
+            "config",
+            dict,
+            RunDPTrain.training_args(),
+            optional=True,
+            default=RunDPTrain.normalize_config({}),
+            doc=doc_config,
+        ),
+        Argument("numb_models", int, optional=True, default=4, doc=doc_numb_models),
+        Argument(
+            "template_script", [List[str], str], optional=False, doc=doc_template_script
+        ),
+        Argument(
+            "init_models_paths",
+            List[str],
+            optional=True,
+            default=None,
+            doc=doc_init_models_paths,
+            alias=["training_iter0_model_path"],
+        ),
+        Argument(
+            "init_models_uri",
+            str,
+            optional=True,
+            default=None,
+            doc=doc_init_models_uri,
+        ),
+        Argument(
+            "optional_files",
+            list,
+            optional=True,
+            default=None,
+            doc=doc_optional_files,
+        ),
+    ]
+
+
 def variant_train():
     doc = "the type of the training"
     return Variant(
         "type",
         [
             Argument("dp", dict, dp_train_args()),
+            Argument("dp-nvnmd", dict, nvnmd_train_args()),
             Argument("dp-dist", dict, dp_dist_train_args()),
         ],
         doc=doc,
@@ -454,6 +501,7 @@ def variant_explore():
         "type",
         [
             Argument("lmp", dict, lmp_args(), doc=doc_lmp),
+            Argument("lmp-nvnmd", dict, lmp_args(), doc=doc_lmp),
             Argument("calypso", dict, caly_args(), doc=doc_calypso),
             Argument("calypso:default", dict, caly_args(), doc=doc_calypso),
             Argument("calypso:merge", dict, caly_args(), doc=doc_calypso),

dpgen2/entrypoint/submit.py

@@ -111,6 +111,8 @@
     RunDPTrain,
     RunLmp,
     RunLmpHDF5,
+    RunNvNMD,
+    RunNvNMDTrain,
     RunRelax,
     RunRelaxHDF5,
     SelectConfs,
@@ -182,6 +184,17 @@ def make_concurrent_learning_op(
             valid_data=valid_data,
             optional_files=train_optional_files,
         )
+    elif train_style == "dp-nvnmd":
+        prep_run_train_op = PrepRunDPTrain(
+            "prep-run-nvnmd-train",
+            PrepDPTrain,
+            RunNvNMDTrain,
+            prep_config=prep_train_config,
+            run_config=run_train_config,
+            upload_python_packages=upload_python_packages,
+            valid_data=valid_data,
+            optional_files=train_optional_files,
+        )
     else:
         raise RuntimeError(f"unknown train_style {train_style}")
     if explore_style == "lmp":
@@ -193,6 +206,15 @@ def make_concurrent_learning_op(
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,
         )
+    elif "lmp-nvnmd" in explore_style:
+        prep_run_explore_op = PrepRunLmp(
+            "prep-run-nvnmd",
+            PrepLmp,
+            RunNvNMD,
+            prep_config=prep_explore_config,
+            run_config=run_explore_config,
+            upload_python_packages=upload_python_packages,
+        )
     elif "calypso" in explore_style:
         expl_mode = explore_style.split(":")[-1] if ":" in explore_style else "default"
         if expl_mode == "merge":
@@ -286,7 +308,7 @@ def make_naive_exploration_scheduler(
     # use npt task group
     explore_style = config["explore"]["type"]
 
-    if explore_style == "lmp":
+    if explore_style in ("lmp", "lmp-nvnmd"):
         return make_lmp_naive_exploration_scheduler(config)
     elif "calypso" in explore_style or explore_style == "diffcsp":
         return make_naive_exploration_scheduler_without_conf(config, explore_style)
@@ -374,6 +396,7 @@ def make_lmp_naive_exploration_scheduler(config):
     output_nopbc = config["explore"]["output_nopbc"]
     conf_filters = get_conf_filters(config["explore"]["filters"])
     use_ele_temp = config["inputs"]["use_ele_temp"]
+    config["explore"]["type"]
     scheduler = ExplorationScheduler()
     # report
     conv_style = convergence.pop("type")
@@ -506,6 +529,16 @@ def workflow_concurrent_learning(
             else None
         )
         config["train"]["numb_models"] = 1
+
+    elif train_style == "dp-nvnmd":
+        init_models_paths = config["train"].get("init_models_paths", None)
+        numb_models = config["train"]["numb_models"]
+        if init_models_paths is not None and len(init_models_paths) != numb_models:
+            raise RuntimeError(
+                f"{len(init_models_paths)} init models provided, which does "
+                "not match numb_models={numb_models}"
+            )
+
     else:
         raise RuntimeError(f"unknown params, train_style: {train_style}")
 

dpgen2/exploration/render/traj_render_lammps.py

@@ -109,6 +109,10 @@ def get_confs(
         conf_filters: Optional["ConfFilters"] = None,
         optional_outputs: Optional[List[Path]] = None,
     ) -> dpdata.MultiSystems:
+        from ase.io import (  # type: ignore
+            read,
+        )
+
         ntraj = len(trajs)
         ele_temp = None
         if optional_outputs:
@@ -123,12 +127,16 @@ def get_confs(
                     traj = StringIO(trajs[ii].get_data())  # type: ignore
                 else:
                     traj = trajs[ii]
-                ss = dpdata.System(traj, fmt=traj_fmt, type_map=type_map)
-                ss.nopbc = self.nopbc
-                if ele_temp:
-                    self.set_ele_temp(ss, ele_temp[ii])
-                ss = ss.sub_system(id_selected[ii])
-                ms.append(ss)
+                # ss = dpdata.System(traj, fmt=traj_fmt, type_map=type_map)
+                ss = read(
+                    str(traj), format="lammps-dump-text", index=":", specorder=type_map
+                )
+                for jj in id_selected[ii]:
+                    s = dpdata.System(ss[jj], fmt="ase/structure", type_map=type_map)
+                    s.nopbc = self.nopbc
+                    if ele_temp:
+                        self.set_ele_temp(s, ele_temp[ii])
+                    ms.append(s)
         if conf_filters is not None:
             ms = conf_filters.check(ms)
         return ms

dpgen2/exploration/task/lmp/lmp_input.py

@@ -50,6 +50,7 @@ def make_lmp_input(
     nopbc: bool = False,
     max_seed: int = 1000000,
     deepmd_version="2.0",
+    nvnmd_version=None,
     trj_seperate_files=True,
     pimd_bead: Optional[str] = None,
 ):
@@ -69,9 +70,9 @@ def make_lmp_input(
     ret += "variable        THERMO_FREQ     equal %d\n" % trj_freq
     ret += "variable        DUMP_FREQ       equal %d\n" % trj_freq
     ret += "variable        TEMP            equal %f\n" % temp
-    if ele_temp_f is not None:
+    if ele_temp_f is not None and nvnmd_version is None:
         ret += "variable        ELE_TEMP        equal %f\n" % ele_temp_f
-    if ele_temp_a is not None:
+    if ele_temp_a is not None and nvnmd_version is None:
         ret += "variable        ELE_TEMP        equal %f\n" % ele_temp_a
     if pres is not None:
         ret += "variable        PRES            equal %f\n" % pres
@@ -106,12 +107,14 @@ def make_lmp_input(
         if pimd_bead is not None
         else lmp_model_devi_name
     )
-    if Version(deepmd_version) < Version("1"):
+    if Version(deepmd_version) < Version("1") and nvnmd_version is None:
         # 0.x
         ret += "pair_style      deepmd %s ${THERMO_FREQ} %s\n" % (
             graph_list,
             model_devi_file_name,
         )
+    elif nvnmd_version is not None:
+        ret += "pair_style      nvnmd %s\n" % ("model.pb")
     else:
         # 1.x
         keywords = ""
@@ -135,17 +138,28 @@ def make_lmp_input(
     ret += "thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz\n"
     ret += "thermo          ${THERMO_FREQ}\n"
     if trj_seperate_files:
-        ret += "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n"
+        if nvnmd_version is None:
+            ret += "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n"
+        else:
+            ret += "dump            1 all custom ${DUMP_FREQ} ${rerun}_traj/*.lammpstrj id type x y z fx fy fz\n"
     else:
         lmp_traj_file_name = (
             lmp_pimd_traj_name % pimd_bead if pimd_bead is not None else lmp_traj_name
         )
-        ret += (
-            "dump            1 all custom ${DUMP_FREQ} %s id type x y z fx fy fz\n"
-            % lmp_traj_file_name
-        )
+        if nvnmd_version is None:
+            ret += (
+                "dump            1 all custom ${DUMP_FREQ} %s id type x y z fx fy fz\n"
+                % lmp_traj_file_name
+            )
+        else:
+            ret += (
+                "dump            1 all custom ${DUMP_FREQ} ${rerun}_%s id type x y z fx fy fz\n"
+                % lmp_traj_file_name
+            )
     ret += "restart         10000 dpgen.restart\n"
     ret += "\n"
+    if nvnmd_version is not None:
+        ret += 'if "${rerun} > 0" then "jump SELF rerun"\n'
     if pka_e is None:
         ret += 'if "${restart} == 0" then "velocity        all create ${TEMP} %d"' % (
             random.randrange(max_seed - 1) + 1
@@ -193,4 +207,12 @@ def make_lmp_input(
     ret += "\n"
     ret += "timestep        %f\n" % dt
     ret += "run             ${NSTEPS} upto\n"
+    if nvnmd_version is not None:
+        ret += "jump SELF end\n"
+        ret += "label rerun\n"
+        if trj_seperate_files:
+            ret += "rerun 0_traj/*.lammpstrj dump x y z fx fy fz add yes\n"
+        else:
+            ret += "rerun 0_%s dump x y z fx fy fz add yes\n" % lmp_traj_name
+        ret += "label end\n"
     return ret

dpgen2/exploration/task/lmp_template_task_group.py

@@ -47,12 +47,14 @@ def set_lmp(
         revisions: dict = {},
         traj_freq: int = 10,
         extra_pair_style_args: str = "",
+        nvnmd_version: Optional[str] = None,
         pimd_bead: Optional[str] = None,
     ) -> None:
         self.lmp_template = Path(lmp_template_fname).read_text().split("\n")
         self.revisions = revisions
         self.traj_freq = traj_freq
         self.extra_pair_style_args = extra_pair_style_args
+        self.nvnmd_version = nvnmd_version
         self.pimd_bead = pimd_bead
         self.lmp_set = True
         self.model_list = sorted([model_name_pattern % ii for ii in range(numb_models)])
@@ -62,10 +64,16 @@ def set_lmp(
             self.traj_freq,
             self.extra_pair_style_args,
             self.pimd_bead,
+            nvnmd_version=self.nvnmd_version,
         )
         self.lmp_template = revise_lmp_input_dump(
-            self.lmp_template, self.traj_freq, self.pimd_bead
+            self.lmp_template,
+            self.traj_freq,
+            self.pimd_bead,
+            nvnmd_version=self.nvnmd_version,
         )
+        if nvnmd_version is not None:
+            self.lmp_template = revise_lmp_input_rerun(self.lmp_template)
         if plm_template_fname is not None:
             self.plm_template = Path(plm_template_fname).read_text().split("\n")
             self.plm_set = True
@@ -158,8 +166,8 @@ def revise_lmp_input_model(
     extra_pair_style_args="",
     pimd_bead=None,
     deepmd_version="1",
+    nvnmd_version=None,
 ):
-    idx = find_only_one_key(lmp_lines, ["pair_style", "deepmd"])
     if extra_pair_style_args:
         extra_pair_style_args = " " + extra_pair_style_args
     graph_list = " ".join(task_model_list)
@@ -168,23 +176,39 @@ def revise_lmp_input_model(
         if pimd_bead is not None
         else lmp_model_devi_name
     )
-    lmp_lines[idx] = "pair_style      deepmd %s out_freq %d out_file %s%s" % (
-        graph_list,
-        trj_freq,
-        model_devi_file_name,
-        extra_pair_style_args,
-    )
+    if nvnmd_version is None:
+        idx = find_only_one_key(lmp_lines, ["pair_style", "deepmd"])
+        lmp_lines[idx] = "pair_style      deepmd %s out_freq %d out_file %s%s" % (
+            graph_list,
+            trj_freq,
+            model_devi_file_name,
+            extra_pair_style_args,
+        )
+    else:
+        idx = find_only_one_key(lmp_lines, ["pair_style", "nvnmd"])
+        lmp_lines[idx] = "pair_style      nvnmd %s %s" % (
+            "model.pb",
+            extra_pair_style_args,
+        )
+
     return lmp_lines
 
 
-def revise_lmp_input_dump(lmp_lines, trj_freq, pimd_bead=None):
+def revise_lmp_input_dump(lmp_lines, trj_freq, pimd_bead=None, nvnmd_version=None):
     idx = find_only_one_key(lmp_lines, ["dump", "dpgen_dump"])
     lmp_traj_file_name = (
         lmp_pimd_traj_name % pimd_bead if pimd_bead is not None else lmp_traj_name
     )
-    lmp_lines[
-        idx
-    ] = f"dump            dpgen_dump all custom {trj_freq} {lmp_traj_file_name} id type x y z"
+    if nvnmd_version is None:
+        lmp_lines[
+            idx
+        ] = f"dump            dpgen_dump all custom {trj_freq} {lmp_traj_file_name} id type x y z"
+    else:
+        lmp_lines[idx] = (
+            "dump            dpgen_dump all custom %s ${rerun}_%s id type x y z fx fy fz"
+            % (trj_freq, lmp_traj_file_name)
+        )
+        lmp_lines.insert(idx + 1, 'if "${rerun} > 0" then "jump SELF rerun"')
     return lmp_lines
 
 
@@ -197,6 +221,14 @@ def revise_lmp_input_plm(lmp_lines, in_plm, out_plm="output.plumed"):
     return lmp_lines
 
 
+def revise_lmp_input_rerun(lmp_lines):
+    lmp_lines.append("jump SELF end")
+    lmp_lines.append("label rerun")
+    lmp_lines.append(f"rerun 0_{lmp_traj_name} dump x y z fx fy fz add yes")
+    lmp_lines.append("label end")
+    return lmp_lines
+
+
 def revise_by_keys(lmp_lines, keys, values):
     for kk, vv in zip(keys, values):  # type: ignore
         for ii in range(len(lmp_lines)):