Support custom variables for the LAMMPS input files (#70)

* fix hti's output of Helmholtz free energy for water * correct error estimation: divide natoms in stead of np.sqrt(natoms) * support custom LAMMPS variables and appending strings in pair_style * keyword 'append' should go to gen_equi_forcefield * add a missing space * support custom LAMMPS variables and appending strings in pair_style for hti * support one-step for run_task in hti * superposition of spring and DP should not be hybrid/overlay * support custom LAMMPS variables and appending strings in pair_style for ti * add custom_variables and append in _gen_lammps_input of ti * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * if custom_variables is not None * use f-string instead of percentage formatting * add unittest for gen_equi_header * correct function name * enable --unsafe-fixes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add unittest for dpti.equi.gen_equi_force_field * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add the pair_coeff * * deepmd deleted mistakenly * fix equi pair_style * pair_coeff deepmd * pair_coeff deepmd --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
deepmodeling · Aug 6, 2024 · 3c91292 · 3c91292
1 parent 1d7b742
commit 3c91292
Show file tree

Hide file tree

Showing 6 changed files with 148 additions and 21 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -23,7 +23,7 @@ repos:
     rev: v0.5.5
     hooks:
     - id: ruff
-      args: ["--fix"]
+      args: ["--fix", "--unsafe-fixes"]
       types_or: [python, pyi, jupyter]
     - id: ruff-format
       types_or: [python, pyi, jupyter]

diff --git a/dpti/equi.py b/dpti/equi.py
@@ -35,7 +35,16 @@
 # def gen_equi_header(nsteps, prt_freq, dump_freq, temp, pres, tau_t, tau_p, mass_map, conf_file):
 # def gen_equi_header(nsteps, prt_freq, dump_freq, temp, pres, tau_t, tau_p, mass_map, conf_file):
 def gen_equi_header(
-    nsteps, thermo_freq, dump_freq, mass_map, temp, tau_t, tau_p, equi_conf, pres=None
+    nsteps,
+    thermo_freq,
+    dump_freq,
+    mass_map,
+    temp,
+    tau_t,
+    tau_p,
+    equi_conf,
+    pres=None,
+    custom_variables=None,
 ):
     ret = ""
     ret += "clear\n"
@@ -45,6 +54,9 @@ def gen_equi_header(
     ret += "variable        DUMP_FREQ       equal %d\n" % dump_freq
     ret += "variable        NREPEAT         equal ${NSTEPS}/${DUMP_FREQ}\n"
     ret += f"variable        TEMP            equal {temp:.6f}\n"
+    if custom_variables is not None:
+        for key, value in custom_variables.items():
+            ret += f"variable        {key}            equal {value}\n"
     # if equi_settings['pres'] is not None :
     if pres is not None:
         ret += f"variable        PRES            equal {pres:.6f}\n"
@@ -64,7 +76,7 @@ def gen_equi_header(
 
 
 # def gen_equi_force_field(model, if_meam=None):
-def gen_equi_force_field(model, if_meam=False, meam_model=None):
+def gen_equi_force_field(model, if_meam=False, meam_model=None, append=None):
     # equi_settings =
     # model = equi_settings['model']
     # assert type(model) is dict, f"equi_settings['model] must be a dict. model:{model}"
@@ -82,7 +94,10 @@ def gen_equi_force_field(model, if_meam=False, meam_model=None):
     ret = ""
     ret += "# --------------------- FORCE FIELDS ---------------------\n"
     if not if_meam:
-        ret += f"pair_style      deepmd {model}\n"
+        ret += f"pair_style      deepmd {model}"
+        if append is not None:
+            ret += " " + append
+        ret += "\n"
         ret += "pair_coeff * *\n"
     else:
         meam_library = meam_model["library"]
@@ -161,6 +176,8 @@ def gen_equi_lammps_input(
     pres=None,
     if_meam=False,
     meam_model=None,
+    custom_variables=None,
+    append=None,
 ):
     if dump_freq is None:
         dump_freq = thermo_freq
@@ -174,9 +191,10 @@ def gen_equi_lammps_input(
         tau_p=tau_p,
         equi_conf=equi_conf,
         pres=pres,
+        custom_variables=custom_variables,
     )
     equi_force_field = gen_equi_force_field(
-        model, if_meam=if_meam, meam_model=meam_model
+        model, if_meam=if_meam, meam_model=meam_model, append=append
     )
     equi_thermo_settings = gen_equi_thermo_settings(timestep=timestep)
     equi_dump_settings = gen_equi_dump_settings(if_dump_avg_posi=if_dump_avg_posi)
@@ -422,6 +440,8 @@ def make_task(
         pres=equi_settings["pres"],
         if_meam=equi_settings["if_meam"],
         meam_model=equi_settings["meam_model"],
+        custom_variables=equi_settings.get("custom_variables", None),
+        append=equi_settings.get("append", None),
     )
 
     with open(os.path.join(task_abs_dir, "in.lammps"), "w") as fp:

diff --git a/dpti/hti.py b/dpti/hti.py
@@ -68,7 +68,7 @@ def _ff_lj_on(lamb, model, sparam):
     return ret
 
 
-def _ff_deep_on(lamb, model, sparam, if_meam=False, meam_model=None):
+def _ff_deep_on(lamb, model, sparam, if_meam=False, meam_model=None, append=None):
     nn = sparam["n"]
     alpha_lj = sparam["alpha_lj"]
     rcut = sparam["rcut"]
@@ -88,9 +88,11 @@ def _ff_deep_on(lamb, model, sparam, if_meam=False, meam_model=None):
         ret += f"pair_style      hybrid/overlay meam lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
         ret += f'pair_coeff      * * meam {meam_model["library"]} {meam_model["element"]} {meam_model["potential"]} {meam_model["element"]}\n'
     else:
-        ret += f"pair_style      hybrid/overlay deepmd {model} lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
+        if append:
+            ret += f"pair_style      hybrid/overlay deepmd {model:s} {append:s} lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
+        else:
+            ret += f"pair_style      hybrid/overlay deepmd {model:s} lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
         ret += "pair_coeff      * * deepmd\n"
-
     element_num = sparam.get("element_num", 1)
     sigma_key_index = filter(
         lambda t: t[0] <= t[1],
@@ -155,7 +157,7 @@ def _ff_deep_on(lamb, model, sparam, if_meam=False, meam_model=None):
 #     return ret
 
 
-def _ff_lj_off(lamb, model, sparam, if_meam=False, meam_model=None):
+def _ff_lj_off(lamb, model, sparam, if_meam=False, meam_model=None, append=None):
     nn = sparam["n"]
     alpha_lj = sparam["alpha_lj"]
     rcut = sparam["rcut"]
@@ -176,7 +178,10 @@ def _ff_lj_off(lamb, model, sparam, if_meam=False, meam_model=None):
         ret += f'pair_coeff      * * meam {meam_model["library"]} {meam_model["element"]} {meam_model["potential"]} {meam_model["element"]}\n'
         # ret += f'pair_coeff      * * meam {meam_model[0]} {meam_model[2]} {meam_model[1]} {meam_model[2]}\n'
     else:
-        ret += f"pair_style      hybrid/overlay deepmd {model} lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
+        if append:
+            ret += f"pair_style      hybrid/overlay deepmd {model:s} {append:s} lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
+        else:
+            ret += f"pair_style      hybrid/overlay deepmd {model:s} lj/cut/soft {nn:f} {alpha_lj:f} {rcut:f}\n"
         ret += "pair_coeff      * * deepmd\n"
 
     element_num = sparam.get("element_num", 1)
@@ -252,19 +257,25 @@ def _ff_spring(lamb, m_spring_k, var_spring):
     return ret
 
 
-def _ff_soft_lj(lamb, model, m_spring_k, step, sparam, if_meam=False, meam_model=None):
+def _ff_soft_lj(
+    lamb, model, m_spring_k, step, sparam, if_meam=False, meam_model=None, append=None
+):
     ret = ""
     ret += "# --------------------- FORCE FIELDS ---------------------\n"
     if step == "lj_on":
         ret += _ff_lj_on(lamb, model, sparam)
         var_spring = False
     elif step == "deep_on":
         # ret += _ff_meam_on(lamb, model, sparam)
-        ret += _ff_deep_on(lamb, model, sparam, if_meam=if_meam, meam_model=meam_model)
+        ret += _ff_deep_on(
+            lamb, model, sparam, if_meam=if_meam, meam_model=meam_model, append=append
+        )
         var_spring = False
     elif step == "spring_off":
         # ret += _ff_meam_lj_off(lamb, model, sparam)
-        ret += _ff_lj_off(lamb, model, sparam, if_meam=if_meam, meam_model=meam_model)
+        ret += _ff_lj_off(
+            lamb, model, sparam, if_meam=if_meam, meam_model=meam_model, append=append
+        )
         var_spring = True
     else:
         raise RuntimeError("unkown step", step)
@@ -274,10 +285,13 @@ def _ff_soft_lj(lamb, model, m_spring_k, step, sparam, if_meam=False, meam_model
     return ret
 
 
-def _ff_two_steps(lamb, model, m_spring_k, step):
+def _ff_two_steps(lamb, model, m_spring_k, step, append=None):
     ret = ""
     ret += "# --------------------- FORCE FIELDS ---------------------\n"
-    ret += f"pair_style      deepmd {model}\n"
+    if append:
+        ret += f"pair_style      deepmd {model:s} {append:s}\n"
+    else:
+        ret += f"pair_style      deepmd {model:s}\n"
     ret += "pair_coeff * *\n"
 
     if step == "both" or step == "spring_off":
@@ -323,6 +337,8 @@ def _gen_lammps_input(
     step="both",
     if_meam=False,
     meam_model=None,
+    custom_variables=None,
+    append=None,
 ):
     ret = ""
     ret += "clear\n"
@@ -336,6 +352,9 @@ def _gen_lammps_input(
     ret += f"variable        TAU_P           equal {tau_p:f}\n"
     ret += f"variable        LAMBDA          equal {lamb:.10e}\n"
     ret += "variable        INV_LAMBDA      equal %.10e\n" % (1 - lamb)
+    if custom_variables is not None:
+        for key, value in custom_variables.items():
+            ret += f"variable {key} equal {value}\n"
     ret += "# ---------------------- INITIALIZAITION ------------------\n"
     ret += "units           metal\n"
     ret += "boundary        p p p\n"
@@ -351,7 +370,7 @@ def _gen_lammps_input(
 
     # force field setting
     if switch == "one-step" or switch == "two-step":
-        ret += _ff_two_steps(lamb, model, m_spring_k, step)
+        ret += _ff_two_steps(lamb, model, m_spring_k, step, append)
     elif switch == "three-step":
         ret += _ff_soft_lj(
             lamb,
@@ -361,6 +380,7 @@ def _gen_lammps_input(
             sparam,
             if_meam=if_meam,
             meam_model=meam_model,
+            append=append,
         )
     else:
         raise RuntimeError("unknow switch", switch)
@@ -666,6 +686,8 @@ def _make_tasks(
     # timestep = jdata['timestep']
     timestep = get_first_matched_key_from_dict(jdata, ["timestep", "dt"])
     spring_k = jdata["spring_k"]
+    custom_variables = jdata.get("custom_variables", None)
+    append = jdata.get("append", None)
 
     sparam = jdata.get("soft_param", {})
     if sparam:
@@ -783,6 +805,8 @@ def _make_tasks(
                 crystal=crystal,
                 if_meam=if_meam,
                 meam_model=meam_model,
+                custom_variables=custom_variables,
+                append=append,
             )
         elif ref == "ideal":
             raise RuntimeError("choose hti_liq.py")
@@ -1375,9 +1399,12 @@ def hti_phase_trans_analyze(job, jdata=None):
 
 
 def run_task(task_dir, machine_file, task_name, no_dp=False):
-    job_work_dir_ = glob.glob(os.path.join(task_dir, task_name + "*"))
-    assert len(job_work_dir_) == 1
-    job_work_dir = job_work_dir_[0]
+    if task_name == "00" or task_name == "01" or task_name == "02":
+        job_work_dir_ = glob.glob(os.path.join(task_dir, task_name + "*"))
+        assert len(job_work_dir_) == 1
+        job_work_dir = job_work_dir_[0]
+    elif task_name == "one-step":
+        job_work_dir = task_dir
     task_dir_list = glob.glob(os.path.join(job_work_dir, "task*"))
     task_dir_list = sorted(task_dir_list)
     work_base_dir = os.getcwd()

diff --git a/dpti/ti.py b/dpti/ti.py
@@ -59,6 +59,8 @@ def _gen_lammps_input(
     copies=None,
     if_meam=False,
     meam_model=None,
+    custom_variables=None,
+    append=None,
 ):
     ret = ""
     ret += "clear\n"
@@ -70,6 +72,9 @@ def _gen_lammps_input(
     ret += f"variable        PRES            equal {pres:f}\n"
     ret += f"variable        TAU_T           equal {tau_t:f}\n"
     ret += f"variable        TAU_P           equal {tau_p:f}\n"
+    if custom_variables is not None:
+        for key, value in custom_variables.items():
+            ret += f"variable {key} equal {value}\n"
     ret += "# ---------------------- INITIALIZAITION ------------------\n"
     ret += "units           metal\n"
     ret += "boundary        p p p\n"
@@ -90,8 +95,12 @@ def _gen_lammps_input(
         ret += "pair_style      meam\n"
         ret += f'pair_coeff      * * {meam_model["library"]} {meam_model["element"]} {meam_model["potential"]} {meam_model["element"]}\n'
     else:
-        ret += f"pair_style      deepmd {model}\n"
-        ret += "pair_coeff * *\n"
+        if append:
+            ret += f"pair_style      deepmd {model:s} {append:s}\n"
+            ret += "pair_coeff * *\n"
+        else:
+            ret += f"pair_style      deepmd {model:s}\n"
+            ret += "pair_coeff * *\n"
     ret += "# --------------------- MD SETTINGS ----------------------\n"
     ret += "neighbor        1.0 bin\n"
     ret += f"timestep        {timestep}\n"
@@ -143,6 +152,8 @@ def make_tasks(iter_name, jdata, if_meam=None):
     if "copies" in jdata:
         copies = jdata["copies"]
     model = jdata["model"]
+    custom_variables = jdata.get("custom_variables", None)
+    append = jdata.get("append", None)
     meam_model = jdata.get("meam_model", None)
     # model = os.path.abspath(model)
     # mass_map = jdata['mass_map']
@@ -266,6 +277,8 @@ def make_tasks(iter_name, jdata, if_meam=None):
                 copies=copies,
                 if_meam=if_meam,
                 meam_model=meam_model,
+                custom_variables=custom_variables,
+                append=append,
             )
             thermo_out = temp_list[ii]
             # with open('thermo.out', 'w') as fp :
@@ -287,6 +300,8 @@ def make_tasks(iter_name, jdata, if_meam=None):
                 copies=copies,
                 if_meam=if_meam,
                 meam_model=meam_model,
+                custom_variables=custom_variables,
+                append=append,
             )
             thermo_out = temp_list[ii]
             # with open('thermo.out', 'w') as fp :
@@ -308,6 +323,8 @@ def make_tasks(iter_name, jdata, if_meam=None):
                 copies=copies,
                 if_meam=if_meam,
                 meam_model=meam_model,
+                custom_variables=custom_variables,
+                append=append,
             )
             thermo_out = pres_list[ii]
         else:

diff --git a/tests/test_equi_gen_force_field.py b/tests/test_equi_gen_force_field.py
@@ -22,6 +22,24 @@ def test_deepmd(self):
         ret2 = dpti.equi.gen_equi_force_field(**input)
         self.assertEqual(ret1, ret2)
 
+    def test_deepmd_append(self):
+        input = {
+            "model": "graph.pb",
+            "if_meam": False,
+            "meam_model": None,
+            "append": "fparam ${TeV}",
+        }
+
+        ret1 = textwrap.dedent(
+            """\
+        # --------------------- FORCE FIELDS ---------------------
+        pair_style      deepmd graph.pb fparam ${TeV}
+        pair_coeff * *
+        """
+        )
+        ret2 = dpti.equi.gen_equi_force_field(**input)
+        self.assertEqual(ret1, ret2)
+
     def test_meam(self):
         input = {"model": None, "if_meam": True, "meam_model": meam_model}
         ret1 = textwrap.dedent(

diff --git a/tests/test_equi_gen_header.py b/tests/test_equi_gen_header.py
@@ -49,6 +49,51 @@ def test_equi_header_npt(self):
         ret2 = dpti.equi.gen_equi_header(**input)
         self.assertEqual(ret1, ret2)
 
+    def test_equi_header_npt_custom_variable(self):
+        input = {
+            "nsteps": 1000000,
+            "thermo_freq": 10,
+            "dump_freq": 100000,
+            "temp": 400,
+            "tau_t": 0.2,
+            "tau_p": 2.0,
+            "mass_map": [118.71],
+            "equi_conf": "conf.lmp",
+            "pres": 200000,
+            "custom_variables": {
+                "kBeV": 8.617333262e-5,
+                "TeV": "${TEMP}*${kBeV}",
+            },
+        }
+
+        ret1 = textwrap.dedent(
+            """\
+        clear
+        # --------------------- VARIABLES-------------------------
+        variable        NSTEPS          equal 1000000
+        variable        THERMO_FREQ     equal 10
+        variable        DUMP_FREQ       equal 100000
+        variable        NREPEAT         equal ${NSTEPS}/${DUMP_FREQ}
+        variable        TEMP            equal 400.000000
+        variable        kBeV            equal 8.617333262e-05
+        variable        TeV            equal ${TEMP}*${kBeV}
+        variable        PRES            equal 200000.000000
+        variable        TAU_T           equal 0.200000
+        variable        TAU_P           equal 2.000000
+        # ---------------------- INITIALIZAITION ------------------
+        units           metal
+        boundary        p p p
+        atom_style      atomic
+        # --------------------- ATOM DEFINITION ------------------
+        box             tilt large
+        read_data       conf.lmp
+        change_box      all triclinic
+        mass            1 118.710000
+        """
+        )
+        ret2 = dpti.equi.gen_equi_header(**input)
+        self.assertEqual(ret1, ret2)
+
     def test_equi_header_nvt(self):
         input = {
             "nsteps": 1000000,