diff --git a/dpti/equi.py b/dpti/equi.py index 504d8832..92bd93ff 100755 --- a/dpti/equi.py +++ b/dpti/equi.py @@ -561,6 +561,16 @@ def _compute_thermo(lmplog, natoms, stat_skip, stat_bsize): return thermo_info +def _get_task_model_file(task_name): + settings_file = os.path.join(task_name, "equi_settings.json") + if not os.path.isfile(settings_file): + return "graph.pb" + with open(settings_file) as fp: + settings = json.load(fp) + model = settings.get("model") + return os.path.basename(model) if model else None + + def _print_thermo_info(info, more_head=""): ptr = f"# thermodynamics {'value':>20s} {'err':>20s} {more_head}\n" ptr += f"# E [eV]: {info['e']:20.8f} {info['e_err']:20.8f}\n" @@ -642,12 +652,16 @@ def run_task(task_name, machine_file): resources=resources, machine=machine, ) + model_file = _get_task_model_file(task_name) + forward_files = ["in.lammps", "*.lmp"] + if model_file: + forward_files.append(model_file) task_list = [ Task( command=f"{mdata['command']} -in in.lammps", task_work_path=ii, - forward_files=["in.lammps", "*.lmp", "graph.pb"], + forward_files=forward_files, backward_files=["log*", "dump.equi", "out.lmp"], ) for ii in task_dir_list diff --git a/dpti/gdi.py b/dpti/gdi.py index 63027977..683b853e 100755 --- a/dpti/gdi.py +++ b/dpti/gdi.py @@ -21,6 +21,7 @@ from dpti.lib.utils import ( create_path, get_first_matched_key_from_dict, + get_model_filename, relative_link_file, ) from dpti.ti import _gen_lammps_input @@ -108,9 +109,9 @@ def _make_tasks_onephase( os.symlink(os.path.relpath(conf_file), "conf.lmp") local_graph_file = graph_file if graph_file: - if not os.path.exists("graph.pb"): - os.symlink(os.path.relpath(graph_abs_file), "graph.pb") - local_graph_file = "graph.pb" + local_graph_file = os.path.basename(graph_file) + if not os.path.exists(local_graph_file): + os.symlink(os.path.relpath(graph_abs_file), local_graph_file) if if_meam: relative_link_file(meam_model["library_abs_path"], "./") @@ -161,9 +162,11 @@ def _has_phase_specific_model(jdata): def _get_phase_graph_file(jdata, phase_idx): + phase_key = "phase_i" if phase_idx == 0 else "phase_ii" + phase_model = _get_phase_model(jdata, phase_key) if _has_phase_specific_model(jdata): - return f"graph.{phase_idx}.pb" - return "graph.pb" + return get_model_filename(phase_model, prefix=f"graph.{phase_idx}") + return get_model_filename(phase_model) def _setup_dpdt(task_path, jdata): @@ -319,15 +322,20 @@ def make_dpdt( resources = Resources.load_from_dict(mdata["resources"]) command = "lmp -i in.lammps" - forward_files = ["conf.lmp", "in.lammps", "graph.pb"] if if_meam: meam_library_basename = os.path.basename(meam_model["library"]) meam_potential_basename = os.path.basename(meam_model["potential"]) - forward_files.extend([meam_library_basename, meam_potential_basename]) backward_files = ["log.lammps", "final.lmp"] task_list = [] for ii in range(2): + forward_files = [ + "conf.lmp", + "in.lammps", + _get_phase_graph_file(jdata, ii), + ] + if if_meam: + forward_files.extend([meam_library_basename, meam_potential_basename]) task = Task( command=command, task_work_path=f"{ii}/", diff --git a/dpti/hti.py b/dpti/hti.py old mode 100755 new mode 100644 index 0477b6f5..4ebd1f59 --- a/dpti/hti.py +++ b/dpti/hti.py @@ -3,6 +3,7 @@ import glob import json import os +import shlex import shutil import numpy as np @@ -22,6 +23,7 @@ compute_nrefine, create_path, get_first_matched_key_from_dict, + get_model_filename, get_task_file_abspath, integrate_range_hti, parse_seq, @@ -551,6 +553,7 @@ def make_tasks(iter_name, jdata, ref="einstein", switch="one-step", if_meam=None equi_conf = os.path.abspath(jdata["equi_conf"]) meam_model = jdata.get("meam_model", None) model = os.path.abspath(jdata["model"]) + model_file = get_model_filename(model) if if_meam is None: if_meam = jdata.get("if_meam", None) @@ -575,7 +578,7 @@ def make_tasks(iter_name, jdata, ref="einstein", switch="one-step", if_meam=None copied_conf = os.path.join(os.path.abspath(iter_name), "conf.lmp") shutil.copyfile(equi_conf, copied_conf) jdata["equi_conf"] = "conf.lmp" - linked_model = os.path.join(os.path.abspath(iter_name), "graph.pb") + linked_model = os.path.join(os.path.abspath(iter_name), model_file) if if_meam: relative_link_file(meam_model["library"], job_abs_dir) @@ -584,7 +587,7 @@ def make_tasks(iter_name, jdata, ref="einstein", switch="one-step", if_meam=None pass shutil.copyfile(model, linked_model) - jdata["model"] = "graph.pb" + jdata["model"] = model_file cwd = os.getcwd() os.chdir(iter_name) with open("in.json", "w") as fp: @@ -703,6 +706,7 @@ def _make_tasks( equi_conf = os.path.abspath(equi_conf) model = jdata["model"] model = os.path.abspath(model) + model_file = get_model_filename(model) # mass_map = jdata['mass_map'] mass_map = get_first_matched_key_from_dict(jdata, ["mass_map", "model_mass_map"]) nsteps = jdata["nsteps"] @@ -767,15 +771,15 @@ def _make_tasks( os.symlink(os.path.relpath(equi_conf), "conf.lmp") os.chdir(cwd) jdata["equi_conf"] = "conf.lmp" - linked_model = os.path.join(os.path.abspath(iter_name), "graph.pb") + linked_model = os.path.join(os.path.abspath(iter_name), model_file) if not link: shutil.copyfile(model, linked_model) else: cwd = os.getcwd() os.chdir(iter_name) - os.symlink(os.path.relpath(model), "graph.pb") + os.symlink(os.path.relpath(model), model_file) os.chdir(cwd) - jdata["model"] = "graph.pb" + jdata["model"] = model_file langevin = jdata.get("langevin", True) cwd = os.getcwd() @@ -789,7 +793,7 @@ def _make_tasks( create_path(work_path) os.chdir(work_path) os.symlink(os.path.relpath(copied_conf), "conf.lmp") - os.symlink(os.path.relpath(linked_model), "graph.pb") + os.symlink(os.path.relpath(linked_model), model_file) if if_meam: meam_library_basename = os.path.basename(meam_model["library"]) meam_potential_basename = os.path.basename(meam_model["potential"]) @@ -813,7 +817,7 @@ def _make_tasks( "conf.lmp", mass_map, ii, - "graph.pb", + model_file, m_spring_k, nsteps, timestep, @@ -1494,12 +1498,22 @@ def _is_completed_lammps_task(task_work_path): return False -def _graph_link_command(task_dir, job_work_dir): +def _get_task_model_file(task_dir): + in_json = os.path.join(task_dir, "in.json") + if not os.path.isfile(in_json): + return "graph.pb" + with open(in_json) as fp: + jdata = json.load(fp) + model = jdata.get("model", "graph.pb") + return os.path.basename(model) if model else None + + +def _graph_link_command(task_dir, job_work_dir, model_file="graph.pb"): graph_relpath = os.path.relpath( - os.path.join(task_dir, "graph.pb"), + os.path.join(task_dir, model_file), os.path.join(job_work_dir, "task.000000"), ) - return f"ln -s {graph_relpath} graph.pb" + return f"ln -sf {shlex.quote(graph_relpath)} {shlex.quote(model_file)}" def run_task(task_dir, machine_file, task_name, no_dp=False): @@ -1528,15 +1542,13 @@ def run_task(task_dir, machine_file, task_name, no_dp=False): resources=resources, machine=machine, ) + model_file = _get_task_model_file(task_dir) - command = ( - f"{mdata['command']} -i in.lammps" - if no_dp - else ( - f"{_graph_link_command(task_dir, job_work_dir)}; " - f"{mdata['command']} -i in.lammps" + command = f"{mdata['command']} -i in.lammps" + if not no_dp and model_file: + command = ( + f"{_graph_link_command(task_dir, job_work_dir, model_file)}; " f"{command}" ) - ) task_list = [ Task( command=command, @@ -1546,8 +1558,8 @@ def run_task(task_dir, machine_file, task_name, no_dp=False): ) for ii in task_dir_list ] - if not no_dp: - submission.forward_common_files = [os.path.join(task_dir, "graph.pb")] + if not no_dp and model_file: + submission.forward_common_files = [os.path.join(task_dir, model_file)] submission.register_task_list(task_list=task_list) submission.run_submission() diff --git a/dpti/hti_ice.py b/dpti/hti_ice.py index 7ab6dccb..1490605f 100755 --- a/dpti/hti_ice.py +++ b/dpti/hti_ice.py @@ -11,6 +11,7 @@ from dpti import einstein, hti from dpti.lib import lmp from dpti.lib.output import tee_stdout +from dpti.lib.utils import get_model_filename def _main(): @@ -218,12 +219,13 @@ def refine_tasks(from_task, to_task, err, print_ref=False): equi_conf = hti.get_task_file_abspath(from_task, jdata["equi_conf"]) model = hti.get_task_file_abspath(from_task, jdata["model"]) + model_file = get_model_filename(model) hti.create_path(to_task) shutil.copyfile(equi_conf, os.path.join(to_task, "conf.lmp")) jdata["equi_conf"] = "conf.lmp" - shutil.copyfile(model, os.path.join(to_task, "graph.pb")) - jdata["model"] = "graph.pb" + shutil.copyfile(model, os.path.join(to_task, model_file)) + jdata["model"] = model_file jdata["switch"] = switch jdata["orig_task"] = from_task jdata["refine_error"] = err diff --git a/dpti/hti_liq.py b/dpti/hti_liq.py index ba046055..52811773 100755 --- a/dpti/hti_liq.py +++ b/dpti/hti_liq.py @@ -19,6 +19,7 @@ compute_nrefine, create_path, get_first_matched_key_from_dict, + get_model_filename, integrate, integrate_sys_err, parse_seq, @@ -318,6 +319,7 @@ def _make_tasks(iter_name, jdata, step, if_meam=False, meam_model=None): ), f"there must be key-value for sigma or {sigma_key_name} in soft_param" job_abs_dir = create_path(iter_name) + model_file = os.path.basename(jdata["model"]) if jdata.get("model") else None if meam_model: relative_link_file(os.path.abspath(meam_model["library"]), job_abs_dir) @@ -329,7 +331,8 @@ def _make_tasks(iter_name, jdata, step, if_meam=False, meam_model=None): os.chdir(iter_name) os.symlink(os.path.join("..", "in.json"), "in.json") os.symlink(os.path.join("..", "conf.lmp"), "conf.lmp") - os.symlink(os.path.join("..", "graph.pb"), "graph.pb") + if model_file: + os.symlink(os.path.join("..", model_file), model_file) os.chdir(cwd) # print(9898, meam_model) @@ -338,7 +341,8 @@ def _make_tasks(iter_name, jdata, step, if_meam=False, meam_model=None): create_path(work_path) os.chdir(work_path) os.symlink(os.path.join("..", "conf.lmp"), "conf.lmp") - os.symlink(os.path.join("..", "graph.pb"), "graph.pb") + if model_file: + os.symlink(os.path.join("..", model_file), model_file) if meam_model: meam_library_basename = os.path.basename(meam_model["library"]) meam_potential_basename = os.path.basename(meam_model["potential"]) @@ -352,7 +356,7 @@ def _make_tasks(iter_name, jdata, step, if_meam=False, meam_model=None): mass_map, ii, soft_param, - "graph.pb", + model_file, nsteps, timestep, "nvt", @@ -378,6 +382,7 @@ def make_tasks(iter_name, jdata, if_meam=None): model = os.path.abspath(jdata["model"]) else: model = None + model_file = get_model_filename(model) if model else None meam_model = jdata.get("meam_model", None) create_path(iter_name) @@ -385,7 +390,7 @@ def make_tasks(iter_name, jdata, if_meam=None): shutil.copyfile(equi_conf, copied_conf) jdata["equi_conf"] = copied_conf if model: - copied_model = os.path.join(os.path.abspath(iter_name), "graph.pb") + copied_model = os.path.join(os.path.abspath(iter_name), model_file) shutil.copyfile(model, copied_model) jdata["model"] = copied_model else: @@ -569,14 +574,15 @@ def refine_tasks(from_task, to_task, err, print_ref=False): equi_conf = hti.get_task_file_abspath(from_task, jdata["equi_conf"]) model = hti.get_task_file_abspath(from_task, jdata["model"]) + model_file = get_model_filename(model) if_meam = jdata.get("if_meam", False) meam_model = jdata.get("meam_model", None) create_path(to_task) shutil.copyfile(equi_conf, os.path.join(to_task, "conf.lmp")) jdata["equi_conf"] = "conf.lmp" - shutil.copyfile(model, os.path.join(to_task, "graph.pb")) - jdata["model"] = "graph.pb" + shutil.copyfile(model, os.path.join(to_task, model_file)) + jdata["model"] = model_file jdata["orig_task"] = from_task jdata["refine_error"] = err diff --git a/dpti/hti_water.py b/dpti/hti_water.py index 8bcdf5bd..61c6b2fb 100755 --- a/dpti/hti_water.py +++ b/dpti/hti_water.py @@ -21,6 +21,7 @@ compute_nrefine, create_path, get_first_matched_key_from_dict, + get_model_filename, get_task_file_abspath, integrate_range, parse_seq, @@ -238,6 +239,7 @@ def _make_tasks(iter_name, jdata, step): pres = jdata["pres"] tau_t = jdata["tau_t"] tau_p = jdata["tau_p"] + model_file = os.path.basename(jdata["model"]) copies = None if "copies" in jdata: copies = jdata["copies"] @@ -247,7 +249,7 @@ def _make_tasks(iter_name, jdata, step): os.chdir(iter_name) os.symlink(os.path.join("..", "in.json"), "in.json") os.symlink(os.path.join("..", "conf.lmp"), "orig.lmp") - os.symlink(os.path.join("..", "graph.pb"), "graph.pb") + os.symlink(os.path.join("..", model_file), model_file) with open("orig.lmp") as f: lines = water.add_bonds(f.read().split("\n")) with open("conf.lmp", "w") as c: @@ -258,13 +260,13 @@ def _make_tasks(iter_name, jdata, step): create_path(work_path) os.chdir(work_path) os.symlink(os.path.join("..", "conf.lmp"), "conf.lmp") - os.symlink(os.path.join("..", "graph.pb"), "graph.pb") + os.symlink(os.path.join("..", model_file), model_file) lmp_str = _gen_lammps_input( step, "conf.lmp", mass_map, all_lambda[idx], - "graph.pb", + model_file, bparam, sparam, nsteps, @@ -357,14 +359,15 @@ def _refine_tasks(from_task, to_task, err, step): def make_tasks(iter_name, jdata): equi_conf = os.path.abspath(jdata["equi_conf"]) model = os.path.abspath(jdata["model"]) + model_file = get_model_filename(model) create_path(iter_name) copied_conf = os.path.join(os.path.abspath(iter_name), "conf.lmp") shutil.copyfile(equi_conf, copied_conf) jdata["equi_conf"] = "conf.lmp" - linked_model = os.path.join(os.path.abspath(iter_name), "graph.pb") + linked_model = os.path.join(os.path.abspath(iter_name), model_file) shutil.copyfile(model, linked_model) - jdata["model"] = "graph.pb" + jdata["model"] = model_file cwd = os.getcwd() os.chdir(iter_name) @@ -383,14 +386,15 @@ def refine_tasks(from_task, to_task, err): jdata = json.load(open(os.path.join(from_task, "in.json"))) equi_conf = get_task_file_abspath(from_task, jdata["equi_conf"]) model = get_task_file_abspath(from_task, jdata["model"]) + model_file = get_model_filename(model) create_path(to_task) copied_conf = os.path.join(os.path.abspath(to_task), "conf.lmp") shutil.copyfile(equi_conf, copied_conf) jdata["equi_conf"] = "conf.lmp" - linked_model = os.path.join(os.path.abspath(to_task), "graph.pb") + linked_model = os.path.join(os.path.abspath(to_task), model_file) shutil.copyfile(model, linked_model) - jdata["model"] = "graph.pb" + jdata["model"] = model_file jdata["orig_task"] = from_task jdata["refine_error"] = err diff --git a/dpti/lib/utils.py b/dpti/lib/utils.py index c551fd9b..87f88dad 100644 --- a/dpti/lib/utils.py +++ b/dpti/lib/utils.py @@ -77,6 +77,13 @@ def relative_link_file(file_path, target_dir): return target_linkfile_path +def get_model_filename(model, prefix="graph"): + if model is None: + return f"{prefix}.pb" + suffix = pathlib.PurePath(str(model)).suffix + return f"{prefix}{suffix}" + + def link_file_in_dict(dct, key_list, target_dir): if not dct: return {} diff --git a/dpti/mti.py b/dpti/mti.py index 17123c85..1a1726dd 100644 --- a/dpti/mti.py +++ b/dpti/mti.py @@ -318,17 +318,32 @@ def make_tasks(iter_name, jdata): json.dump(settings, f, indent=4) +def _get_task_model_file(task_name): + settings_file = os.path.join(task_name, "mti_settings.json") + if not os.path.isfile(settings_file): + return "graph.pb" + with open(settings_file) as fp: + settings = json.load(fp) + model = settings.get("model") + return os.path.basename(model) if model else None + + def run_task(task_name, jdata, machine_file): job_type = jdata["job_type"] nprocs_per_bead = jdata.get("nprocs_per_bead", 1) + model_file = _get_task_model_file(task_name) if job_type == "nbead_convergence": task_dir_list = glob.glob( os.path.join(task_name, "task.*/mass_scale_y.*/nbead.*") ) - link_model = "ln -s ../../../graph.pb" + link_model = ( + f'ln -sf "../../../{model_file}" "{model_file}"; ' if model_file else "" + ) elif job_type == "mass_ti": task_dir_list = glob.glob(os.path.join(task_name, "task.*/mass_scale_y.*")) - link_model = "ln -s ../../graph.pb" + link_model = ( + f'ln -sf "../../{model_file}" "{model_file}"; ' if model_file else "" + ) else: raise RuntimeError( "Unknow job_type. Only nbead_convergence and mass_ti are supported." @@ -360,9 +375,13 @@ def run_task(task_name, jdata, machine_file): ) task = Task( - command=f"{link_model}; if ls *.restart.100000 1> /dev/null 2>&1; then {task_exec} -in in.lammps -p {nbead}x{nprocs_per_bead} -log log -v restart 1; else {task_exec} -in in.lammps -p {nbead}x{nprocs_per_bead} -log log -v restart 0; fi", + command=f"{link_model}if ls *.restart.100000 1> /dev/null 2>&1; then {task_exec} -in in.lammps -p {nbead}x{nprocs_per_bead} -log log -v restart 1; else {task_exec} -in in.lammps -p {nbead}x{nprocs_per_bead} -log log -v restart 0; fi", task_work_path=ii, - forward_files=["in.lammps", "*.lmp", "graph.pb"], + forward_files=[ + forward_file + for forward_file in ["in.lammps", "*.lmp", model_file] + if forward_file + ], backward_files=["log*", "*out.lmp", "*.dump"], ) diff --git a/dpti/old_equi.py b/dpti/old_equi.py index 0320aeb2..8b1debbd 100755 --- a/dpti/old_equi.py +++ b/dpti/old_equi.py @@ -14,7 +14,12 @@ # import dpti from dpti.lib.lammps import get_last_dump, get_natoms, get_thermo from dpti.lib.lmp import from_system_data -from dpti.lib.utils import block_avg, create_path, get_task_file_abspath +from dpti.lib.utils import ( + block_avg, + create_path, + get_model_filename, + get_task_file_abspath, +) from dpti.lib.water import compute_bonds, posi_diff # from lib import dump @@ -218,11 +223,12 @@ def make_task( os.symlink(os.path.realpath(equi_conf), "conf.lmp") else: open("conf.lmp", "w").write(npt_equi_conf(npt_conf)) - os.symlink(os.path.realpath(model), "graph.pb") + model_file = get_model_filename(model) + os.symlink(os.path.realpath(model), model_file) lmp_str = _gen_lammps_input( "conf.lmp", model_mass_map, - "graph.pb", + model_file, nsteps, dt, ens, diff --git a/dpti/relax.py b/dpti/relax.py index ec01c8bf..b76ff18b 100755 --- a/dpti/relax.py +++ b/dpti/relax.py @@ -5,7 +5,7 @@ import os from lib.lammps import get_last_dump -from lib.utils import create_path, cvt_conf +from lib.utils import create_path, cvt_conf, get_model_filename def _gen_lammps_relax(conf_file, mass_map, model, pres, thermo_freq=100, dump_freq=100): @@ -60,8 +60,9 @@ def make_task(iter_name, jdata, pres): with open("in.json", "w") as fp: json.dump(jdata, fp, indent=4) os.symlink(os.path.relpath(equi_conf), "conf.lmp") - os.symlink(os.path.relpath(model), "graph.pb") - lmp_str = _gen_lammps_relax("conf.lmp", model_mass_map, "graph.pb", pres) + model_file = get_model_filename(model) + os.symlink(os.path.relpath(model), model_file) + lmp_str = _gen_lammps_relax("conf.lmp", model_mass_map, model_file, pres) with open("in.lammps", "w") as fp: fp.write(lmp_str) os.chdir(cwd) diff --git a/dpti/ti.py b/dpti/ti.py old mode 100755 new mode 100644 index 554afb94..88028234 --- a/dpti/ti.py +++ b/dpti/ti.py @@ -1046,6 +1046,16 @@ def _is_completed_lammps_task(task_work_path): return False +def _get_task_model_file(task_name): + settings_file = os.path.join(task_name, "ti_settings.json") + if not os.path.isfile(settings_file): + return "graph.pb" + with open(settings_file) as fp: + settings = json.load(fp) + model = settings.get("model", "graph.pb") + return os.path.basename(model) if model else None + + def run_task(task_name, machine_file): task_dir_list = glob.glob(os.path.join(task_name, "task.*")) task_dir_list = sorted(task_dir_list) @@ -1064,10 +1074,15 @@ def run_task(task_name, machine_file): resources=resources, machine=machine, ) + model_file = _get_task_model_file(task_name) task_list = [ Task( - command=f"ln -s ../graph.pb graph.pb; {mdata['command']} -in in.lammps", + command=( + f'ln -sf "../{model_file}" "{model_file}"; {mdata["command"]} -in in.lammps' + if model_file + else f"{mdata['command']} -in in.lammps" + ), task_work_path=ii, forward_files=["in.lammps", "*.lmp"], backward_files=["log*", "final.lmp", "traj.dump"], @@ -1075,7 +1090,8 @@ def run_task(task_name, machine_file): for ii in task_dir_list ] - submission.forward_common_files = [os.path.join(task_name, "graph.pb")] + if model_file: + submission.forward_common_files = [os.path.join(task_name, model_file)] submission.register_task_list(task_list=task_list) submission.run_submission() diff --git a/tests/test_equi_make_task.py b/tests/test_equi_make_task.py index d1b66400..93ab1426 100644 --- a/tests/test_equi_make_task.py +++ b/tests/test_equi_make_task.py @@ -114,6 +114,14 @@ def test_water_npt(self, patch_random): f2 = os.path.join(test_dir, file) self.assertEqual(get_file_md5(f1), get_file_md5(f2), msg=(f1, f2)) + def test_run_task_model_file_from_settings(self): + test_dir = os.path.join(self.test_dir, "model_file_lookup") + os.makedirs(test_dir) + with open(os.path.join(test_dir, "equi_settings.json"), "w") as fp: + json.dump({"model": "/tmp/model.pth"}, fp) + + self.assertEqual(dpti.equi._get_task_model_file(test_dir), "model.pth") + @classmethod def tearDownClass(cls): shutil.rmtree("tmp_equi/") diff --git a/tests/test_gdi_make_task.py b/tests/test_gdi_make_task.py index 552da98f..59d8b45c 100644 --- a/tests/test_gdi_make_task.py +++ b/tests/test_gdi_make_task.py @@ -204,7 +204,7 @@ def test_setup_dpdt_silica_phase_specific_input_layout(self): self.assertEqual(dpti.gdi._get_phase_graph_file(jdata, 1), "graph.1.pb") @patch("numpy.random.default_rng") - def test_deepmd_uses_local_graph_name(self, patch_random): + def test_deepmd_uses_phase_graph_name(self, patch_random): patch_random.return_value = MagicMock(integers=MagicMock(return_value=7858)) test_dir = os.path.join(self.test_dir, "deepmd_local_graph") json_file = os.path.join(self.benchmark_dir, "deepmd", "pb.json") @@ -225,8 +225,44 @@ def test_deepmd_uses_local_graph_name(self, patch_random): with open(os.path.join(test_dir, "in.lammps")) as fp: lmp_input = fp.read() - self.assertIn("pair_style deepmd graph.pb", lmp_input) - self.assertNotIn("pair_style deepmd graph.0.pb", lmp_input) + self.assertIn("pair_style deepmd graph.0.pb", lmp_input) + self.assertNotIn("pair_style deepmd graph.pb", lmp_input) + + def test_setup_dpdt_phase_specific_model_suffixes(self): + parent_dir = os.path.join(self.test_dir, "phase_model_suffixes") + task_dir = os.path.join(parent_dir, "gdi_job") + os.makedirs(parent_dir) + shutil.copyfile("conf.lmp", os.path.join(parent_dir, "conf.0.lmp")) + shutil.copyfile("alpha.lmp", os.path.join(parent_dir, "conf.1.lmp")) + shutil.copyfile("graph.pb", os.path.join(parent_dir, "solid_model.pth")) + shutil.copyfile("beta.lmp", os.path.join(parent_dir, "liquid_model.pb")) + jdata = { + "phase_i": { + "name": "PHASE_0", + "equi_conf": "conf.0.lmp", + "model": "solid_model.pth", + }, + "phase_ii": { + "name": "PHASE_1", + "equi_conf": "conf.1.lmp", + "model": "liquid_model.pb", + }, + "mass_map": [118.71], + "nsteps": 5000, + "timestep": 0.002, + "tau_t": 0.1, + "tau_p": 1.0, + "thermo_freq": 10, + "stat_skip": 100, + "stat_bsize": 10, + } + + dpti.gdi._setup_dpdt(task_dir, jdata) + + self.assertTrue(os.path.isfile(os.path.join(task_dir, "graph.0.pth"))) + self.assertTrue(os.path.isfile(os.path.join(task_dir, "graph.1.pb"))) + self.assertEqual(dpti.gdi._get_phase_graph_file(jdata, 0), "graph.0.pth") + self.assertEqual(dpti.gdi._get_phase_graph_file(jdata, 1), "graph.1.pb") @classmethod def tearDownClass(cls): diff --git a/tests/test_hti_make_task.py b/tests/test_hti_make_task.py index 4d87a616..e3a348a3 100644 --- a/tests/test_hti_make_task.py +++ b/tests/test_hti_make_task.py @@ -98,12 +98,40 @@ def test_deepmd_one_step(self, patch_random): def test_graph_link_command(self): self.assertEqual( dpti.hti._graph_link_command("hti", "hti"), - "ln -s ../graph.pb graph.pb", + "ln -sf ../graph.pb graph.pb", ) self.assertEqual( dpti.hti._graph_link_command("hti", "hti/00.deep_on"), - "ln -s ../../graph.pb graph.pb", + "ln -sf ../../graph.pb graph.pb", ) + self.assertEqual( + dpti.hti._graph_link_command("hti", "hti", "graph.pth"), + "ln -sf ../graph.pth graph.pth", + ) + + @patch("numpy.random.default_rng") + def test_deepmd_one_step_keeps_model_suffix(self, patch_random): + patch_random.return_value = MagicMock(integers=MagicMock(return_value=7858)) + benchmark_dir = os.path.join(self.benchmark_dir, "one_step") + test_dir = os.path.join(self.test_dir, "one_step_pth") + model_file = os.path.join(self.test_dir, "model.pth") + shutil.copyfile("graph.pb", model_file) + + json_file = os.path.join(benchmark_dir, "jdata.json") + with open(json_file) as f: + jdata = json.load(f) + jdata["model"] = model_file + + dpti.hti.make_tasks(iter_name=test_dir, jdata=jdata, switch="one-step") + + self.assertTrue(os.path.isfile(os.path.join(test_dir, "graph.pth"))) + self.assertTrue( + os.path.isfile(os.path.join(test_dir, "task.000006", "graph.pth")) + ) + with open(os.path.join(test_dir, "task.000006", "in.lammps")) as fp: + lmp_input = fp.read() + self.assertIn("pair_style deepmd graph.pth", lmp_input) + self.assertNotIn("deepmd graph.pb", lmp_input) @patch("numpy.random.default_rng") def test_meam_three_step(self, patch_random): diff --git a/tests/test_lib_utils.py b/tests/test_lib_utils.py index 24154445..3c63570b 100644 --- a/tests/test_lib_utils.py +++ b/tests/test_lib_utils.py @@ -5,7 +5,13 @@ import numpy as np from numpy.testing import assert_almost_equal -from dpti.lib.utils import block_avg, integrate_range_hti, parse_seq, relative_link_file +from dpti.lib.utils import ( + block_avg, + get_model_filename, + integrate_range_hti, + parse_seq, + relative_link_file, +) lambda_seq = [ "0.00:0.05:0.010", @@ -173,5 +179,21 @@ def tearDownClass(cls): shutil.rmtree("relative_link_file_test_dir/") +class TestGetModelFilename(unittest.TestCase): + def test_preserves_backend_suffix(self): + self.assertEqual(get_model_filename("graph.pb"), "graph.pb") + self.assertEqual(get_model_filename("/tmp/model.pth"), "graph.pth") + self.assertEqual(get_model_filename("model.savedmodel"), "graph.savedmodel") + + def test_phase_prefix(self): + self.assertEqual( + get_model_filename("/tmp/model.pth", prefix="graph.0"), + "graph.0.pth", + ) + + def test_default_model_name(self): + self.assertEqual(get_model_filename(None), "graph.pb") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_ti_make_task.py b/tests/test_ti_make_task.py index b2db1610..6c456449 100644 --- a/tests/test_ti_make_task.py +++ b/tests/test_ti_make_task.py @@ -68,6 +68,14 @@ def test_meam_path_p(self, patch_random): f2 = os.path.join(test_dir, file) self.assertEqual(get_file_md5(f1), get_file_md5(f2), msg=(f1, f2)) + def test_run_task_model_file_from_settings(self): + test_dir = os.path.join(self.test_dir, "model_file_lookup") + os.makedirs(test_dir) + with open(os.path.join(test_dir, "ti_settings.json"), "w") as fp: + json.dump({"model": "/tmp/model.pth"}, fp) + + self.assertEqual(dpti.ti._get_task_model_file(test_dir), "model.pth") + @classmethod def tearDownClass(cls): shutil.rmtree("tmp_ti/") diff --git a/workflow/DpFreeEnergy.py b/workflow/DpFreeEnergy.py index 610ae958..fd5d5611 100644 --- a/workflow/DpFreeEnergy.py +++ b/workflow/DpFreeEnergy.py @@ -12,6 +12,24 @@ from dpti import equi, hti, hti_liq, ti +def _model_file_from_json(job_work_dir, json_file): + path = os.path.join(job_work_dir, json_file) + if not os.path.isfile(path): + return "graph.pb" + with open(path) as fp: + jdata = json.load(fp) + model = jdata.get("model") + return os.path.basename(model) if model else None + + +def _forward_files_with_model(job_work_dir, json_file, base_files): + forward_files = list(base_files) + model_file = _model_file_from_json(job_work_dir, json_file) + if model_file: + forward_files.append(model_file) + return forward_files + + @task() def all_start_check(): context = get_current_context() @@ -97,7 +115,9 @@ def NPT_sim(job_work_dir): task = Task( command="lmp -i in.lammps", task_work_path="./", - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "equi_settings.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps", "dump.equi"], ) # Files to transfer forward and backwards context = get_current_context() @@ -159,7 +179,9 @@ def NVT_sim(job_work_dir): task = Task( command="lmp -i in.lammps", task_work_path="./", - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "equi_settings.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps"], ) submission.register_task_list([task]) @@ -236,7 +258,9 @@ def HTI_sim(job_work_dir): Task( command="lmp -i in.lammps", task_work_path=ii, - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "in.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps"], ) for ii in task_dir_list @@ -322,7 +346,9 @@ def TI_sim(job_work_dir): Task( command="lmp -i in.lammps", task_work_path=ii, - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "ti_settings.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps"], ) for ii in task_dir_list diff --git a/workflow/DpFreeEnergyWater.py b/workflow/DpFreeEnergyWater.py index b3f58598..e748fd93 100644 --- a/workflow/DpFreeEnergyWater.py +++ b/workflow/DpFreeEnergyWater.py @@ -16,6 +16,24 @@ from dpti import equi, hti, hti_ice, hti_liq, ti +def _model_file_from_json(job_work_dir, json_file): + path = os.path.join(job_work_dir, json_file) + if not os.path.isfile(path): + return "graph.pb" + with open(path) as fp: + jdata = json.load(fp) + model = jdata.get("model") + return os.path.basename(model) if model else None + + +def _forward_files_with_model(job_work_dir, json_file, base_files): + forward_files = list(base_files) + model_file = _model_file_from_json(job_work_dir, json_file) + if model_file: + forward_files.append(model_file) + return forward_files + + def get_empty_submission(job_work_dir): context = get_current_context() dag_run = context["params"] @@ -118,7 +136,9 @@ def NPT_sim(job_work_dir): task = Task( command="lmp -i in.lammps", task_work_path="./", - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "equi_settings.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps", "dump.equi"], ) submission = get_empty_submission(job_work_dir) @@ -174,7 +194,9 @@ def NVT_sim(job_work_dir): task = Task( command="lmp -i in.lammps", task_work_path="./", - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "equi_settings.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps", "out.lmp"], ) submission.register_task_list([task]) @@ -255,7 +277,8 @@ def HTI_sim(HTI_init_info): for ii in task_dir_list ] submission = get_empty_submission(job_work_dir) - submission.forward_common_files = ["graph.pb"] + model_file = _model_file_from_json(job_work_dir, "in.json") + submission.forward_common_files = [model_file] if model_file else [] submission.register_task_list(task_list=task_list) submission.run_submission() return job_work_dir @@ -345,7 +368,9 @@ def TI_sim(job_work_dir): Task( command="lmp -i in.lammps", task_work_path=ii, - forward_files=["in.lammps", "*lmp", "graph.pb"], + forward_files=_forward_files_with_model( + job_work_dir, "ti_settings.json", ["in.lammps", "*lmp"] + ), backward_files=["log.lammps"], ) for ii in task_dir_list