Template lammps task group (#90)

* implementing template lmp and plm

* implement plm

* add make task group from config

* implement task_group from config to submit

* add aliases for template fnames

* fix bugs in submit

* fix bugs in submit

* update example

* docs and examples

* add missing file

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

docs/dpgen2_configs.rst

@@ -0,0 +1,50 @@
+DPGEN2 configurations
+=====================
+
+Op configs
+----------
+
+RunDPTrain
+^^^^^^^^^^
+.. _rundptrainargs:
+.. dargs::
+   :module: dpgen2.op.run_dp_train
+   :func: config_args
+
+RunLmp
+^^^^^^
+.. _runlmpargs:
+.. dargs::
+   :module: dpgen2.op.run_lmp
+   :func: config_args
+
+RunVasp
+^^^^^^^
+.. _runvaspargs:
+.. dargs::
+   :module: dpgen2.op.run_vasp
+   :func: config_args
+
+
+Alloy configs
+-------------
+.. _alloy configs:
+.. dargs::
+   :module: dpgen2.utils.alloy_conf
+   :func: generate_alloy_conf_args
+
+
+Task group configs
+------------------
+.. _task_group_configs:
+.. dargs::
+   :module: dpgen2.exploration.task
+   :func: task_group_args
+
+
+Step configs
+------------
+.. _stepconfigargs:
+.. dargs::
+   :module: dpgen2.utils.step_config
+   :func: step_conf_args

docs/index.rst

@@ -16,9 +16,7 @@ DPGEN2 is the 2nd generation of the Deep Potential GENerator.
    cli
    input.md
    submit_args
-   op_configs
-   alloy_configs
-   step_configs
+   dpgen2_configs
 
 .. _developer-guide:
 

docs/input.md

@@ -91,8 +91,29 @@ This section defines how the configuration space is explored.
 		}
 	],
 	"stages":	[
-		{ "_idx": 0, "ensemble": "nvt", "nsteps": 20, "press": null, "conf_idx": [0], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 },
-		{ "_idx": 1, "ensemble": "nvt", "nsteps": 20, "press": null, "conf_idx": [1], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 }
+	    [
+		{
+		    "_comment" : "stage 0, task group 0",
+		    "type" : "lmp-md",
+		    "ensemble": "nvt", "nsteps":  50, "temps": [50, 100], "trj_freq": 10,
+		    "conf_idx": [0], "n_sample" : 3
+		},
+		{
+		    "_comment" : "stage 1, task group 0",
+		    "type" : "lmp-template",
+		    "lmp" : "template.lammps", "plm" : "template.plumed", 
+		    "trj_freq" : 10, "revisions" : {"V_NSTEPS" : [40], "V_TEMP" : [150, 200]},
+		    "conf_idx": [0], "n_sample" : 3
+		}
+	    ],
+	    [
+		{
+		    "_comment" : "stage 1, task group 0",
+		    "type" : "lmp-md",
+		    "ensemble": "npt", "nsteps":  50, "press": [1e0], "temps": [50, 100, 200], "trj_freq": 10,
+		    "conf_idx": [1], "n_sample" : 3
+		}
+	    ],
 	],
 }
 ```
@@ -103,7 +124,9 @@ The `"configurations"` provides the initial configurations (coordinates of atoms
 - `list[str]`: The strings provides the path to the configuration files.
 - `dict`: Automatic alloy configuration generator. See [the detailed doc](alloy_configs) of the allowed keys.
 
-The `"stages"` defines the exploration stages. It is a list of `dict`s, with each `dict` defining a stage. The `"ensemble"`, `"nsteps"`, `"press"`, `"temps"`, `"traj_freq"` keys are self-explanatory. `"conf_idx"` pickes initial configurations of DPMD simulations from the `"configurations"` list, it provides the index of the element in the `"configurations"` list. `"n_sample"` tells the number of confgiruations randomly sampled from the set picked by `"conf_idx"` for each thermodynamic state. All configurations picked by `"conf_idx"` has the same possibility to be sampled. The default value of `"n_sample"` is `null`, in this case all picked configurations are sampled. In the example, each stage have 3 samples and 2 thermodynamic states (NVT, T=50 and 100K), then each iteration run 3x2=6 NVT DPMD simulatins.
+The `"stages"` defines the exploration stages. It is of type `list[list[dict]]`. The outer `list` enumerate the exploration stages, the inner list enumerate the task groups of the stage. Each `dict` defines a stage. See [the full documentation of the target group](task_group_configs) for writting task groups.
+
+`"n_sample"` tells the number of confgiruations randomly sampled from the set picked by `"conf_idx"` from `configurations` for each exploration task. All configurations has the equal possibility to be sampled. The default value of `"n_sample"` is `null`, in this case all picked configurations are sampled. In the example, we have 3 samples for stage 0 task group 0 and 2 thermodynamic states (NVT, T=50 and 100K), then the task group has 3x2=6 NVT DPMD tasks.
 
 
 ### FP

dpgen2/constants.py

@@ -7,6 +7,8 @@
 lmp_task_pattern = 'task.' + lmp_index_pattern
 lmp_conf_name = 'conf.lmp'
 lmp_input_name = 'in.lammps'
+plm_input_name = 'input.plumed'
+plm_output_name = 'output.plumed'
 lmp_traj_name = 'traj.dump'
 lmp_log_name = 'log.lammps'
 lmp_model_devi_name = 'model_devi.out'

dpgen2/entrypoint/submit.py

@@ -56,6 +56,8 @@
     ExplorationStage,
     ExplorationTask,
     NPTTaskGroup,
+    LmpTemplateTaskGroup,
+    make_task_group_from_config,
 )
 from dpgen2.exploration.selector import (
     ConfSelectorLammpsFrames,
@@ -217,42 +219,33 @@ def make_naive_exploration_scheduler(
     fatal_at_max = config.get('fatal_at_max', True) if old_style else config['explore']['fatal_at_max']
     scheduler = ExplorationScheduler()
 
-    for job in model_devi_jobs:
-        # task group
-        tgroup = NPTTaskGroup()
-        ##  ignore the expansion of sys_idx
-        # get all file names of md initial configurations
-        try:
-            sys_idx = job['sys_idx']
-        except KeyError:
-            sys_idx = job['conf_idx']
-        conf_list = []        
-        for ii in sys_idx:
-            conf_list += make_conf_list(sys_configs[ii], type_map)
-        # add the list to task group
-        n_sample = job.get('n_sample')
-        tgroup.set_conf(
-            conf_list,
-            n_sample=n_sample,
-        )
-        temps = job['temps']
-        press = job['press']
-        trj_freq = job['trj_freq']
-        nsteps = job['nsteps']
-        ensemble = job['ensemble']
-        # add md settings
-        tgroup.set_md(
-            numb_models,
-            mass_map,
-            temps = temps,
-            press = press,
-            ens = ensemble,
-            nsteps = nsteps,
-        )
-        tasks = tgroup.make_task()
+    for job_ in model_devi_jobs:
+        if not isinstance(job_, list):
+            job = [job_]
+        else:
+            job = job_        
         # stage
         stage = ExplorationStage()
-        stage.add_task_group(tasks)
+        for jj in job:
+            n_sample = jj.pop('n_sample')
+            ##  ignore the expansion of sys_idx
+            # get all file names of md initial configurations
+            try:
+                sys_idx = jj.pop('sys_idx')
+            except KeyError:
+                sys_idx = jj.pop('conf_idx')
+            conf_list = []        
+            for ii in sys_idx:
+                conf_list += make_conf_list(sys_configs[ii], type_map)
+            # make task group
+            tgroup = make_task_group_from_config(numb_models, mass_map, jj)
+            # add the list to task group
+            tgroup.set_conf(
+                conf_list,
+                n_sample=n_sample,
+            )
+            tasks = tgroup.make_task()
+            stage.add_task_group(tasks)
         # trust level
         trust_level = TrustLevel(
             config['model_devi_f_trust_lo'] if old_style else config['explore']['f_trust_lo'],

dpgen2/exploration/task/__init__.py

@@ -2,9 +2,21 @@
     ExplorationTask,
     ExplorationTaskGroup,
 )
+from .conf_sampling_task_group import (
+    ConfSamplingTaskGroup,
+)
 from .npt_task_group import (
     NPTTaskGroup,
 )
+from .lmp_template_task_group import (
+    LmpTemplateTaskGroup,
+)
+from .make_task_group_from_config import(
+    normalize as normalize_task_group_config,
+    make_task_group_from_config,
+    variant_task_group,
+    task_group_args,
+)
 from .stage import (
     ExplorationStage,
 )

dpgen2/exploration/task/conf_sampling_task_group.py

@@ -0,0 +1,63 @@
+import itertools, random
+from typing import (
+    List,
+)
+from . import (
+    ExplorationTask,
+    ExplorationTaskGroup,
+)
+from dpgen2.constants import (
+    lmp_conf_name, 
+    lmp_input_name,
+    model_name_pattern,
+)
+
+class ConfSamplingTaskGroup(ExplorationTaskGroup):
+    def __init__(
+            self, 
+    ):
+        super().__init__()
+        self.conf_set = False
+
+    def set_conf(
+            self,
+            conf_list : List[str],
+            n_sample : int = None,
+            random_sample : bool = False,
+    ):
+        """
+        Set the configurations of exploration
+
+        Parameters
+        ----------
+        conf_list       str
+                        A list of file contents
+        n_sample        int
+                        Number of samples drawn from the conf list each time 
+                        `make_task` is called. If set to `None`, 
+                        `n_sample` is set to length of the conf_list.
+        random_sample   bool
+                        If true the confs are randomly sampled, otherwise are
+                        consecutively sampled from the conf_list
+        """
+        self.conf_list = conf_list
+        if n_sample is None:
+            self.n_sample = len(self.conf_list)
+        else:
+            self.n_sample = n_sample
+        self.random_sample = random_sample
+        self.conf_queue = []
+        self.conf_set = True
+
+    def _sample_confs(
+            self,
+    ):
+        confs = []
+        for ii in range(self.n_sample):
+            if len(self.conf_queue) == 0:
+                add_list = self.conf_list.copy()
+                if self.random_sample:
+                    random.shuffle(add_list)
+                self.conf_queue += add_list
+            confs.append(self.conf_queue.pop(0))
+        return confs