Support PIMD for LAMMPS (#263)

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

## Release Notes

- **New Features**
- Introduced new parameters for LAMMPS input file generation, enhancing
flexibility in handling PIMD tasks.
- Added a function to merge multiple trajectory and model deviation
files into consolidated outputs.
- Enhanced testing capabilities with new templates for PIMD simulations.

- **Bug Fixes**
  - Minor corrections in documentation strings for clarity.

- **Tests**
- Added new tests to ensure the functionality of merging PIMD files is
working as expected.
- Included tests to validate the PIMD template integration within the
existing test framework.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: zjgemi <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Han Wang <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

dpgen2/constants.py

@@ -12,8 +12,10 @@
 plm_input_name = "input.plumed"
 plm_output_name = "output.plumed"
 lmp_traj_name = "traj.dump"
+lmp_pimd_traj_name = "traj.%s.dump"
 lmp_log_name = "log.lammps"
 lmp_model_devi_name = "model_devi.out"
+lmp_pimd_model_devi_name = "model_devi.%s.out"
 fp_index_pattern = "%06d"
 fp_task_pattern = "task." + fp_index_pattern
 fp_default_log_name = "fp.log"

dpgen2/exploration/task/lmp/lmp_input.py

@@ -12,6 +12,9 @@
 )
 
 from dpgen2.constants import (
+    lmp_model_devi_name,
+    lmp_pimd_model_devi_name,
+    lmp_pimd_traj_name,
     lmp_traj_name,
 )
 
@@ -48,6 +51,7 @@ def make_lmp_input(
     max_seed: int = 1000000,
     deepmd_version="2.0",
     trj_seperate_files=True,
+    pimd_bead: Optional[str] = None,
 ):
     if (ele_temp_f is not None or ele_temp_a is not None) and Version(
         deepmd_version
@@ -97,9 +101,17 @@ def make_lmp_input(
     graph_list = ""
     for ii in graphs:
         graph_list += ii + " "
+    model_devi_file_name = (
+        lmp_pimd_model_devi_name % pimd_bead
+        if pimd_bead is not None
+        else lmp_model_devi_name
+    )
     if Version(deepmd_version) < Version("1"):
         # 0.x
-        ret += "pair_style      deepmd %s ${THERMO_FREQ} model_devi.out\n" % graph_list
+        ret += "pair_style      deepmd %s ${THERMO_FREQ} %s\n" % (
+            graph_list,
+            model_devi_file_name,
+        )
     else:
         # 1.x
         keywords = ""
@@ -113,9 +125,10 @@ def make_lmp_input(
             keywords += "fparam ${ELE_TEMP}"
         if ele_temp_a is not None:
             keywords += "aparam ${ELE_TEMP}"
-        ret += (
-            "pair_style      deepmd %s out_freq ${THERMO_FREQ} out_file model_devi.out %s\n"
-            % (graph_list, keywords)
+        ret += "pair_style      deepmd %s out_freq ${THERMO_FREQ} out_file %s %s\n" % (
+            graph_list,
+            model_devi_file_name,
+            keywords,
         )
     ret += "pair_coeff      * *\n"
     ret += "\n"
@@ -124,9 +137,12 @@ def make_lmp_input(
     if trj_seperate_files:
         ret += "dump            1 all custom ${DUMP_FREQ} 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_name
+            % lmp_traj_file_name
         )
     ret += "restart         10000 dpgen.restart\n"
     ret += "\n"

dpgen2/exploration/task/lmp_template_task_group.py

@@ -11,6 +11,9 @@
 from dpgen2.constants import (
     lmp_conf_name,
     lmp_input_name,
+    lmp_model_devi_name,
+    lmp_pimd_model_devi_name,
+    lmp_pimd_traj_name,
     lmp_traj_name,
     model_name_pattern,
     plm_input_name,
@@ -44,20 +47,25 @@ def set_lmp(
         revisions: dict = {},
         traj_freq: int = 10,
         extra_pair_style_args: str = "",
+        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.pimd_bead = pimd_bead
         self.lmp_set = True
         self.model_list = sorted([model_name_pattern % ii for ii in range(numb_models)])
         self.lmp_template = revise_lmp_input_model(
             self.lmp_template,
             self.model_list,
             self.traj_freq,
             self.extra_pair_style_args,
+            self.pimd_bead,
+        )
+        self.lmp_template = revise_lmp_input_dump(
+            self.lmp_template, self.traj_freq, self.pimd_bead
         )
-        self.lmp_template = revise_lmp_input_dump(self.lmp_template, self.traj_freq)
         if plm_template_fname is not None:
             self.plm_template = Path(plm_template_fname).read_text().split("\n")
             self.plm_set = True
@@ -144,29 +152,39 @@ def find_only_one_key(lmp_lines, key):
 
 
 def revise_lmp_input_model(
-    lmp_lines, task_model_list, trj_freq, extra_pair_style_args="", deepmd_version="1"
+    lmp_lines,
+    task_model_list,
+    trj_freq,
+    extra_pair_style_args="",
+    pimd_bead=None,
+    deepmd_version="1",
 ):
     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)
-    lmp_lines[idx] = (
-        "pair_style      deepmd %s out_freq %d out_file model_devi.out%s"
-        % (
-            graph_list,
-            trj_freq,
-            extra_pair_style_args,
-        )
+    model_devi_file_name = (
+        lmp_pimd_model_devi_name % pimd_bead
+        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,
     )
     return lmp_lines
 
 
-def revise_lmp_input_dump(lmp_lines, trj_freq):
+def revise_lmp_input_dump(lmp_lines, trj_freq, pimd_bead=None):
     idx = find_only_one_key(lmp_lines, ["dump", "dpgen_dump"])
-    lmp_lines[idx] = (
-        f"dump            dpgen_dump all custom %d {lmp_traj_name} id type x y z"
-        % trj_freq
+    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"
     return lmp_lines
 
 

dpgen2/exploration/task/make_task_group_from_config.py

@@ -47,6 +47,7 @@ def npt_task_group_args():
     doc_relative_v_epsilon = "Calculate relative virial model deviation"
     doc_ele_temp_f = "The electron temperature set by frame style"
     doc_ele_temp_a = "The electron temperature set by atomistic style"
+    doc_pimd_bead = "Bead index for PIMD, None for non-PIMD"
 
     return [
         Argument("conf_idx", list, optional=False, doc=doc_conf_idx, alias=["sys_idx"]),
@@ -108,6 +109,13 @@ def npt_task_group_args():
             default=None,
             doc=doc_ele_temp_a,
         ),
+        Argument(
+            "pimd_bead",
+            str,
+            optional=True,
+            default=None,
+            doc=doc_pimd_bead,
+        ),
     ]
 
 
@@ -117,6 +125,7 @@ def lmp_template_task_group_args():
     doc_revisions = "The revisions. Should be a dict providing the key - list of desired values pair. Key is the word to be replaced in the templates, and it may appear in both the lammps and plumed input templates. All values in the value list will be enmerated."
     doc_traj_freq = "The frequency of dumping configurations and thermodynamic states"
     doc_extra_pair_style_args = "The extra arguments for pair_style"
+    doc_pimd_bead = "Bead index for PIMD, None for non-PIMD"
 
     return [
         Argument("conf_idx", list, optional=False, doc=doc_conf_idx, alias=["sys_idx"]),
@@ -158,6 +167,13 @@ def lmp_template_task_group_args():
             default="",
             doc=doc_extra_pair_style_args,
         ),
+        Argument(
+            "pimd_bead",
+            str,
+            optional=True,
+            default=None,
+            doc=doc_pimd_bead,
+        ),
     ]
 
 

dpgen2/exploration/task/npt_task_group.py

@@ -49,6 +49,7 @@ def set_md(
         relative_v_epsilon: Optional[float] = None,
         ele_temp_f: Optional[float] = None,
         ele_temp_a: Optional[float] = None,
+        pimd_bead: Optional[str] = None,
     ):
         """
         Set MD parameters
@@ -72,6 +73,7 @@ def set_md(
         self.ele_temp_f = ele_temp_f
         self.ele_temp_a = ele_temp_a
         self.md_set = True
+        self.pimd_bead = pimd_bead
 
     def make_task(
         self,
@@ -131,6 +133,7 @@ def _make_lmp_task(
                 self.ele_temp_a,
                 self.no_pbc,
                 trj_seperate_files=False,
+                pimd_bead=self.pimd_bead,
             ),
         )
         return task

dpgen2/op/run_lmp.py

@@ -1,3 +1,4 @@
+import glob
 import json
 import logging
 import os
@@ -194,6 +195,7 @@ def execute(
                     with open("job.json", "w") as f:
                         json.dump(data, f, indent=4)
 
+        merge_pimd_files()
         ret_dict = {
             "log": work_dir / lmp_log_name,
             "traj": work_dir / lmp_traj_name,
@@ -356,3 +358,18 @@ def freeze_model(input_model, frozen_model, head=None):
             )
         )
         raise TransientError("freeze failed")
+
+
+def merge_pimd_files():
+    traj_files = glob.glob("traj.*.dump")
+    if len(traj_files) > 0:
+        with open(lmp_traj_name, "w") as f:
+            for traj_file in sorted(traj_files):
+                with open(traj_file, "r") as f2:
+                    f.write(f2.read())
+    model_devi_files = glob.glob("model_devi.*.out")
+    if len(model_devi_files) > 0:
+        with open(lmp_model_devi_name, "w") as f:
+            for model_devi_file in sorted(model_devi_files):
+                with open(model_devi_file, "r") as f2:
+                    f.write(f2.read())

tests/exploration/test_lmp_templ_task_group.py

@@ -197,6 +197,84 @@
 )
 
 
+in_lmp_pimd_template = textwrap.dedent(
+    """variable        NSTEPS          equal V_NSTEPS
+variable        THERMO_FREQ     equal 10
+variable        DUMP_FREQ       equal 10
+variable        TEMP            equal V_TEMP
+variable        PRES            equal 0.0
+variable        TAU_T           equal 0.100000
+variable        TAU_P           equal 0.500000
+variable        ibead           uloop 4 pad
+
+units           metal
+boundary        p p p
+atom_style      atomic
+
+neighbor        1.0 bin
+
+box             tilt large
+read_data       conf.lmp
+change_box      all triclinic
+mass            1 27.000000
+mass            2 24.000000
+
+pair_style      deepmd
+pair_coeff      * *
+
+thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz
+thermo          ${THERMO_FREQ}
+
+dump            dpgen_dump
+
+velocity        all create ${TEMP} 826513
+fix             1 all pimd/langevin ensemble npt integrator baoab temp ${TEMP} thermostat PILE_L 1234 tau ${TAU_T} iso ${PRES} barostat BZP taup ${TAU_P}
+
+timestep        0.001
+run             ${NSTEPS}
+"""
+)
+
+
+expected_lmp_pimd_template = textwrap.dedent(
+    """variable        NSTEPS          equal 1000
+variable        THERMO_FREQ     equal 10
+variable        DUMP_FREQ       equal 10
+variable        TEMP            equal 300
+variable        PRES            equal 0.0
+variable        TAU_T           equal 0.100000
+variable        TAU_P           equal 0.500000
+variable        ibead           uloop 4 pad
+
+units           metal
+boundary        p p p
+atom_style      atomic
+
+neighbor        1.0 bin
+
+box             tilt large
+read_data       conf.lmp
+change_box      all triclinic
+mass            1 27.000000
+mass            2 24.000000
+
+pair_style      deepmd model.000.pb model.001.pb model.002.pb model.003.pb out_freq 20 out_file model_devi.${ibead}.out
+pair_coeff      * *
+
+thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz
+thermo          ${THERMO_FREQ}
+
+dump            dpgen_dump all custom 20 traj.${ibead}.dump id type x y z
+
+velocity        all create ${TEMP} 826513
+fix             1 all pimd/langevin ensemble npt integrator baoab temp ${TEMP} thermostat PILE_L 1234 tau ${TAU_T} iso ${PRES} barostat BZP taup ${TAU_P}
+
+timestep        0.001
+run             ${NSTEPS}
+"""
+)
+
+
 class TestLmpTemplateTaskGroup(unittest.TestCase):
     def setUp(self):
         self.lmp_template_fname = Path("lmp.template")
@@ -215,11 +293,14 @@ def setUp(self):
         }
         self.rev_empty = {}
         self.traj_freq = 20
+        self.lmp_pimd_template_fname = Path("lmp.pimd.template")
+        self.lmp_pimd_template_fname.write_text(in_lmp_pimd_template)
 
     def tearDown(self):
         os.remove(self.lmp_template_fname)
         os.remove(self.lmp_plm_template_fname)
         os.remove(self.plm_template_fname)
+        os.remove(self.lmp_pimd_template_fname)
 
     def test_lmp(self):
         task_group = LmpTemplateTaskGroup()
@@ -333,3 +414,25 @@ def test_lmp_empty(self):
                 ee,
             )
             idx += 1
+
+    def test_lmp_pimd(self):
+        task_group = LmpTemplateTaskGroup()
+        task_group.set_conf(["foo"])
+        task_group.set_lmp(
+            self.numb_models,
+            self.lmp_pimd_template_fname,
+            revisions={"V_NSTEPS": [1000], "V_TEMP": [300]},
+            traj_freq=self.traj_freq,
+            pimd_bead="${ibead}",
+        )
+        task_group.make_task()
+        ngroup = len(task_group)
+        self.assertEqual(
+            ngroup,
+            1,
+        )
+        ee = expected_lmp_pimd_template.split("\n")
+        self.assertEqual(
+            task_group[0].files()[lmp_input_name].split("\n"),
+            ee,
+        )