Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history 
for more information, see https://pre-commit.ci
  • Loading branch information
@pre-commit-ci
pre-commit-ci[bot] committed Oct 28, 2023
1 parent 6cfca84 commit 84aed7b
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 23 deletions.
28 changes: 18 additions & 10 deletions dpgen/generator/lib/lammps.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,18 +52,20 @@ def make_lammps_input(
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.")
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)
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)
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:
Expand Down Expand Up @@ -93,7 +95,7 @@ def make_lammps_input(
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])
Expand Down Expand Up @@ -144,13 +146,15 @@ def make_lammps_input(
ret += "\n"
if pka_e is None:
if nbeads is None:
ret += 'if "${restart} == 0" then "velocity all create ${TEMP} %d"' % (
random.randrange(max_seed - 1) + 1
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
)
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
Expand Down Expand Up @@ -199,7 +203,11 @@ def make_lammps_input(
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")
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"
Expand Down
59 changes: 46 additions & 13 deletions dpgen/generator/run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2148,36 +2148,69 @@ def _read_model_devi_file(
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)))
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], "r") as f:
with open(model_devi_files_sorted[0]) as f:
first_line = f.readline()
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."
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_contents[ibead][:, 0] = model_devi_contents[ibead][
:, 0
] + ibead * (last_step + 1)
model_devi = np.concatenate(model_devi_contents, axis=0)
np.savetxt(os.path.join(task_path, "model_devi.out"), model_devi, fmt="%16.6e", header=first_line, comments="")

np.savetxt(
os.path.join(task_path, "model_devi.out"),
model_devi,
fmt="%16.6e",
header=first_line,
comments="",
)

if not model_devi_merge_traj:
num_digits = np.ceil(np.log10(num_beads+1)).astype(int)
num_digits = np.ceil(np.log10(num_beads + 1)).astype(int)
traj_files_sorted = []
for ibead in range(num_beads):
# pattern = os.path.join(task_path, f"*lammpstrj{N_str}")
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."
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 * (last_step+1)}.lammpstrj")
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 * (last_step+1)}.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:
Expand Down

0 comments on commit 84aed7b

Please sign in to comment.