diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 5c6d5f49f..e1d220e42 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -262,6 +262,7 @@ def model_devi_jobs_args() -> list[Argument]: doc_press = "Pressure (Bar) in MD. Required when ensemble is npt." doc_trj_freq = "Frequecy of trajectory saved in MD." doc_nsteps = "Running steps of MD. It is not optional when not using a template." + doc_nbeads = "Number of beads in PIMD. If not given, classical MD will be performed. Only supported for LAMMPS version >= 20230615." doc_ensemble = "Determining which ensemble used in MD, options include “npt” and “nvt”. It is not optional when not using a template." doc_neidelay = "delay building until this many steps since last build." doc_taut = "Coupling time of thermostat (ps)." @@ -280,6 +281,7 @@ def model_devi_jobs_args() -> list[Argument]: Argument("press", list[float], optional=True, doc=doc_press), Argument("trj_freq", int, optional=False, doc=doc_trj_freq), Argument("nsteps", int, optional=True, doc=doc_nsteps), + Argument("nbeads", int, optional=True, doc=doc_nbeads), Argument("ensemble", str, optional=True, doc=doc_ensemble), Argument("neidelay", int, optional=True, doc=doc_neidelay), Argument("taut", float, optional=True, doc=doc_taut), diff --git a/dpgen/generator/lib/lammps.py b/dpgen/generator/lib/lammps.py index 0da9c8346..d9ff4a493 100644 --- a/dpgen/generator/lib/lammps.py +++ b/dpgen/generator/lib/lammps.py @@ -37,6 +37,7 @@ def make_lammps_input( max_seed=1000000, nopbc=False, deepmd_version="0.1", + nbeads=None, ): if (ele_temp_f is not None or ele_temp_a is not None) and Version( deepmd_version @@ -49,9 +50,22 @@ def make_lammps_input( "the frame style ele_temp and atom style ele_temp should not be set at the same time" ) ret = "variable NSTEPS equal %d\n" % nsteps + if nbeads is not None: + if nbeads <= 0: + raise ValueError( + "The number of beads should be positive. Check your nbeads setting." + ) + power = 1 + while power < nbeads: + power *= 10 + ret += "variable ibead uloop %d pad\n" % (power - 1) + if nbeads is not None: + ret += "atom_modify map yes\n" ret += "variable THERMO_FREQ equal %d\n" % trj_freq ret += "variable DUMP_FREQ equal %d\n" % trj_freq ret += "variable TEMP equal %f\n" % temp + if nbeads is not None: + ret += "variable TEMP_NBEADS equal %f\n" % (temp * nbeads) if ele_temp_f is not None: ret += "variable ELE_TEMP equal %f\n" % ele_temp_f if ele_temp_a is not None: @@ -72,10 +86,16 @@ def make_lammps_input( ret += "neigh_modify delay %d\n" % neidelay ret += "\n" ret += "box tilt large\n" - ret += ( - 'if "${restart} > 0" then "read_restart dpgen.restart.*" else "read_data %s"\n' - % conf_file - ) + if nbeads is None: + ret += ( + 'if "${restart} > 0" then "read_restart dpgen.restart.*" else "read_data %s"\n' + % conf_file + ) + else: + ret += ( + 'if "${restart} > 0" then "read_restart dpgen.restart${ibead}.*" else "read_data %s"\n' + % conf_file + ) ret += "change_box all triclinic\n" for jj in range(len(mass_map)): ret += "mass %d %f\n" % (jj + 1, mass_map[jj]) @@ -98,23 +118,43 @@ def make_lammps_input( keywords += "fparam ${ELE_TEMP}" if ele_temp_a is not None: keywords += "aparam ${ELE_TEMP}" - ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords}\n" + if nbeads is None: + ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords}\n" + else: + ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi${{ibead}}.out {keywords}\n" ret += "pair_coeff * *\n" ret += "\n" ret += "thermo_style custom step temp pe ke etotal press vol lx ly lz xy xz yz\n" ret += "thermo ${THERMO_FREQ}\n" model_devi_merge_traj = jdata.get("model_devi_merge_traj", False) - if model_devi_merge_traj is True: - ret += "dump 1 all custom ${DUMP_FREQ} all.lammpstrj id type x y z fx fy fz\n" - ret += 'if "${restart} > 0" then "dump_modify 1 append yes"\n' + if nbeads is None: + if model_devi_merge_traj is True: + ret += "dump 1 all custom ${DUMP_FREQ} all.lammpstrj id type x y z fx fy fz\n" + ret += 'if "${restart} > 0" then "dump_modify 1 append yes"\n' + else: + ret += "dump 1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n" else: - ret += "dump 1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n" - ret += "restart 10000 dpgen.restart\n" + if model_devi_merge_traj is True: + ret += "dump 1 all custom ${DUMP_FREQ} all.lammpstrj${ibead} id type x y z fx fy fz\n" + ret += 'if "${restart} > 0" then "dump_modify 1 append yes"\n' + else: + ret += "dump 1 all custom ${DUMP_FREQ} traj/*.lammpstrj${ibead} id type x y z fx fy fz\n" + if nbeads is None: + ret += "restart 10000 dpgen.restart\n" + else: + ret += "restart 10000 dpgen.restart${ibead}\n" ret += "\n" if pka_e is None: - ret += 'if "${restart} == 0" then "velocity all create ${TEMP} %d"' % ( - random.randrange(max_seed - 1) + 1 - ) + if nbeads is None: + ret += ( + 'if "${restart} == 0" then "velocity all create ${TEMP} %d"' + % (random.randrange(max_seed - 1) + 1) + ) + else: + ret += ( + 'if "${restart} == 0" then "velocity all create ${TEMP_NBEADS} %d"' + % (random.randrange(max_seed - 1) + 1) + ) else: sys = dpdata.System(conf_file, fmt="lammps/lmp") sys_data = sys.data @@ -140,18 +180,34 @@ def make_lammps_input( assert pres is not None if nopbc: raise RuntimeError("ensemble %s is conflicting with nopbc" % ensemble) - if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso": - ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P}\n" - elif ensemble == "npt-a" or ensemble == "npt-aniso": - ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} aniso ${PRES} ${PRES} ${TAU_P}\n" - elif ensemble == "npt-t" or ensemble == "npt-tri": - ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} tri ${PRES} ${PRES} ${TAU_P}\n" - elif ensemble == "nvt": - ret += "fix 1 all nvt temp ${TEMP} ${TEMP} ${TAU_T}\n" - elif ensemble == "nve": - ret += "fix 1 all nve\n" + if nbeads is None: + if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso": + ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P}\n" + elif ensemble == "npt-a" or ensemble == "npt-aniso": + ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} aniso ${PRES} ${PRES} ${TAU_P}\n" + elif ensemble == "npt-t" or ensemble == "npt-tri": + ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} tri ${PRES} ${PRES} ${TAU_P}\n" + elif ensemble == "nvt": + ret += "fix 1 all nvt temp ${TEMP} ${TEMP} ${TAU_T}\n" + elif ensemble == "nve": + ret += "fix 1 all nve\n" + else: + raise RuntimeError("unknown emsemble " + ensemble) else: - raise RuntimeError("unknown emsemble " + ensemble) + if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso": + ret += "fix 1 all pimd/langevin fmmode physical ensemble npt integrator obabo thermostat PILE_L ${ibead} temp ${TEMP} tau ${TAU_T} scale 1.0 barostat BZP iso ${PRES} taup ${TAU_P}\n" + elif ensemble == "npt-a" or ensemble == "npt-aniso": + ret += "fix 1 all pimd/langevin fmmode physical ensemble npt integrator obabo thermostat PILE_L ${ibead} temp ${TEMP} tau ${TAU_T} scale 1.0 barostat BZP aniso ${PRES} taup ${TAU_P}\n" + elif ensemble == "nvt": + ret += "fix 1 all pimd/langevin fmmode physical ensemble nvt integrator obabo thermostat PILE_L ${ibead} temp ${TEMP} tau ${TAU_T} scale 1.0\n" + elif ensemble == "nve": + ret += "fix 1 all pimd/langevin fmmode physical ensemble nve integrator obabo temp ${TEMP}\n" + else: + raise RuntimeError( + "unknown emsemble " + + ensemble + + " for fix pimd/langevin\nrefer to https://docs.lammps.org/fix_pimd.html for more information" + ) if nopbc: ret += "velocity all zero linear\n" ret += "fix fm all momentum 1 linear 1 1 1\n" diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index da1454244..034918efc 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -18,6 +18,7 @@ import os import queue import random +import re import shutil import sys import warnings @@ -915,7 +916,11 @@ def parse_cur_job(cur_job): dt = _get_param_alias(cur_job, ["dt"]) else: dt = None - return ensemble, nsteps, trj_freq, temps, press, pka_e, dt + if "nbeads" in cur_job: + nbeads = _get_param_alias(cur_job, ["nbeads"]) + else: + nbeads = None + return ensemble, nsteps, trj_freq, temps, press, pka_e, dt, nbeads def expand_matrix_values(target_list, cur_idx=0): @@ -1457,7 +1462,21 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems): if iter_index >= len(model_devi_jobs): return False cur_job = model_devi_jobs[iter_index] - ensemble, nsteps, trj_freq, temps, press, pka_e, dt = parse_cur_job(cur_job) + ensemble, nsteps, trj_freq, temps, press, pka_e, dt, nbeads = parse_cur_job(cur_job) + model_devi_f_avg_relative = jdata.get("model_devi_f_avg_relative", False) + model_devi_merge_traj = jdata.get("model_devi_merge_traj", False) + if (nbeads is not None) and model_devi_f_avg_relative: + raise RuntimeError( + "model_devi_f_avg_relative has not been supported for pimd. Set model_devi_f_avg_relative to False." + ) + if (nbeads is not None) and (model_devi_merge_traj): + raise RuntimeError( + "model_devi_merge_traj has not been supported for pimd. Set model_devi_merge_traj to False." + ) + if (nbeads is not None) and (not (nsteps % trj_freq == 0)): + raise RuntimeError( + "trj_freq should be a factor of nsteps for pimd. Please check your input." + ) if dt is not None: model_devi_dt = dt sys_idx = expand_idx(cur_job["sys_idx"]) @@ -1560,6 +1579,7 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems): ele_temp_a=te_a, nopbc=nopbc, deepmd_version=deepmd_version, + nbeads=nbeads, ) job = {} job["ensemble"] = ensemble @@ -1916,12 +1936,23 @@ def run_md_model_devi(iter_index, jdata, mdata): model_devi_engine = jdata.get("model_devi_engine", "lammps") if model_devi_engine == "lammps": - command = f"{{ if [ ! -f dpgen.restart.10000 ]; then {model_devi_exec} -i input.lammps -v restart 0; else {model_devi_exec} -i input.lammps -v restart 1; fi }}" - command = "/bin/sh -c '%s'" % command + nbeads = jdata["model_devi_jobs"][iter_index].get("nbeads") + if nbeads is None: + command = f"{{ if [ ! -f dpgen.restart.10000 ]; then {model_devi_exec} -i input.lammps -v restart 0; else {model_devi_exec} -i input.lammps -v restart 1; fi }}" + else: + command = f"{{ all_exist=true; for i in $(seq -w 1 {nbeads}); do [[ ! -f dpgen.restart${{i}}.10000 ]] && {{ all_exist=false; break; }}; done; $all_exist && {{ {model_devi_exec} -p {nbeads}x1 -i input.lammps -v restart 1; }} || {{ {model_devi_exec} -p {nbeads}x1 -i input.lammps -v restart 0; }} }}" + command = "/bin/bash -c '%s'" % command commands = [command] forward_files = ["conf.lmp", "input.lammps"] - backward_files = ["model_devi.out", "model_devi.log"] + backward_files = ["model_devi.log"] + if nbeads is None: + backward_files += ["model_devi.out"] + else: + num_digits = np.ceil(np.log10(nbeads + 1)).astype(int) + backward_files += [ + f"model_devi{i+1:0{num_digits}d}.out" for i in range(nbeads) + ] if model_devi_merge_traj: backward_files += ["all.lammpstrj"] else: @@ -2131,6 +2162,75 @@ def _read_model_devi_file( model_devi_f_avg_relative: bool = False, model_devi_merge_traj: bool = False, ): + model_devi_files = glob.glob(os.path.join(task_path, "model_devi*.out")) + model_devi_files_sorted = sorted( + model_devi_files, key=lambda x: int(re.search(r"(\d+)", x).group(1)) + ) + if len(model_devi_files_sorted) > 1: + with open(model_devi_files_sorted[0]) as f: + first_line = f.readline() + if not (first_line.startswith("#")): + first_line = "#" + num_beads = len(model_devi_files_sorted) + model_devi_contents = [] + for file in model_devi_files_sorted: + model_devi_contents.append(np.loadtxt(file)) + assert all( + model_devi_content.shape[0] == model_devi_contents[0].shape[0] + for model_devi_content in model_devi_contents + ), "Not all beads generated the same number of lines in the model_devi$\{ibead\}.out file. Check your pimd task carefully." + for file in model_devi_files_sorted: + os.remove(file) + last_step = model_devi_contents[0][-1, 0] + for ibead in range(1, num_beads): + model_devi_contents[ibead][:, 0] = model_devi_contents[ibead][ + :, 0 + ] + ibead * (last_step + 1) + model_devi = np.concatenate(model_devi_contents, axis=0) + num_columns = model_devi.shape[1] + formats = ["%12d"] + ["%22.6e"] * (num_columns - 1) + np.savetxt( + os.path.join(task_path, "model_devi.out"), + model_devi, + fmt=formats, + header=first_line.rstrip(), + comments="", + ) + + if not model_devi_merge_traj: + num_digits = np.ceil(np.log10(num_beads + 1)).astype(int) + traj_files_sorted = [] + for ibead in range(num_beads): + traj_files = glob.glob( + os.path.join( + task_path, "traj", f"*lammpstrj{ibead+1:0{num_digits}d}" + ) + ) + traj_files_sorted.append( + sorted( + traj_files, + key=lambda x: int( + re.search(r"^(\d+)\.lammpstrj", os.path.basename(x)).group( + 1 + ) + ), + ) + ) + assert all( + len(traj_list) == len(traj_files_sorted[0]) + for traj_list in traj_files_sorted + ), "Not all beads generated the same number of frames. Check your pimd task carefully." + for ibead in range(num_beads): + for itraj in range(len(traj_files_sorted[0])): + base_path, original_filename = os.path.split( + traj_files_sorted[ibead][itraj] + ) + frame_number = int(original_filename.split(".")[0]) + new_filename = os.path.join( + base_path, + f"{frame_number + ibead * (int(last_step)+1):d}.lammpstrj", + ) + os.rename(traj_files_sorted[ibead][itraj], new_filename) model_devi = np.loadtxt(os.path.join(task_path, "model_devi.out")) if model_devi_f_avg_relative: if model_devi_merge_traj is True: diff --git a/tests/generator/context.py b/tests/generator/context.py index 7a6f1f2d3..a58b0dddb 100644 --- a/tests/generator/context.py +++ b/tests/generator/context.py @@ -20,6 +20,7 @@ from dpgen.util import setup_ele_temp # noqa: F401 param_file = "param-mg-vasp.json" +param_pimd_file = "param-mg-pimd-vasp.json" param_file_merge_traj = "param-mg-vasp_merge_traj.json" param_file_v1 = "param-mg-vasp-v1.json" param_file_v1_et = "param-mg-vasp-v1-et.json" diff --git a/tests/generator/param-mg-pimd-vasp.json b/tests/generator/param-mg-pimd-vasp.json new file mode 100644 index 000000000..efd75dd5c --- /dev/null +++ b/tests/generator/param-mg-pimd-vasp.json @@ -0,0 +1,153 @@ +{ + "type_map": [ + "Mg", + "Al" + ], + "mass_map": [ + 24, + 27 + ], + "init_data_prefix": "data", + "init_data_sys": [ + "deepmd" + ], + "init_batch_size": [ + 16 + ], + "sys_configs": [ + [ + "data/mg.fcc.02x02x02/01.scale_pert/sys-0032/scale*/000000/POSCAR" + ], + [ + "data/mg.fcc.02x02x02/01.scale_pert/sys-0032/scale*/000001/POSCAR" + ] + ], + "_comment1": "0 1 2 3", + "_comment2": "4 5 6 7", + "sys_batch_size": [ + 1, + 1 + ], + "_comment3": " 00.train ", + "numb_models": 4, + "default_training_param": { + "model": { + "descriptor": { + "seed": 0, + "type": "se_a", + "sel": [ + 90 + ], + "rcut": 6.0, + "rcut_smth": 2.0, + "neuron": [ + 25, + 50, + 100 + ], + "axis_neuron": 12, + "resnet_dt": false + }, + "fitting_net": { + "seed": 0, + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true + } + }, + "learning_rate": { + "type": "exp", + "decay_steps": 2000, + "start_lr": 0.002, + "stop_lr": 7.010533249765748e-08 + }, + "loss": { + "start_pref_e": 0.02, + "limit_pref_e": 2, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0.0, + "limit_pref_v": 0.0 + }, + "training": { + "seed": 0, + "stop_batch": 400000, + "disp_file": "lcurve.out", + "disp_freq": 2000, + "save_freq": 20000, + "save_ckpt": "model.ckpt", + "disp_training": true, + "time_training": true, + "profiling": false, + "training_data": { + "systems": [], + "set_prefix": "set", + "batch_size": 1 + } + } + }, + "_comment9": " 01.model_devi ", + "_comment10": "model_devi_skip: the first x of the recorded frames", + "model_devi_dt": 0.002, + "model_devi_skip": 0, + "model_devi_f_trust_lo": 0.05, + "model_devi_f_trust_hi": 0.15, + "model_devi_clean_traj": false, + "model_devi_jobs": [ + { + "sys_idx": [ + 0, + 1 + ], + "temps": [ + 50, + 100 + ], + "press": [ + 1.0, + 2.0 + ], + "nbeads": 4, + "trj_freq": 10, + "nsteps": 10, + "ensemble": "npt", + "_idx": "00" + } + ], + "_comment11": " 02.fp ", + "fp_style": "vasp", + "shuffle_poscar": false, + "fp_task_max": 100, + "fp_task_min": 10, + "fp_pp_path": ".", + "fp_pp_files": [ + "vasp/potcars/POTCAR.mg", + "vasp/potcars/POTCAR.al" + ], + "_comment12": " user provided vasp script ", + "user_fp_params": { + "PREC": "A", + "ENCUT": 600, + "ISYM": 0, + "ALGO": "fast", + "EDIFF": 1e-05, + "LREAL": "A", + "NPAR": 1, + "KPAR": 1, + "NELMIN": 4, + "ISIF": 2, + "ISMEAR": 1, + "SIGMA": 0.25, + "IBRION": -1, + "NSW": 0, + "LWAVE": false, + "LCHARG": false, + "PSTRESS": 0, + "KSPACING": 0.16, + "KGAMMA": false + }, + "_comment13": " that's all " +} diff --git a/tests/generator/test_make_md.py b/tests/generator/test_make_md.py index c6138ee52..bf09508ee 100644 --- a/tests/generator/test_make_md.py +++ b/tests/generator/test_make_md.py @@ -7,25 +7,31 @@ import unittest import dpdata +import numpy as np -from dpgen.generator.run import parse_cur_job_sys_revmat +from dpgen.generator.run import _read_model_devi_file, parse_cur_job_sys_revmat sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) __package__ = "generator" +import tempfile + from .comp_sys import test_atom_names, test_atom_types, test_cell, test_coord from .context import ( find_only_one_key, machine_file, + machine_file_v1, make_model_devi, my_file_cmp, param_amber_file, param_file, + param_pimd_file, parse_cur_job, parse_cur_job_revmat, revise_by_keys, revise_lmp_input_dump, revise_lmp_input_model, revise_lmp_input_plm, + run_model_devi, ) @@ -95,7 +101,7 @@ def _check_traj_dir(testCase, idx): def _get_lammps_pt(lmp_input): with open(lmp_input) as fp: for ii in fp: - if "variable" in ii and "TEMP" in ii: + if "variable" in ii and "TEMP" in ii and "TEMP_NBEADS" not in ii: lt = float(ii.split()[3]) if "variable" in ii and "PRES" in ii: lp = float(ii.split()[3]) @@ -107,7 +113,7 @@ def _check_pt(testCase, idx, jdata): tasks = glob.glob(os.path.join(md_dir, "task.*")) tasks.sort() cur_job = jdata["model_devi_jobs"][idx] - ensemble, nsteps, trj_freq, temps, press, pka_e, dt = parse_cur_job(cur_job) + ensemble, nsteps, trj_freq, temps, press, pka_e, dt, nbeads = parse_cur_job(cur_job) testCase.assertTrue(ensemble, "npt") # get poscars sys_idx = cur_job["sys_idx"] @@ -136,6 +142,8 @@ class TestMakeModelDevi(unittest.TestCase): def tearDown(self): if os.path.isdir("iter.000000"): shutil.rmtree("iter.000000") + if os.path.isdir("test_model_devi_pimd"): + shutil.rmtree("test_model_devi_pimd") def test_make_model_devi(self): if os.path.isdir("iter.000000"): @@ -152,6 +160,20 @@ def test_make_model_devi(self): _check_pt(self, 0, jdata) # shutil.rmtree('iter.000000') + def test_make_model_devi_pimd(self): + if os.path.isdir("iter.000000"): + shutil.rmtree("iter.000000") + with open(param_pimd_file) as fp: + jdata = json.load(fp) + with open(machine_file) as fp: + mdata = json.load(fp) + _make_fake_models(0, jdata["numb_models"]) + make_model_devi(0, jdata, mdata) + _check_pb(self, 0) + _check_confs(self, 0, jdata) + _check_traj_dir(self, 0) + _check_pt(self, 0, jdata) + def test_make_model_devi_nopbc_npt(self): if os.path.isdir("iter.000000"): shutil.rmtree("iter.000000") @@ -183,6 +205,66 @@ def test_make_model_devi_nopbc_nvt(self): _check_pt(self, 0, jdata) # shutil.rmtree('iter.000000') + def test_run_model_devi_pimd(self): + if os.path.isdir("iter.000000"): + shutil.rmtree("iter.000000") + with open(param_pimd_file) as fp: + jdata = json.load(fp) + with open(machine_file_v1) as fp: + mdata = json.load(fp) + _make_fake_models(0, jdata["numb_models"]) + make_model_devi(0, jdata, mdata) + with tempfile.TemporaryDirectory() as remote_root: + run_model_devi( + 0, + jdata, + { + "api_version": "1.0", + "model_devi_command": ( + "touch model_devi1.out model_devi2.out model_devi3.out model_devi4.out" + "&& echo lmp" + ), + "model_devi_machine": { + "batch_type": "shell", + "local_root": "./", + "remote_root": remote_root, + "context_type": "local", + }, + "model_devi_group_size": 1, + "model_devi_resources": { + "group_size": 1, + "cpu_per_node": 4, + }, + }, + ) + + def test_read_model_devi_file_pimd(self): + path = "test_model_devi_pimd" + if os.path.isdir(path): + shutil.rmtree(path) + os.makedirs(path, exist_ok=True) + os.makedirs(os.path.join(path, "traj"), exist_ok=True) + for i in range(4): + for j in range(0, 5, 2): + with open(os.path.join(path, f"traj/{j}.lammpstrj{i+1}"), "a"): + pass + model_devi_array = np.zeros([3, 7]) + model_devi_array[:, 0] = np.array([0, 2, 4]) + for i in range(4): + np.savetxt( + os.path.join(path, f"model_devi{i+1}.out"), model_devi_array, fmt="%d" + ) + _read_model_devi_file(path) + model_devi_total_array = np.zeros([12, 7]) + total_steps = np.array([0, 2, 4, 5, 7, 9, 10, 12, 14, 15, 17, 19]) + model_devi_total_array[:, 0] = total_steps + model_devi_out = np.loadtxt(os.path.join(path, "model_devi.out")) + np.testing.assert_array_almost_equal(model_devi_out, model_devi_total_array) + for istep in total_steps: + self.assertTrue( + os.path.isfile(os.path.join(path, f"traj/{istep}.lammpstrj")) + ) + class TestMakeModelDeviRevMat(unittest.TestCase): def tearDown(self):