refactor configurations

dpgen2/conf/__init__.py

@@ -0,0 +1,14 @@
+from .conf_generator import (
+    ConfGenerator,
+)
+from .alloy_conf import (
+    AlloyConfGenerator,
+)
+from .file_conf import (
+    FileConfGenerator,
+)
+
+conf_styles = {
+    "alloy" : AlloyConfGenerator,
+    "file" : FileConfGenerator,
+}

dpgen2/utils/alloy_conf.py → dpgen2/conf/alloy_conf.py

@@ -11,6 +11,95 @@
     Variant,
 )
 from .unit_cells import generate_unit_cell
+from .conf_generator import ConfGenerator
+
+class AlloyConfGenerator(ConfGenerator):
+    """
+    Parameters
+    ----------
+    numb_confs  int
+                Number of configurations to generate
+    lattice     Union[dpdata.System, Tuple[str,float]]
+                Lattice of the alloy confs. can be 
+                `dpdata.System`: lattice in `dpdata.System`
+                `Tuple[str, float]`: pair of lattice type and lattice constant.
+                lattice type can be "bcc", "fcc", "hcp", "sc" or "diamond"
+    replicate   Union[List[int], Tuple[int], int]
+                replicate of the lattice
+    concentration   List[List[float]] or List[float] or None
+                If `List[float]`, the concentrations of each element. The length of
+                the list should be the same as the `type_map`.
+                If `List[List[float]]`, a list of concentrations (`List[float]`) is
+                randomly picked from the List.
+                If `None`, the elements are assumed to be of equal concentration.
+    cell_pert_frac  float
+                fraction of cell perturbation
+    atom_pert_dist  float
+                the atom perturbation distance (unit angstrom).
+    """
+    def __init__ (
+            self,
+            numb_confs,
+            lattice : Union[dpdata.System, Tuple[str,float]],
+            replicate : Union[List[int], Tuple[int], int, None] = None,
+            concentration: Union[List[List[float]], List[float], None] = None,
+            cell_pert_frac: float = 0.0,
+            atom_pert_dist: float = 0.0,
+    ):
+        self.numb_confs = numb_confs
+        self.lattice = lattice
+        self.replicate = replicate
+        self.concentration = concentration
+        self.cell_pert_frac = cell_pert_frac
+        self.atom_pert_dist = atom_pert_dist
+
+    def generate(
+            self,
+            type_map,            
+    ) -> dpdata.MultiSystems:
+        r"""Method of generating configurations.
+
+        Parameters
+        ----------
+        type_map: List[str]
+                The type map.
+
+        Returns
+        -------
+        confs:  dpdata.MultiSystems
+                The returned configurations in `dpdata.MultiSystems` format
+
+        """
+        ms = dpdata.MultiSystems(type_map=type_map)
+        ac = AlloyConf(self.lattice, type_map, replicate=self.replicate)
+        systems = ac.generate_systems(
+            self.numb_confs, 
+            concentration=self.concentration,
+            cell_pert_frac=self.cell_pert_frac,
+            atom_pert_dist=self.atom_pert_dist,
+        )
+        for ss in systems:
+            ms.append(ss)
+        return ms    
+
+    @staticmethod
+    def args() -> List[Argument]:
+        doc_numb_confs = 'The number of configurations to generate'
+        doc_lattice = 'The lattice. Should be a list providing [ "lattice_type", lattice_const ], or a list providing [ "/path/to/dpdata/system", "fmt" ]. The two styles are distinguished by the type of the second element.'
+        doc_replicate = 'The number of replicates in each direction'
+        doc_concentration = 'The concentration of each element. If None all elements have the same concentration'
+        doc_cell_pert_frac = 'The faction of cell perturbation'
+        doc_atom_pert_dist = 'The distance of atomic position perturbation'
+
+        return [
+            Argument("numb_confs", int, optional=True, default=1, doc=doc_numb_confs),
+            Argument("lattice", [list,tuple], doc=doc_lattice),
+            Argument("replicate", list, optional=True, default=None, doc=doc_replicate),
+            Argument("concentration", list, optional=True, default=None, doc=doc_concentration),
+            Argument("cell_pert_frac", float, optional=True, default=0.0, doc=doc_cell_pert_frac),
+            Argument("atom_pert_dist", float, optional=True, default=0.0, doc=doc_atom_pert_dist),
+        ]
+
 
 class AlloyConf():
     """

dpgen2/conf/conf_generator.py

@@ -0,0 +1,95 @@
+import dpdata
+import dargs
+import tempfile
+from pathlib import Path
+from typing import (
+    List, Dict,
+)
+from abc import (
+    ABC, abstractmethod,
+)
+
+class ConfGenerator(ABC):
+    @abstractmethod
+    def generate(
+            self,
+            type_map,            
+    ) -> dpdata.MultiSystems:
+        r"""Method of generating configurations.
+
+        Parameters
+        ----------
+        type_map: List[str]
+                The type map.
+
+        Returns
+        -------
+        confs:  dpdata.MultiSystems
+                The returned configurations in `dpdata.MultiSystems` format
+
+        """
+        pass
+
+
+    def get_file_content(
+            self, 
+            type_map,
+            fmt='lammps/lmp',
+    ) -> List[str]:
+        r"""Get the file content of configurations
+
+        Parameters
+        ----------
+        type_map: List[str]
+                The type map.
+
+        Returns
+        -------
+        conf_list: List[str]
+                A list of file content of configurations.
+
+        """
+        ret = []
+        ms = self.generate(type_map)
+        for ii in range(len(ms)):
+            ss = ms[ii]
+            for jj in range(ss.get_nframes()):
+                with tempfile.NamedTemporaryFile() as ft:
+                    tf = Path(ft.name)
+                    ss[jj].to(fmt, tf)
+                    ret.append(tf.read_text())
+        return ret
+
+
+    @staticmethod
+    @abstractmethod
+    def args() -> List[dargs.Argument]:
+        pass
+
+
+    @classmethod
+    def normalize_config(
+            cls, 
+            data: Dict={}, 
+            strict: bool=True,
+    ) -> Dict:
+        r"""Normalized the argument.
+
+        Parameters
+        ----------
+        data: Dict
+            The input dict of arguments.
+        strict: bool
+            Strictly check the arguments.
+
+        Returns
+        -------
+        data: Dict
+            The normalized arguments.
+
+        """
+        ta = cls.args()
+        base = dargs.Argument("base", dict, ta)
+        data = base.normalize_value(data, trim_pattern="_*")
+        base.check_value(data, strict=strict)
+        return data

dpgen2/conf/file_conf.py

@@ -0,0 +1,58 @@
+import os
+import dpdata
+import glob
+from pathlib import Path
+from .conf_generator import ConfGenerator
+from typing import (
+    Optional, Union, List, Tuple
+)
+from dargs import (
+    Argument,
+    Variant,
+)
+
+
+class FileConfGenerator(ConfGenerator):
+    def __init__(
+            self,
+            files : Union[str,List[str]],
+            fmt : str = 'auto',
+            prefix : Optional[str] = None,
+    ):
+        if not isinstance(files, list):
+            assert(isinstance(files, str))
+            files = [files]
+        if prefix is not None:
+            pfiles = [Path(prefix) / Path(ii) for ii in files]
+        else:
+            pfiles = [Path(ii) for ii in files]
+        self.files = []
+        for ii in pfiles:
+            ff = glob.glob(str(ii.absolute()))
+            ff.sort()
+            self.files += ff
+        self.fmt = fmt
+
+
+    def generate(
+            self,
+            type_map,
+    ) -> dpdata.MultiSystems:
+        ms = dpdata.MultiSystems(type_map=type_map)
+        for ff in self.files:
+            ms.append(dpdata.System(ff, fmt=self.fmt))
+        return ms
+
+
+    @staticmethod
+    def args() -> List[Argument]:
+        doc_files = "The paths to the configuration files. widecards are supported."
+        doc_prefix = "The prefix of file paths."
+        doc_fmt = "The format (dpdata accepted formats) of the files."
+
+        return [
+            Argument("files", [str, list], optional=False, doc=doc_files),
+            Argument("prefix", str, optional=True, default=None, doc=doc_prefix),
+            Argument("fmt", str, optional=True, default='auto', doc=doc_fmt),
+        ]
+

dpgen2/entrypoint/args.py

@@ -3,6 +3,7 @@
     Argument,
     Variant,
 )
+import dpgen2
 from dpgen2.constants import default_image
 from dflow.plugins.lebesgue import LebesgueExecutor
 from dpgen2.op.run_dp_train import RunDPTrain
@@ -12,6 +13,7 @@
     normalize_step_dict,
 )
 from dpgen2.fp import fp_styles
+from dpgen2.conf import conf_styles
 
 def dp_train_args():
     doc_numb_models = "Number of models trained for evaluating the model deviation"
@@ -53,7 +55,10 @@ def lmp_args():
         Argument("v_trust_lo", float, optional=True, default=None, doc=doc_v_trust_lo),
         Argument("v_trust_hi", float, optional=True, default=None, doc=doc_v_trust_hi),
         Argument("configuration_prefix", str, optional=True, default=None, doc=doc_configuration_prefix),
-        Argument("configurations", list, optional=False, doc=doc_configuration, alias=["configuration"]),
+        Argument("configurations", list, 
+                 [], [variant_conf()], 
+                 optional=False, repeat=True,                 
+                 doc=doc_configuration, alias=["configuration"]),
         Argument("stages", list, optional=False, doc=doc_stages),
     ]
 
@@ -99,6 +104,16 @@ def variant_fp():
     return Variant("type", fp_list, doc=doc)
 
 
+def variant_conf():
+    doc = "the type of the configuration generator"
+    var_list = []
+    for kk in conf_styles.keys():
+        var_list.append(
+            Argument(kk, dict, conf_styles[kk].args())
+        )
+    return Variant("type", var_list, doc=doc,)
+
+
 def input_args():
     doc_type_map = 'The type map. e.g. ["Al", "Mg"]. Al and Mg will have type 0 and 1, respectively.'
     doc_mass_map = "The mass map. e.g. [27., 24.]. Al and Mg will be set with mass 27. and 24. amu, respectively."

dpgen2/entrypoint/common.py

@@ -3,8 +3,6 @@
 from dpgen2.utils import (
     dump_object_to_file,
     load_object_from_file,
-    normalize_alloy_conf_dict,
-    generate_alloy_conf_file_content,
     sort_slice_ops,
     print_keys_in_nice_format,
     workflow_config_from_dict,

dpgen2/entrypoint/submit.py

@@ -43,7 +43,12 @@
 from dpgen2.flow import (
     ConcurrentLearning,
 )
-from dpgen2.fp import fp_styles
+from dpgen2.fp import (
+    fp_styles,
+)
+from dpgen2.conf import (
+    conf_styles,
+)
 from dpgen2.exploration.scheduler import (
     ExplorationScheduler,
     ConvergenceCheckStageScheduler,
@@ -66,8 +71,6 @@
 from dpgen2.utils import (
     dump_object_to_file,
     load_object_from_file,
-    normalize_alloy_conf_dict,
-    generate_alloy_conf_file_content,
     sort_slice_ops,
     print_keys_in_nice_format,
     workflow_config_from_dict,
@@ -169,32 +172,6 @@ def make_concurrent_learning_op (
     return dpgen_op
 
 
-def make_conf_list(
-        conf_list,
-        type_map,
-        fmt='vasp/poscar'
-):
-    # load confs from files
-    if isinstance(conf_list, list):
-        conf_list_fname = []
-        for jj in conf_list:
-            confs = sorted(glob.glob(jj))
-            conf_list_fname = conf_list_fname + confs
-        conf_list = []
-        for ii in conf_list_fname:
-            ss = dpdata.System(ii, type_map=type_map, fmt=fmt)
-            ss_str = dpdata.lammps.lmp.from_system_data(ss, 0)
-            conf_list.append(ss_str)
-    # generate alloy confs
-    elif isinstance(conf_list, dict):
-        conf_list['type_map'] = type_map
-        i_dict = normalize_alloy_conf_dict(conf_list)
-        conf_list = generate_alloy_conf_file_content(**i_dict)
-    else:
-        raise RuntimeError('unknown input format of conf_list: ', type(conf_list))
-    return conf_list
-
-
 def make_naive_exploration_scheduler(
         config,
         old_style = False,
@@ -218,13 +195,15 @@ def make_naive_exploration_scheduler(
 
     sys_configs_lmp = []
     for sys_config in sys_configs:
-        sys_configs_lmp.append(make_conf_list(sys_config, type_map))
+        conf_style = sys_config.pop("type")        
+        generator = conf_styles[conf_style](**sys_config)
+        sys_configs_lmp.append(generator.get_file_content(type_map))
 
     for job_ in model_devi_jobs:
         if not isinstance(job_, list):
             job = [job_]
         else:
-            job = job_        
+            job = job_
         # stage
         stage = ExplorationStage()
         for jj in job: