diff --git a/dpgen2/constants.py b/dpgen2/constants.py index fe1bffbf..6d5d0197 100644 --- a/dpgen2/constants.py +++ b/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" diff --git a/dpgen2/exploration/task/lmp/lmp_input.py b/dpgen2/exploration/task/lmp/lmp_input.py index 5898000f..c2a22b60 100644 --- a/dpgen2/exploration/task/lmp/lmp_input.py +++ b/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" diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py index b82e1695..1a44cb8e 100644 --- a/dpgen2/exploration/task/lmp_template_task_group.py +++ b/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,11 +47,13 @@ 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( @@ -56,8 +61,11 @@ def set_lmp( 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 diff --git a/dpgen2/exploration/task/make_task_group_from_config.py b/dpgen2/exploration/task/make_task_group_from_config.py index c467fd8e..3b793c58 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/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, + ), ] diff --git a/dpgen2/exploration/task/npt_task_group.py b/dpgen2/exploration/task/npt_task_group.py index 4c999638..27c1e001 100644 --- a/dpgen2/exploration/task/npt_task_group.py +++ b/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 diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 2822a325..3ff366c9 100644 --- a/dpgen2/op/run_lmp.py +++ b/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()) diff --git a/tests/exploration/test_lmp_templ_task_group.py b/tests/exploration/test_lmp_templ_task_group.py index 5ef67f72..6211137d 100644 --- a/tests/exploration/test_lmp_templ_task_group.py +++ b/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, + ) diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index b727fb76..5b7f4542 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -6,6 +6,7 @@ Path, ) +import dpdata import numpy as np from dflow.python import ( OP, @@ -35,6 +36,7 @@ from dpgen2.op.run_lmp import ( RunLmp, get_ele_temp, + merge_pimd_files, set_models, ) from dpgen2.utils import ( @@ -286,3 +288,64 @@ def test_get_ele_temp(self): def tearDown(self): if os.path.exists("log"): os.remove("log") + + +class TestMergePIMDFiles(unittest.TestCase): + def test_merge_pimd_files(self): + for i in range(1, 3): + with open("traj.%s.dump" % i, "w") as f: + f.write( + """ITEM: TIMESTEP +0 +ITEM: NUMBER OF ATOMS +3 +ITEM: BOX BOUNDS xy xz yz pp pp pp +0.0000000000000000e+00 1.2444661140399999e+01 0.0000000000000000e+00 +0.0000000000000000e+00 1.2444661140399999e+01 0.0000000000000000e+00 +0.0000000000000000e+00 1.2444661140399999e+01 0.0000000000000000e+00 +ITEM: ATOMS id type x y z +1 8 7.23489 0.826309 4.61669 +2 1 8.04419 0.520382 5.14395 +3 1 6.48126 0.446895 4.99766 +ITEM: TIMESTEP +10 +ITEM: NUMBER OF ATOMS +3 +ITEM: BOX BOUNDS xy xz yz pp pp pp +0.0000000000000000e+00 1.2444661140399999e+01 0.0000000000000000e+00 +0.0000000000000000e+00 1.2444661140399999e+01 0.0000000000000000e+00 +0.0000000000000000e+00 1.2444661140399999e+01 0.0000000000000000e+00 +ITEM: ATOMS id type x y z +1 8 7.23103 0.814939 4.59892 +2 1 7.96453 0.61699 5.19158 +3 1 6.43661 0.370311 5.09854 +""" + ) + for i in range(1, 3): + with open("model_devi.%s.out" % i, "w") as f: + f.write( + """# step max_devi_v min_devi_v avg_devi_v max_devi_f min_devi_f avg_devi_f + 0 9.023897e-17 3.548771e-17 5.237314e-17 8.196123e-16 1.225653e-16 3.941002e-16 + 10 1.081667e-16 4.141596e-17 7.534462e-17 9.070597e-16 1.067947e-16 4.153524e-16 +""" + ) + + merge_pimd_files() + self.assertTrue(os.path.exists(lmp_traj_name)) + self.assertTrue(os.path.exists(lmp_model_devi_name)) + s = dpdata.System(lmp_traj_name, fmt="lammps/dump") + assert len(s) == 4 + model_devi = np.loadtxt(lmp_model_devi_name) + assert model_devi.shape[0] == 4 + + def tearDown(self): + for f in [ + lmp_traj_name, + "traj.1.dump", + "traj.2.dump", + lmp_model_devi_name, + "model_devi.1.out", + "model_devi.2.out", + ]: + if os.path.exists(f): + os.remove(f)