Preserve model backend suffix in generated tasks

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

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}/",

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

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

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,14 +382,15 @@ 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)
     copied_conf = os.path.join(os.path.abspath(iter_name), "conf.lmp")
     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