From 55162c202ea22832ad552ddd126db717768fe0b5 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 25 Aug 2025 10:50:23 +0800 Subject: [PATCH 1/2] feat: add input_extra_files for lmp-md and lmp-template --- .../task/lmp_template_task_group.py | 12 +++++++++ .../task/make_task_group_from_config.py | 26 ++++++++++++++----- dpgen2/exploration/task/npt_task_group.py | 15 +++++++++++ 3 files changed, 46 insertions(+), 7 deletions(-) diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py index 1a44cb8e..9185523c 100644 --- a/dpgen2/exploration/task/lmp_template_task_group.py +++ b/dpgen2/exploration/task/lmp_template_task_group.py @@ -48,12 +48,19 @@ 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 +141,11 @@ 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..ac32ebdb 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -48,6 +48,7 @@ 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"]), @@ -60,9 +61,7 @@ def npt_task_group_args(): ), Argument("temps", list, optional=False, doc=doc_temps, alias=["Ts"]), Argument("press", list, optional=True, doc=doc_press, alias=["Ps"]), - Argument( - "ens", str, optional=True, default="nve", doc=doc_ens, alias=["ensemble"] - ), + Argument("ens", str, optional=True, default="nve", doc=doc_ens, alias=["ensemble"]), Argument("dt", float, optional=True, default=1e-3, doc=doc_dt), Argument("nsteps", int, optional=True, default=100, doc=doc_nsteps), Argument( @@ -70,7 +69,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), @@ -78,9 +77,7 @@ def npt_task_group_args(): Argument("pka_e", float, optional=True, default=None, doc=doc_pka_e), Argument("neidelay", int, optional=True, default=None, doc=doc_neidelay), Argument("no_pbc", bool, optional=True, default=False, doc=doc_no_pbc), - Argument( - "use_clusters", bool, optional=True, default=False, doc=doc_use_clusters - ), + Argument("use_clusters", bool, optional=True, default=False, doc=doc_use_clusters), Argument( "relative_f_epsilon", float, @@ -116,6 +113,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 +130,7 @@ 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 +179,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..d01b92ed 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,12 @@ 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 +146,9 @@ 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 From e7ce9968ccd497ff725aeb409ee167a9b3121a40 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 Aug 2025 04:18:25 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../exploration/task/lmp_template_task_group.py | 12 ++++++++---- .../task/make_task_group_from_config.py | 16 ++++++++++++---- dpgen2/exploration/task/npt_task_group.py | 8 ++++++-- 3 files changed, 26 insertions(+), 10 deletions(-) diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py index 9185523c..d1f8e9fc 100644 --- a/dpgen2/exploration/task/lmp_template_task_group.py +++ b/dpgen2/exploration/task/lmp_template_task_group.py @@ -57,7 +57,9 @@ def set_lmp( 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] + 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 = [] @@ -141,11 +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): + 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 ac32ebdb..2859ac8f 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -48,7 +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)" + 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"]), @@ -61,7 +63,9 @@ def npt_task_group_args(): ), Argument("temps", list, optional=False, doc=doc_temps, alias=["Ts"]), Argument("press", list, optional=True, doc=doc_press, alias=["Ps"]), - Argument("ens", str, optional=True, default="nve", doc=doc_ens, alias=["ensemble"]), + Argument( + "ens", str, optional=True, default="nve", doc=doc_ens, alias=["ensemble"] + ), Argument("dt", float, optional=True, default=1e-3, doc=doc_dt), Argument("nsteps", int, optional=True, default=100, doc=doc_nsteps), Argument( @@ -77,7 +81,9 @@ def npt_task_group_args(): Argument("pka_e", float, optional=True, default=None, doc=doc_pka_e), Argument("neidelay", int, optional=True, default=None, doc=doc_neidelay), Argument("no_pbc", bool, optional=True, default=False, doc=doc_no_pbc), - Argument("use_clusters", bool, optional=True, default=False, doc=doc_use_clusters), + Argument( + "use_clusters", bool, optional=True, default=False, doc=doc_use_clusters + ), Argument( "relative_f_epsilon", float, @@ -130,7 +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)" + 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"]), diff --git a/dpgen2/exploration/task/npt_task_group.py b/dpgen2/exploration/task/npt_task_group.py index d01b92ed..ada5a15a 100644 --- a/dpgen2/exploration/task/npt_task_group.py +++ b/dpgen2/exploration/task/npt_task_group.py @@ -78,7 +78,9 @@ def set_md( 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] + 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 = [] @@ -148,7 +150,9 @@ def _make_lmp_task( ) # Add extra files to the task - for file_name, file_content in zip(self.input_extra_files, self.input_extra_file_contents): + 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