add support for fpop ABACUS (#174)

Signed-off-by: zjgemi <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Han Wang <[email protected]>

dpgen2/entrypoint/submit.py

@@ -206,7 +206,7 @@ def make_naive_exploration_scheduler(
     sys_configs = (
         config["sys_configs"] if old_style else config["explore"]["configurations"]
     )
-    sys_prefix = config.get("sys_prefix")
+    sys_prefix = config["explore"].get("configuration_prefix")
     if sys_prefix is not None:
         for ii in range(len(sys_configs)):
             if isinstance(sys_configs[ii], list):

dpgen2/fp/__init__.py

@@ -1,3 +1,8 @@
+from .abacus import (
+    FpOpAbacusInputs,
+    PrepFpOpAbacus,
+    RunFpOpAbacus,
+)
 from .deepmd import (
     DeepmdInputs,
     PrepDeepmd,
@@ -30,4 +35,9 @@
         "prep": PrepDeepmd,
         "run": RunDeepmd,
     },
+    "fpop_abacus": {
+        "inputs": FpOpAbacusInputs,
+        "prep": PrepFpOpAbacus,
+        "run": RunFpOpAbacus,
+    },
 }

dpgen2/fp/abacus.py

@@ -0,0 +1,197 @@
+from pathlib import (
+    Path,
+)
+from typing import (
+    List,
+)
+
+import dpdata
+from dargs import (
+    Argument,
+)
+from dflow.python import (
+    OP,
+    OPIO,
+    Artifact,
+    BigParameter,
+    OPIOSign,
+)
+
+try:
+    from fpop.abacus import (
+        AbacusInputs,
+        PrepAbacus,
+        RunAbacus,
+    )
+except ModuleNotFoundError:
+    AbacusInputs = PrepAbacus = RunAbacus = object
+
+from ..constants import (
+    fp_default_out_data_name,
+)
+
+
+class FpOpAbacusInputs(AbacusInputs):  # type: ignore
+    @staticmethod
+    def args():
+        doc_input_file = "A template INPUT file."
+        doc_pp_files = (
+            "The pseudopotential files for the elements. "
+            'For example: {"H": "/path/to/H.upf", "O": "/path/to/O.upf"}.'
+        )
+        doc_element_mass = (
+            "Specify the mass of some elements. "
+            'For example: {"H": 1.0079, "O": 15.9994}.'
+        )
+        doc_kpt_file = "The KPT file, by default None."
+        doc_orb_files = (
+            "The numerical orbital fiels for the elements, "
+            "by default None. "
+            'For example: {"H": "/path/to/H.orb", "O": "/path/to/O.orb"}.'
+        )
+        doc_deepks_descriptor = "The deepks descriptor file, by default None."
+        doc_deepks_model = "The deepks model file, by default None."
+        return [
+            Argument("input_file", str, optional=False, doc=doc_input_file),
+            Argument("pp_files", dict, optional=False, doc=doc_pp_files),
+            Argument(
+                "element_mass", dict, optional=True, default=None, doc=doc_element_mass
+            ),
+            Argument("kpt_file", str, optional=True, default=None, doc=doc_kpt_file),
+            Argument("orb_files", dict, optional=True, default=None, doc=doc_orb_files),
+            Argument(
+                "deepks_descriptor",
+                str,
+                optional=True,
+                default=None,
+                doc=doc_deepks_descriptor,
+            ),
+            Argument(
+                "deepks_model", str, optional=True, default=None, doc=doc_deepks_model
+            ),
+        ]
+
+
+class PrepFpOpAbacus(OP):
+    @classmethod
+    def get_input_sign(cls):
+        return OPIOSign(
+            {
+                "config": BigParameter(dict),
+                "type_map": List[str],
+                "confs": Artifact(List[Path]),
+            }
+        )
+
+    @classmethod
+    def get_output_sign(cls):
+        return OPIOSign(
+            {
+                "task_names": BigParameter(List[str]),
+                "task_paths": Artifact(List[Path]),
+            }
+        )
+
+    @OP.exec_sign_check
+    def execute(
+        self,
+        ip: OPIO,
+    ) -> OPIO:
+        confs = []
+        # remove atom types with 0 atom from type map, for abacus need pp_files
+        # for all atom types in the type map
+        for p in ip["confs"]:
+            for f in p.rglob("type.raw"):
+                system = f.parent
+                s = dpdata.System(system, fmt="deepmd/npy")
+                atom_numbs = []
+                atom_names = []
+                for numb, name in zip(s["atom_numbs"], s["atom_names"]):  # type: ignore https://github.com/microsoft/pyright/issues/5620
+                    if numb > 0:
+                        atom_numbs.append(numb)
+                        atom_names.append(name)
+                if atom_names != s["atom_names"]:
+                    for i, t in enumerate(s["atom_types"]):  # type: ignore https://github.com/microsoft/pyright/issues/5620
+                        s["atom_types"][i] = atom_names.index(s["atom_names"][t])  # type: ignore https://github.com/microsoft/pyright/issues/5620
+                    s.data["atom_numbs"] = atom_numbs
+                    s.data["atom_names"] = atom_names
+                    target = "output/%s" % system
+                    s.to("deepmd/npy", target)
+                    confs.append(Path(target))
+                else:
+                    confs.append(system)
+        op_in = OPIO(
+            {
+                "inputs": ip["config"]["inputs"],
+                "type_map": ip["type_map"],
+                "confs": confs,
+                "prep_image_config": ip["config"].get("prep", {}),
+            }
+        )
+        op = PrepAbacus()
+        return op.execute(op_in)  # type: ignore in the case of not importing fpop
+
+
+class RunFpOpAbacus(OP):
+    @classmethod
+    def get_input_sign(cls):
+        return OPIOSign(
+            {
+                "config": BigParameter(dict),
+                "task_name": BigParameter(str),
+                "task_path": Artifact(Path),
+            }
+        )
+
+    @classmethod
+    def get_output_sign(cls):
+        return OPIOSign(
+            {
+                "log": Artifact(Path),
+                "labeled_data": Artifact(Path),
+            }
+        )
+
+    @OP.exec_sign_check
+    def execute(
+        self,
+        ip: OPIO,
+    ) -> OPIO:
+        run_config = ip["config"].get("run", {})
+        op_in = OPIO(
+            {
+                "task_name": ip["task_name"],
+                "task_path": ip["task_path"],
+                "backward_list": [],
+                "run_image_config": run_config,
+            }
+        )
+        op = RunAbacus()
+        op_out = op.execute(op_in)  # type: ignore in the case of not importing fpop
+        workdir = op_out["backward_dir"].parent
+
+        # convert the output to deepmd/npy format
+        sys = dpdata.LabeledSystem(str(workdir), fmt="abacus/scf")
+        out_name = run_config.get("out", fp_default_out_data_name)
+        sys.to("deepmd/npy", workdir / out_name)
+
+        return OPIO(
+            {
+                "log": workdir / "log",
+                "labeled_data": workdir / out_name,
+            }
+        )
+
+    @staticmethod
+    def args():
+        doc_cmd = "The command of abacus"
+        doc_out = (
+            "The output dir name of labeled data. "
+            "In `deepmd/npy` format provided by `dpdata`."
+        )
+        return [
+            Argument("command", str, optional=True, default="abacus", doc=doc_cmd),
+            Argument(
+                "out", str, optional=True, default=fp_default_out_data_name, doc=doc_out
+            ),
+        ]

pyproject.toml

@@ -21,7 +21,9 @@ dependencies = [
 	     'dargs>=0.3.1',
 	     'scipy',
 	     'lbg',
-          'packaging',
+	     'packaging',
+	     'fpop',
+	     'dpgui',
 ]
 requires-python = ">=3.7"
 readme = "README.md"