model_devi: add support for pimd (#1366)

Yi-FanLi · pre-commit-ci[bot] · njzjz · web-flow · commit 9ddef6ac3119 · 2023-11-02T13:26:11.000+08:00
Add support for LAMMPS's fix pimd/langevin in model deviation tasks.

---------

Signed-off-by: Yifan Li李一帆 &lt;yifanl0716@gmail.com&gt;
Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
Co-authored-by: Jinzhe Zeng &lt;jinzhe.zeng@rutgers.edu&gt;
diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py
@@ -262,6 +262,7 @@ def model_devi_jobs_args() -> list[Argument]:
     doc_press = "Pressure (Bar) in MD. Required when ensemble is npt."
     doc_trj_freq = "Frequecy of trajectory saved in MD."
     doc_nsteps = "Running steps of MD. It is not optional when not using a template."
+    doc_nbeads = "Number of beads in PIMD. If not given, classical MD will be performed. Only supported for LAMMPS version >= 20230615."
     doc_ensemble = "Determining which ensemble used in MD, options include “npt” and “nvt”. It is not optional when not using a template."
     doc_neidelay = "delay building until this many steps since last build."
     doc_taut = "Coupling time of thermostat (ps)."
@@ -280,6 +281,7 @@ def model_devi_jobs_args() -> list[Argument]:
         Argument("press", list[float], optional=True, doc=doc_press),
         Argument("trj_freq", int, optional=False, doc=doc_trj_freq),
         Argument("nsteps", int, optional=True, doc=doc_nsteps),
+        Argument("nbeads", int, optional=True, doc=doc_nbeads),
         Argument("ensemble", str, optional=True, doc=doc_ensemble),
         Argument("neidelay", int, optional=True, doc=doc_neidelay),
         Argument("taut", float, optional=True, doc=doc_taut),
diff --git a/dpgen/generator/lib/lammps.py b/dpgen/generator/lib/lammps.py
@@ -37,6 +37,7 @@ def make_lammps_input(
     max_seed=1000000,
     nopbc=False,
     deepmd_version="0.1",
+    nbeads=None,
 ):
     if (ele_temp_f is not None or ele_temp_a is not None) and Version(
         deepmd_version
@@ -49,9 +50,22 @@ def make_lammps_input(
             "the frame style ele_temp and atom style ele_temp should not be set at the same time"
         )
     ret = "variable        NSTEPS          equal %d\n" % nsteps
+    if nbeads is not None:
+        if nbeads <= 0:
+            raise ValueError(
+                "The number of beads should be positive. Check your nbeads setting."
+            )
+        power = 1
+        while power < nbeads:
+            power *= 10
+        ret += "variable        ibead           uloop %d pad\n" % (power - 1)
+    if nbeads is not None:
+        ret += "atom_modify        map yes\n"
     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 nbeads is not None:
+        ret += "variable        TEMP_NBEADS            equal %f\n" % (temp * nbeads)
     if ele_temp_f is not None:
         ret += "variable        ELE_TEMP        equal %f\n" % ele_temp_f
     if ele_temp_a is not None:
@@ -72,10 +86,16 @@ def make_lammps_input(
         ret += "neigh_modify    delay %d\n" % neidelay
     ret += "\n"
     ret += "box          tilt large\n"
-    ret += (
-        'if "${restart} > 0" then "read_restart dpgen.restart.*" else "read_data %s"\n'
-        % conf_file
-    )
+    if nbeads is None:
+        ret += (
+            'if "${restart} > 0" then "read_restart dpgen.restart.*" else "read_data %s"\n'
+            % conf_file
+        )
+    else:
+        ret += (
+            'if "${restart} > 0" then "read_restart dpgen.restart${ibead}.*" else "read_data %s"\n'
+            % conf_file
+        )
     ret += "change_box   all triclinic\n"
     for jj in range(len(mass_map)):
         ret += "mass            %d %f\n" % (jj + 1, mass_map[jj])
@@ -98,23 +118,43 @@ def make_lammps_input(
             keywords += "fparam ${ELE_TEMP}"
         if ele_temp_a is not None:
             keywords += "aparam ${ELE_TEMP}"
-        ret += f"pair_style      deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords}\n"
+        if nbeads is None:
+            ret += f"pair_style      deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords}\n"
+        else:
+            ret += f"pair_style      deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi${{ibead}}.out {keywords}\n"
     ret += "pair_coeff      * *\n"
     ret += "\n"
     ret += "thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz\n"
     ret += "thermo          ${THERMO_FREQ}\n"
     model_devi_merge_traj = jdata.get("model_devi_merge_traj", False)
-    if model_devi_merge_traj is True:
-        ret += "dump            1 all custom ${DUMP_FREQ} all.lammpstrj id type x y z fx fy fz\n"
-        ret += 'if "${restart} > 0" then "dump_modify     1 append yes"\n'
+    if nbeads is None:
+        if model_devi_merge_traj is True:
+            ret += "dump            1 all custom ${DUMP_FREQ} all.lammpstrj id type x y z fx fy fz\n"
+            ret += 'if "${restart} > 0" then "dump_modify     1 append yes"\n'
+        else:
+            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} traj/*.lammpstrj id type x y z fx fy fz\n"
-    ret += "restart         10000 dpgen.restart\n"
+        if model_devi_merge_traj is True:
+            ret += "dump            1 all custom ${DUMP_FREQ} all.lammpstrj${ibead} id type x y z fx fy fz\n"
+            ret += 'if "${restart} > 0" then "dump_modify     1 append yes"\n'
+        else:
+            ret += "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj${ibead} id type x y z fx fy fz\n"
+    if nbeads is None:
+        ret += "restart         10000 dpgen.restart\n"
+    else:
+        ret += "restart         10000 dpgen.restart${ibead}\n"
     ret += "\n"
     if pka_e is None:
-        ret += 'if "${restart} == 0" then "velocity        all create ${TEMP} %d"' % (
-            random.randrange(max_seed - 1) + 1
-        )
+        if nbeads is None:
+            ret += (
+                'if "${restart} == 0" then "velocity        all create ${TEMP} %d"'
+                % (random.randrange(max_seed - 1) + 1)
+            )
+        else:
+            ret += (
+                'if "${restart} == 0" then "velocity        all create ${TEMP_NBEADS} %d"'
+                % (random.randrange(max_seed - 1) + 1)
+            )
     else:
         sys = dpdata.System(conf_file, fmt="lammps/lmp")
         sys_data = sys.data
@@ -140,18 +180,34 @@ def make_lammps_input(
         assert pres is not None
         if nopbc:
             raise RuntimeError("ensemble %s is conflicting with nopbc" % ensemble)
-    if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso":
-        ret += "fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P}\n"
-    elif ensemble == "npt-a" or ensemble == "npt-aniso":
-        ret += "fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} aniso ${PRES} ${PRES} ${TAU_P}\n"
-    elif ensemble == "npt-t" or ensemble == "npt-tri":
-        ret += "fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} tri ${PRES} ${PRES} ${TAU_P}\n"
-    elif ensemble == "nvt":
-        ret += "fix             1 all nvt temp ${TEMP} ${TEMP} ${TAU_T}\n"
-    elif ensemble == "nve":
-        ret += "fix             1 all nve\n"
+    if nbeads is None:
+        if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso":
+            ret += "fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P}\n"
+        elif ensemble == "npt-a" or ensemble == "npt-aniso":
+            ret += "fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} aniso ${PRES} ${PRES} ${TAU_P}\n"
+        elif ensemble == "npt-t" or ensemble == "npt-tri":
+            ret += "fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} tri ${PRES} ${PRES} ${TAU_P}\n"
+        elif ensemble == "nvt":
+            ret += "fix             1 all nvt temp ${TEMP} ${TEMP} ${TAU_T}\n"
+        elif ensemble == "nve":
+            ret += "fix             1 all nve\n"
+        else:
+            raise RuntimeError("unknown emsemble " + ensemble)
     else:
-        raise RuntimeError("unknown emsemble " + ensemble)
+        if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso":
+            ret += "fix 1 all pimd/langevin fmmode physical ensemble npt integrator obabo thermostat PILE_L ${ibead} temp ${TEMP} tau ${TAU_T} scale 1.0 barostat BZP iso ${PRES} taup ${TAU_P}\n"
+        elif ensemble == "npt-a" or ensemble == "npt-aniso":
+            ret += "fix 1 all pimd/langevin fmmode physical ensemble npt integrator obabo thermostat PILE_L ${ibead} temp ${TEMP} tau ${TAU_T} scale 1.0 barostat BZP aniso ${PRES} taup ${TAU_P}\n"
+        elif ensemble == "nvt":
+            ret += "fix 1 all pimd/langevin fmmode physical ensemble nvt integrator obabo thermostat PILE_L ${ibead} temp ${TEMP} tau ${TAU_T} scale 1.0\n"
+        elif ensemble == "nve":
+            ret += "fix 1 all pimd/langevin fmmode physical ensemble nve integrator obabo temp ${TEMP}\n"
+        else:
+            raise RuntimeError(
+                "unknown emsemble "
+                + ensemble
+                + " for fix pimd/langevin\nrefer to https://docs.lammps.org/fix_pimd.html for more information"
+            )
     if nopbc:
         ret += "velocity        all zero linear\n"
         ret += "fix             fm all momentum 1 linear 1 1 1\n"
diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py
@@ -18,6 +18,7 @@
 import os
 import queue
 import random
+import re
 import shutil
 import sys
 import warnings
@@ -915,7 +916,11 @@ def parse_cur_job(cur_job):
         dt = _get_param_alias(cur_job, ["dt"])
     else:
         dt = None
-    return ensemble, nsteps, trj_freq, temps, press, pka_e, dt
+    if "nbeads" in cur_job:
+        nbeads = _get_param_alias(cur_job, ["nbeads"])
+    else:
+        nbeads = None
+    return ensemble, nsteps, trj_freq, temps, press, pka_e, dt, nbeads
 
 
 def expand_matrix_values(target_list, cur_idx=0):
@@ -1457,7 +1462,21 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems):
     if iter_index >= len(model_devi_jobs):
         return False
     cur_job = model_devi_jobs[iter_index]
-    ensemble, nsteps, trj_freq, temps, press, pka_e, dt = parse_cur_job(cur_job)
+    ensemble, nsteps, trj_freq, temps, press, pka_e, dt, nbeads = parse_cur_job(cur_job)
+    model_devi_f_avg_relative = jdata.get("model_devi_f_avg_relative", False)
+    model_devi_merge_traj = jdata.get("model_devi_merge_traj", False)
+    if (nbeads is not None) and model_devi_f_avg_relative:
+        raise RuntimeError(
+            "model_devi_f_avg_relative has not been supported for pimd. Set model_devi_f_avg_relative to False."
+        )
+    if (nbeads is not None) and (model_devi_merge_traj):
+        raise RuntimeError(
+            "model_devi_merge_traj has not been supported for pimd. Set model_devi_merge_traj to False."
+        )
+    if (nbeads is not None) and (not (nsteps % trj_freq == 0)):
+        raise RuntimeError(
+            "trj_freq should be a factor of nsteps for pimd. Please check your input."
+        )
     if dt is not None:
         model_devi_dt = dt
     sys_idx = expand_idx(cur_job["sys_idx"])
@@ -1560,6 +1579,7 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems):
                         ele_temp_a=te_a,
                         nopbc=nopbc,
                         deepmd_version=deepmd_version,
+                        nbeads=nbeads,
                     )
                     job = {}
                     job["ensemble"] = ensemble
@@ -1916,12 +1936,23 @@ def run_md_model_devi(iter_index, jdata, mdata):
 
     model_devi_engine = jdata.get("model_devi_engine", "lammps")
     if model_devi_engine == "lammps":
-        command = f"{{ if [ ! -f dpgen.restart.10000 ]; then {model_devi_exec} -i input.lammps -v restart 0; else {model_devi_exec} -i input.lammps -v restart 1; fi }}"
-        command = "/bin/sh -c '%s'" % command
+        nbeads = jdata["model_devi_jobs"][iter_index].get("nbeads")
+        if nbeads is None:
+            command = f"{{ if [ ! -f dpgen.restart.10000 ]; then {model_devi_exec} -i input.lammps -v restart 0; else {model_devi_exec} -i input.lammps -v restart 1; fi }}"
+        else:
+            command = f"{{ all_exist=true; for i in $(seq -w 1 {nbeads}); do [[ ! -f dpgen.restart${{i}}.10000 ]] && {{ all_exist=false; break; }}; done; $all_exist && {{ {model_devi_exec} -p {nbeads}x1 -i input.lammps -v restart 1; }} || {{ {model_devi_exec} -p {nbeads}x1 -i input.lammps -v restart 0; }} }}"
+        command = "/bin/bash -c '%s'" % command
         commands = [command]
 
         forward_files = ["conf.lmp", "input.lammps"]
-        backward_files = ["model_devi.out", "model_devi.log"]
+        backward_files = ["model_devi.log"]
+        if nbeads is None:
+            backward_files += ["model_devi.out"]
+        else:
+            num_digits = np.ceil(np.log10(nbeads + 1)).astype(int)
+            backward_files += [
+                f"model_devi{i+1:0{num_digits}d}.out" for i in range(nbeads)
+            ]
         if model_devi_merge_traj:
             backward_files += ["all.lammpstrj"]
         else:
@@ -2131,6 +2162,75 @@ def _read_model_devi_file(
     model_devi_f_avg_relative: bool = False,
     model_devi_merge_traj: bool = False,
 ):
+    model_devi_files = glob.glob(os.path.join(task_path, "model_devi*.out"))
+    model_devi_files_sorted = sorted(
+        model_devi_files, key=lambda x: int(re.search(r"(\d+)", x).group(1))
+    )
+    if len(model_devi_files_sorted) > 1:
+        with open(model_devi_files_sorted[0]) as f:
+            first_line = f.readline()
+        if not (first_line.startswith("#")):
+            first_line = "#"
+        num_beads = len(model_devi_files_sorted)
+        model_devi_contents = []
+        for file in model_devi_files_sorted:
+            model_devi_contents.append(np.loadtxt(file))
+        assert all(
+            model_devi_content.shape[0] == model_devi_contents[0].shape[0]
+            for model_devi_content in model_devi_contents
+        ), "Not all beads generated the same number of lines in the model_devi$\{ibead\}.out file. Check your pimd task carefully."
+        for file in model_devi_files_sorted:
+            os.remove(file)
+        last_step = model_devi_contents[0][-1, 0]
+        for ibead in range(1, num_beads):
+            model_devi_contents[ibead][:, 0] = model_devi_contents[ibead][
+                :, 0
+            ] + ibead * (last_step + 1)
+        model_devi = np.concatenate(model_devi_contents, axis=0)
+        num_columns = model_devi.shape[1]
+        formats = ["%12d"] + ["%22.6e"] * (num_columns - 1)
+        np.savetxt(
+            os.path.join(task_path, "model_devi.out"),
+            model_devi,
+            fmt=formats,
+            header=first_line.rstrip(),
+            comments="",
+        )
+
+        if not model_devi_merge_traj:
+            num_digits = np.ceil(np.log10(num_beads + 1)).astype(int)
+            traj_files_sorted = []
+            for ibead in range(num_beads):
+                traj_files = glob.glob(
+                    os.path.join(
+                        task_path, "traj", f"*lammpstrj{ibead+1:0{num_digits}d}"
+                    )
+                )
+                traj_files_sorted.append(
+                    sorted(
+                        traj_files,
+                        key=lambda x: int(
+                            re.search(r"^(\d+)\.lammpstrj", os.path.basename(x)).group(
+                                1
+                            )
+                        ),
+                    )
+                )
+            assert all(
+                len(traj_list) == len(traj_files_sorted[0])
+                for traj_list in traj_files_sorted
+            ), "Not all beads generated the same number of frames. Check your pimd task carefully."
+            for ibead in range(num_beads):
+                for itraj in range(len(traj_files_sorted[0])):
+                    base_path, original_filename = os.path.split(
+                        traj_files_sorted[ibead][itraj]
+                    )
+                    frame_number = int(original_filename.split(".")[0])
+                    new_filename = os.path.join(
+                        base_path,
+                        f"{frame_number + ibead * (int(last_step)+1):d}.lammpstrj",
+                    )
+                    os.rename(traj_files_sorted[ibead][itraj], new_filename)
     model_devi = np.loadtxt(os.path.join(task_path, "model_devi.out"))
     if model_devi_f_avg_relative:
         if model_devi_merge_traj is True:
diff --git a/tests/generator/context.py b/tests/generator/context.py
@@ -20,6 +20,7 @@
 from dpgen.util import setup_ele_temp  # noqa: F401
 
 param_file = "param-mg-vasp.json"
+param_pimd_file = "param-mg-pimd-vasp.json"
 param_file_merge_traj = "param-mg-vasp_merge_traj.json"
 param_file_v1 = "param-mg-vasp-v1.json"
 param_file_v1_et = "param-mg-vasp-v1-et.json"
diff --git a/tests/generator/param-mg-pimd-vasp.json b/tests/generator/param-mg-pimd-vasp.json
diff --git a/tests/generator/test_make_md.py b/tests/generator/test_make_md.py
