feat(generator): support generic conditioning parameters

dpgen/generator/arginfo.py

@@ -36,6 +36,10 @@ def basic_args() -> list[Argument]:
 - 0: no electron temperature. \n\n\
 - 1: eletron temperature as frame parameter. \n\n\
 - 2: electron temperature as atom parameter."
+    doc_conditioning = (
+        "Generic DeePMD conditioning declarations. Use fparam/aparam lists of "
+        "{name, source, dim} mappings; source values are read from FP job.json."
+    )
 
     return [
         Argument("type_map", list[str], optional=False, doc=doc_type_map),
@@ -47,6 +51,7 @@ def basic_args() -> list[Argument]:
             doc=doc_mass_map,
         ),
         Argument("use_ele_temp", int, optional=True, default=0, doc=doc_use_ele_temp),
+        Argument("conditioning", dict, optional=True, default={}, doc=doc_conditioning),
     ]
 
 
@@ -350,7 +355,7 @@ def model_devi_jobs_args() -> list[Argument]:
         model_devi_jobs_rev_mat_args(),
         Argument("sys_rev_mat", dict, optional=True, doc=doc_sys_rev_mat),
         Argument("sys_idx", list[int], optional=False, doc=doc_sys_idx),
-        Argument("temps", list[float], optional=True, doc=doc_temps),
+        Argument("temps", [list[float], list[dict]], optional=True, doc=doc_temps),
         Argument("press", list[float], optional=True, doc=doc_press),
         Argument("trj_freq", int, optional=False, doc=doc_trj_freq),
         Argument("nsteps", int, optional=True, doc=doc_nsteps),

dpgen/generator/lib/lammps.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import random
+import re
 
 import dpdata
 import numpy as np
@@ -17,6 +18,29 @@ def _sample_sphere():
         return vv / vn
 
 
+def conditioning_variable_names(conditioning, location, item):
+    """Return stable LAMMPS variable names for one conditioning item."""
+    base = re.sub(r"[^A-Za-z0-9_]", "_", item["source"]).upper()
+    prefix = f"DPGEN_{location.upper()}_{base}"
+    if item["dim"] == 1:
+        return [prefix]
+    return [f"{prefix}_{idx}" for idx in range(item["dim"])]
+
+
+def conditioning_keywords(conditioning):
+    """Return pair-style fparam/aparam keywords for conditioning."""
+    parts = []
+    for location in ("fparam", "aparam"):
+        variables = []
+        for item in conditioning.get(location, []):
+            variables.extend(conditioning_variable_names(conditioning, location, item))
+        if variables:
+            parts.append(
+                location + " " + " ".join(f"${{{name}}}" for name in variables)
+            )
+    return " ".join(parts)
+
+
 def make_lammps_input(
     ensemble,
     conf_file,
@@ -34,6 +58,8 @@ def make_lammps_input(
     pka_e=None,
     ele_temp_f=None,
     ele_temp_a=None,
+    conditioning=None,
+    param_values=None,
     max_seed=1000000,
     nopbc=False,
     deepmd_version="0.1",
@@ -70,6 +96,27 @@ def make_lammps_input(
         ret += f"variable        ELE_TEMP        equal {ele_temp_f:f}\n"
     if ele_temp_a is not None:
         ret += f"variable        ELE_TEMP        equal {ele_temp_a:f}\n"
+    if conditioning:
+        param_values = param_values or {}
+        for location in ("fparam", "aparam"):
+            for item in conditioning.get(location, []):
+                value = param_values.get(item["source"], param_values.get(item["name"]))
+                if value is None:
+                    raise ValueError(
+                        f"missing value for conditioning source {item['source']!r}"
+                    )
+                values = np.asarray(value, dtype=float).reshape(-1)
+                if values.size != item["dim"]:
+                    raise ValueError(
+                        f"conditioning source {item['source']!r} has width "
+                        f"{values.size}, expected {item['dim']}"
+                    )
+                for name, scalar in zip(
+                    conditioning_variable_names(conditioning, location, item),
+                    values,
+                    strict=True,
+                ):
+                    ret += f"variable        {name:<16} equal {scalar:f}\n"
     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"
@@ -135,6 +182,10 @@ def make_lammps_input(
             keywords += "fparam ${ELE_TEMP}"
         if ele_temp_a is not None:
             keywords += "aparam ${ELE_TEMP}"
+        if conditioning:
+            conditioning_text = conditioning_keywords(conditioning)
+            if conditioning_text:
+                keywords += conditioning_text + " "
 
         if d3_enabled:
             # Use hybrid/overlay with D3