add arguments for lammps revise_template

dpgen/generator/arginfo.py

@@ -101,8 +101,38 @@ def training_args() -> List[Argument]:
 
 
 # Exploration
+def model_devi_jobs_template_args() -> Argument:
+    doc_template = ('Give an input file template for the supported engine software adopted in 01.model_devi. '
+                   'Through user-defined template, any freedom (function) that is permitted by the engine '
+                   'software could be inherited (invoked) in the workflow.')
+    doc_template_lmp = 'The path to input.lammps template'
+    doc_template_plm = 'The path to input.plumed template'
+
+    args = [
+        Argument("lmp", str, optional=True, doc=doc_template_lmp),
+        Argument("plm", str, optional=True, doc=doc_template_plm),
+    ]
+    return Argument("template", list, args, [], optional=True, repeat=False, doc=doc_template)
+
+
+def model_devi_jobs_rev_mat_args() -> Argument:
+    doc_rev_mat = ('revise matrix for revising variable(s) defined in the template into the specific values (iteration-resolved).'
+                   ' Values will be broadcasted for all tasks within the iteration invoking this key.')
+    doc_rev_mat_lmp = 'revise matrix for revising variable(s) defined in the lammps template into the specific values (iteration-resolved).'
+    doc_rev_mat_plm = 'revise matrix for revising variable(s) defined in the plumed template into specific values(iteration-resolved)'
+
+    args = [
+        Argument("lmp", dict, optional=True, doc=doc_rev_mat_lmp),
+        Argument("plm", dict, optional=True, doc=doc_rev_mat_plm),
+    ]
+    return Argument("rev_mat", list, args, [], optional=True, repeat=False, doc=doc_rev_mat)
+
+
 def model_devi_jobs_args() -> List[Argument]:
     # this may be not correct
+    doc_sys_rev_mat = ('system-resolved revise matrix for revising variable(s) defined in the template into specific values. '
+                       'Values should be individually assigned to each system adopted by this iteration, through a dictionary '
+                       'where first-level keys are values of sys_idx of this iteration.')
     doc_sys_idx = 'Systems to be selected as the initial structure of MD and be explored. The index corresponds exactly to the sys_configs.'
     doc_temps = 'Temperature (K) in MD.'
     doc_press = 'Pressure (Bar) in MD. Required when ensemble is npt.'
@@ -118,8 +148,11 @@ def model_devi_jobs_args() -> List[Argument]:
     doc_model_devi_v_trust_hi = 'Upper bound of virial for the selection. If dict, should be set for each index in sys_idx, respectively. Should be used with DeePMD-kit v2.x.'
 
     args = [
+        model_devi_jobs_template_args(), 
+        model_devi_jobs_rev_mat_args(),
+        Argument("sys_rev_mat", dict, optional=True, doc=doc_sys_rev_mat),
         Argument("sys_idx", list, optional=False, doc=doc_sys_idx),
-        Argument("temps", list, optional=False, doc=doc_temps),
+        Argument("temps", list, optional=True, doc=doc_temps),
         Argument("press", list, optional=True, doc=doc_press),
         Argument("trj_freq", int, optional=False, doc=doc_trj_freq),
         Argument("nsteps", int, optional=False, doc=doc_nsteps),
@@ -160,6 +193,8 @@ def model_devi_lmp_args() -> List[Argument]:
     doc_model_devi_clean_traj = 'If type of model_devi_clean_traj is bool type then it denote whether to clean traj folders in MD since they are too large. If it is Int type, then the most recent n iterations of traj folders will be retained, others will be removed.'
     doc_model_devi_merge_traj = 'If model_devi_merge_traj is set as True, only all.lammpstrj will be generated, instead of lots of small traj files.'
     doc_model_devi_nopbc = 'Assume open boundary condition in MD simulations.'
+    doc_model_devi_plumed = '' # looking forward to update
+    doc_model_devi_plumed_path = '' # looking forward to update
     doc_shuffle_poscar = 'Shuffle atoms of each frame before running simulations. The purpose is to sample the element occupation of alloys.'
     doc_use_relative = 'Calculate relative force model deviation.'
     doc_epsilon = 'The level parameter for computing the relative force model deviation.'
@@ -194,10 +229,14 @@ def model_devi_lmp_args() -> List[Argument]:
                  doc=doc_model_devi_f_avg_relative),
         Argument("model_devi_clean_traj", [
                  bool, int], optional=True, default=True , doc=doc_model_devi_clean_traj),
-        Argument("model_devi_merge_traj", [
-                 bool], optional=True, default=False , doc=doc_model_devi_merge_traj),
+        Argument("model_devi_merge_traj", 
+                 bool, optional=True, default=False , doc=doc_model_devi_merge_traj),
         Argument("model_devi_nopbc", bool, optional=True, default=False,
                  doc=doc_model_devi_nopbc),
+        Argument("model_devi_plumed", 
+                 bool, optional=True, default=False , doc=doc_model_devi_plumed),
+        Argument("model_devi_plumed_path", 
+                 bool, optional=True, default=False , doc=doc_model_devi_plumed_path),
         Argument("shuffle_poscar", bool, optional=True, default=False, doc=doc_shuffle_poscar),
         Argument("use_relative", bool, optional=True, default=False, doc=doc_use_relative),
         Argument("epsilon", float, optional=True, doc=doc_epsilon),