model_devi: add support for pimd (#1366)

Add support for LAMMPS's fix pimd/langevin in model deviation tasks.

---------

Signed-off-by: Yifan Li李一帆 <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <[email protected]>

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),

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"

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:

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"