diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py index 1a44cb8e..d1f8e9fc 100644 --- a/dpgen2/exploration/task/lmp_template_task_group.py +++ b/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 diff --git a/dpgen2/exploration/task/make_task_group_from_config.py b/dpgen2/exploration/task/make_task_group_from_config.py index 3b793c58..2859ac8f 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/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, + ), ] diff --git a/dpgen2/exploration/task/npt_task_group.py b/dpgen2/exploration/task/npt_task_group.py index 27c1e001..ada5a15a 100644 --- a/dpgen2/exploration/task/npt_task_group.py +++ b/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