Preserve model backend suffix in generated tasks

dpti/equi.py

@@ -604,6 +604,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"
@@ -686,7 +696,9 @@ def run_task(task_name, machine_file):
             [os.path.basename(ii) for ii in normalize_template_ff_files(jdata)]
         )
     else:
-        forward_files.append("graph.pb")
+        model_file = _get_task_model_file(task_name)
+        if model_file:
+            forward_files.append(model_file)
     work_base_dir = os.getcwd()
     with open(machine_file) as f:
         mdata = json.load(f)

dpti/gdi.py

@@ -22,6 +22,7 @@
 from dpti.lib.utils import (
     create_path,
     get_first_matched_key_from_dict,
+    get_model_filename,
     get_template_ff_file,
     normalize_template_ff_files,
     read_template_ff,
@@ -120,9 +121,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 template_ff_files is not None:
         for file_path in template_ff_files:
             relative_link_file(file_path, "./")
@@ -198,9 +199,11 @@ def _has_phase_specific_template_ff(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 _get_phase_template_ff_name(jdata, phase_idx):
@@ -210,14 +213,16 @@ def _get_phase_template_ff_name(jdata, phase_idx):
 
 
 def _get_phase_forward_files(jdata, phase_key):
+    """Return files staged for one GDI phase and its selected force field."""
+    phase_idx = 0 if phase_key == "phase_i" else 1
     forward_files = ["conf.lmp", "in.lammps"]
     if _get_phase_template_ff_file(jdata, phase_key) is not None:
         forward_files.extend(
             os.path.basename(ii)
             for ii in _get_phase_template_ff_files(jdata, phase_key)
         )
     else:
-        forward_files.append("graph.pb")
+        forward_files.append(_get_phase_graph_file(jdata, phase_idx))
     return forward_files
 
 

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,
     get_template_ff_file,
     integrate_range_hti,
@@ -629,6 +631,7 @@ def make_tasks(iter_name, jdata, ref="einstein", switch="one-step", if_meam=None
     model = jdata.get("model")
     if model is not None:
         model = os.path.abspath(model)
+    model_file = get_model_filename(model) if model is not None else None
     template_ff_file = None if if_meam else get_template_ff_file(jdata)
     if template_ff_file is not None:
         template_ff_file = os.path.abspath(template_ff_file)
@@ -665,9 +668,9 @@ def make_tasks(iter_name, jdata, ref="einstein", switch="one-step", if_meam=None
             relative_link_file(meam_model["library"], job_abs_dir)
             relative_link_file(meam_model["potential"], job_abs_dir)
         elif model is not None:
-            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
         elif template_ff_file is not None:
             template_destination = os.path.join(
                 os.path.abspath(iter_name), os.path.basename(template_ff_file)
@@ -803,6 +806,7 @@ def _make_tasks(
     model = jdata.get("model")
     if model is not None:
         model = os.path.abspath(model)
+    model_file = get_model_filename(model) if model is not None else None
     template_ff_file = None if if_meam else get_template_ff_file(jdata)
     template_ff = None
     if template_ff_file is not None:
@@ -879,15 +883,15 @@ def _make_tasks(
     linked_model = None
     linked_template_files = []
     if model is not None:
-        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
     elif template_ff is not None:
         template_destination = os.path.join(
             os.path.abspath(iter_name), os.path.basename(template_ff_file)
@@ -929,7 +933,7 @@ def _make_tasks(
         os.chdir(work_path)
         os.symlink(os.path.relpath(copied_conf), "conf.lmp")
         if linked_model is not None:
-            os.symlink(os.path.relpath(linked_model), "graph.pb")
+            os.symlink(os.path.relpath(linked_model), model_file)
         for template_file in linked_template_files:
             os.symlink(
                 os.path.relpath(template_file, os.path.abspath(".")),
@@ -958,7 +962,7 @@ def _make_tasks(
                 "conf.lmp",
                 mass_map,
                 ii,
-                "graph.pb",
+                model_file,
                 m_spring_k,
                 nsteps,
                 timestep,
@@ -1649,12 +1653,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):
@@ -1692,15 +1706,13 @@ def run_task(task_dir, machine_file, task_name, no_dp=False):
         resources=resources,
         machine=machine,
     )
+    model_file = None if uses_template else _get_task_model_file(task_dir)
 
-    command = (
-        f"{mdata['command']} -i in.lammps -screen none"
-        if no_dp or uses_template
-        else (
-            f"{_graph_link_command(task_dir, job_work_dir)}; "
-            f"{mdata['command']} -i in.lammps -screen none"
+    command = f"{mdata['command']} -i in.lammps -screen none"
+    if not no_dp and model_file:
+        command = (
+            f"{_graph_link_command(task_dir, job_work_dir, model_file)}; {command}"
         )
-    )
     forward_files = ["in.lammps", "conf.lmp"]
     if uses_template:
         forward_files.extend(template_ff_files)
@@ -1713,8 +1725,8 @@ def run_task(task_dir, machine_file, task_name, no_dp=False):
         )
         for ii in task_dir_list
     ]
-    if not no_dp and not uses_template:
-        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():
@@ -225,8 +226,9 @@ def refine_tasks(from_task, to_task, err, print_ref=False):
     shutil.copyfile(equi_conf, os.path.join(to_task, "conf.lmp"))
     jdata["equi_conf"] = "conf.lmp"
     if model is not None:
-        shutil.copyfile(model, os.path.join(to_task, "graph.pb"))
-        jdata["model"] = "graph.pb"
+        model_file = get_model_filename(model)
+        shutil.copyfile(model, os.path.join(to_task, model_file))
+        jdata["model"] = model_file
     else:
         template_source = hti.get_task_file_abspath(from_task, jdata["template_ff"])
         shutil.copyfile(

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
 

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,
     get_template_ff_file,
     integrate_range,
@@ -273,6 +274,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(model) if model is not None else None
     copies = None
     if "copies" in jdata:
         copies = jdata["copies"]
@@ -283,7 +285,7 @@ def _make_tasks(iter_name, jdata, step):
     os.symlink(os.path.join("..", "in.json"), "in.json")
     os.symlink(os.path.join("..", "conf.lmp"), "orig.lmp")
     if model is not None:
-        os.symlink(os.path.join("..", "graph.pb"), "graph.pb")
+        os.symlink(os.path.join("..", model_file), model_file)
     for template_file in normalize_template_ff_files(jdata):
         basename = os.path.basename(template_file)
         os.symlink(os.path.join("..", basename), basename)
@@ -298,7 +300,7 @@ def _make_tasks(iter_name, jdata, step):
         os.chdir(work_path)
         os.symlink(os.path.join("..", "conf.lmp"), "conf.lmp")
         if model is not None:
-            os.symlink(os.path.join("..", "graph.pb"), "graph.pb")
+            os.symlink(os.path.join("..", model_file), model_file)
         for template_file in normalize_template_ff_files(jdata):
             basename = os.path.basename(template_file)
             os.symlink(os.path.join("..", basename), basename)
@@ -307,7 +309,7 @@ def _make_tasks(iter_name, jdata, step):
             "conf.lmp",
             mass_map,
             all_lambda[idx],
-            "graph.pb",
+            model_file,
             bparam,
             sparam,
             nsteps,
@@ -421,9 +423,10 @@ def make_tasks(iter_name, jdata):
     shutil.copyfile(equi_conf, copied_conf)
     jdata["equi_conf"] = "conf.lmp"
     if model is not None:
-        linked_model = os.path.join(os.path.abspath(iter_name), "graph.pb")
+        model_file = get_model_filename(model)
+        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
     else:
         copied_template = os.path.join(
             os.path.abspath(iter_name), os.path.basename(template_ff_file)
@@ -464,9 +467,10 @@ def refine_tasks(from_task, to_task, err):
     shutil.copyfile(equi_conf, copied_conf)
     jdata["equi_conf"] = "conf.lmp"
     if model is not None:
-        linked_model = os.path.join(os.path.abspath(to_task), "graph.pb")
+        model_file = get_model_filename(model)
+        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
     else:
         template_source = get_task_file_abspath(from_task, jdata["template_ff"])
         template_destination = os.path.join(