feat: add input_extra_files for lmp-md and lmp-template

dpgen2/exploration/task/lmp_template_task_group.py

@@ -48,12 +48,21 @@ def set_lmp(
         traj_freq: int = 10,
         extra_pair_style_args: str = "",
         pimd_bead: Optional[str] = None,
+        input_extra_files: Optional[List[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
+        if input_extra_files is not None:
+            self.input_extra_files = [Path(ii).name for ii in input_extra_files]
+            self.input_extra_file_contents = [
+                Path(ii).read_text() for ii in input_extra_files
+            ]
+        else:
+            self.input_extra_files = []
+            self.input_extra_file_contents = []
         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(
@@ -134,6 +143,13 @@ def _make_lmp_task(
                 plm_input_name,
                 plm_cont,
             )
+
+        # Add extra files to the task
+        for file_name, file_content in zip(
+            self.input_extra_files, self.input_extra_file_contents
+        ):
+            task.add_file(file_name, file_content)
+
         return task
 
 

dpgen2/exploration/task/make_task_group_from_config.py

@@ -48,6 +48,9 @@ def npt_task_group_args():
     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"
+    doc_input_extra_files = (
+        "Extra files that may be needed during exploration (e.g., ZBL parameter files)"
+    )
 
     return [
         Argument("conf_idx", list, optional=False, doc=doc_conf_idx, alias=["sys_idx"]),
@@ -70,7 +73,7 @@ def npt_task_group_args():
             int,
             optional=True,
             default=10,
-            doc=doc_nsteps,
+            doc=doc_traj_freq,
             alias=["t_freq", "trj_freq", "traj_freq"],
         ),
         Argument("tau_t", float, optional=True, default=5e-2, doc=doc_tau_t),
@@ -116,6 +119,13 @@ def npt_task_group_args():
             default=None,
             doc=doc_pimd_bead,
         ),
+        Argument(
+            "input_extra_files",
+            list,
+            optional=True,
+            default=[],
+            doc=doc_input_extra_files,
+        ),
     ]
 
 
@@ -126,6 +136,9 @@ def lmp_template_task_group_args():
     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"
+    doc_input_extra_files = (
+        "Extra files that may be needed during exploration (e.g., ZBL parameter files)"
+    )
 
     return [
         Argument("conf_idx", list, optional=False, doc=doc_conf_idx, alias=["sys_idx"]),
@@ -174,6 +187,13 @@ def lmp_template_task_group_args():
             default=None,
             doc=doc_pimd_bead,
         ),
+        Argument(
+            "input_extra_files",
+            list,
+            optional=True,
+            default=[],
+            doc=doc_input_extra_files,
+        ),
     ]
 
 

dpgen2/exploration/task/npt_task_group.py

@@ -1,5 +1,8 @@
 import itertools
 import random
+from pathlib import (
+    Path,
+)
 from typing import (
     List,
     Optional,
@@ -50,6 +53,7 @@ def set_md(
         ele_temp_f: Optional[float] = None,
         ele_temp_a: Optional[float] = None,
         pimd_bead: Optional[str] = None,
+        input_extra_files: Optional[List[str]] = None,
     ):
         """
         Set MD parameters
@@ -72,6 +76,14 @@ def set_md(
         self.relative_v_epsilon = relative_v_epsilon
         self.ele_temp_f = ele_temp_f
         self.ele_temp_a = ele_temp_a
+        if input_extra_files is not None:
+            self.input_extra_files = [Path(ii).name for ii in input_extra_files]
+            self.input_extra_file_contents = [
+                Path(ii).read_text() for ii in input_extra_files
+            ]
+        else:
+            self.input_extra_files = []
+            self.input_extra_file_contents = []
         self.md_set = True
         self.pimd_bead = pimd_bead
 
@@ -136,4 +148,11 @@ def _make_lmp_task(
                 pimd_bead=self.pimd_bead,
             ),
         )
+
+        # Add extra files to the task
+        for file_name, file_content in zip(
+            self.input_extra_files, self.input_extra_file_contents
+        ):
+            task.add_file(file_name, file_content)
+
         return task