Adding Knowledge Distillation (#122)

Adding `Knowledge Distillation` in dpgen2, including running `lammps`
and `fp` with the teacher model. Additional data selection strategies
and convergence criteria are not yet finished.

---------

Co-authored-by: Sikai Yao <yaosk@dp.tech>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

dpgen2/entrypoint/args.py

@@ -17,10 +17,33 @@
 from dpgen2.exploration.report import conv_styles
 
 
+def dp_dist_train_args():
+    doc_config = "Configuration of training"
+    doc_template_script = "File names of the template training script. It can be a `List[Dict]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. "
+    dock_student_model_path = "The path of student model"
+
+    return [
+        Argument(
+            "config",
+            dict,
+            RunDPTrain.training_args(),
+            optional=True,
+            default=RunDPTrain.normalize_config({}),
+            doc=doc_config,
+        ),
+        Argument(
+            "template_script", [list, str], optional=False, doc=doc_template_script
+        ),
+        Argument(
+            "student_model_path", str, optional=False, doc=dock_student_model_path
+        ),
+    ]
+
+
 def dp_train_args():
     doc_numb_models = "Number of models trained for evaluating the model deviation"
     doc_config = "Configuration of training"
-    doc_template_script = "File names of the template training script. It can be a `List[Dict]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `Dict`, the models share the same template training script. "
+    doc_template_script = "File names of the template training script. It can be a `List[Dict]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. "
     doc_init_models_paths = "the paths to initial models"
 
     return [
@@ -52,6 +75,7 @@ def variant_train():
         "type",
         [
             Argument("dp", dict, dp_train_args()),
+            Argument("dp-dist", dict, dp_dist_train_args()),
         ],
         doc=doc,
     )

dpgen2/entrypoint/submit.py

@@ -84,6 +84,7 @@
     workflow_config_from_dict,
     matched_step_key,
     bohrium_config_from_dict,
+    BinaryFileInput,
     get_subkey,
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
@@ -126,7 +127,7 @@ def make_concurrent_learning_op(
     cl_step_config: dict = default_config,
     upload_python_packages: Optional[List[os.PathLike]] = None,
 ):
-    if train_style == "dp":
+    if train_style in ("dp", "dp-dist"):
         prep_run_train_op = PrepRunDPTrain(
             "prep-run-dp-train",
             PrepDPTrain,
@@ -358,11 +359,21 @@ def workflow_concurrent_learning(
         else config["step_configs"]["cl_step_config"]
     )
     upload_python_packages = config.get("upload_python_packages", None)
-    init_models_paths = (
-        config.get("training_iter0_model_path", None)
-        if old_style
-        else config["train"].get("init_models_paths", None)
-    )
+
+    if train_style == "dp":
+        init_models_paths = (
+            config.get("training_iter0_model_path", None)
+            if old_style
+            else config["train"].get("init_models_paths", None)
+        )
+    elif train_style == "dp-dist" and not old_style:
+        init_models_paths = [config["train"].get("student_model_path", None)]
+        config["train"]["numb_models"] = 1
+    else:
+        raise RuntimeError(
+            f"unknown params, train_style: {train_style}, old_style: {old_style}"
+        )
+
     if upload_python_packages is not None and isinstance(upload_python_packages, str):
         upload_python_packages = [upload_python_packages]
     if upload_python_packages is not None:
@@ -403,6 +414,14 @@ def workflow_concurrent_learning(
     lmp_config = (
         config.get("lmp_config", {}) if old_style else config["explore"]["config"]
     )
+    if "teacher_model_path" in lmp_config:
+        assert os.path.exists(
+            lmp_config["teacher_model_path"]
+        ), f"No such file: {lmp_config['teacher_model_path']}"
+        lmp_config["teacher_model_path"] = BinaryFileInput(
+            lmp_config["teacher_model_path"], "pb"
+        )
+
     fp_config = config.get("fp_config", {}) if old_style else {}
     if old_style:
         potcar_names = config["fp_pp_files"]
@@ -420,6 +439,16 @@ def workflow_concurrent_learning(
 
     fp_config["inputs"] = fp_inputs
     fp_config["run"] = config["fp"]["run_config"]
+    if fp_style == "deepmd":
+        assert (
+            "teacher_model_path" in fp_config["run"]
+        ), f"Cannot find 'teacher_model_path' in config['fp']['run_config'] when fp_style == 'deepmd'"
+        assert os.path.exists(
+            fp_config["run"]["teacher_model_path"]
+        ), f"No such file: {fp_config['run']['teacher_model_path']}"
+        fp_config["run"]["teacher_model_path"] = BinaryFileInput(
+            fp_config["run"]["teacher_model_path"], "pb"
+        )
 
     init_data_prefix = (
         config.get("init_data_prefix")

dpgen2/fp/__init__.py

@@ -8,6 +8,7 @@
     PrepGaussian,
     RunGaussian,
 )
+from .deepmd import DeepmdInputs, PrepDeepmd, RunDeepmd
 
 fp_styles = {
     "vasp": {
@@ -20,4 +21,9 @@
         "prep": PrepGaussian,
         "run": RunGaussian,
     },
+    "deepmd": {
+        "inputs": DeepmdInputs,
+        "prep": PrepDeepmd,
+        "run": RunDeepmd,
+    },
 }

dpgen2/fp/deepmd.py

@@ -0,0 +1,214 @@
+"""Prep and Run Gaussian tasks."""
+from dflow.python import TransientError, FatalError
+from typing import Tuple, List, Any, Optional
+import dpdata
+from dargs import (
+    dargs,
+    Argument,
+)
+
+from .prep_fp import PrepFp
+from .run_fp import RunFp
+from ..utils import BinaryFileInput
+from dpgen2.constants import fp_default_out_data_name, fp_default_log_name
+from dpgen2.utils.run_command import run_command
+from pathlib import Path
+import os
+
+import numpy as np
+
+# global static variables
+deepmd_input_path = "one_frame_input"
+
+# global static variables
+deepmd_temp_path = "one_frame_temp"
+
+# global static variables
+deepmd_teacher_model = "teacher_model.pb"
+
+
+class DeepmdInputs:
+    @staticmethod
+    def args() -> List[Argument]:
+        return []
+
+    def __init__(self, **kwargs: Any):
+        self.data = kwargs
+
+
+class PrepDeepmd(PrepFp):
+    def prep_task(
+        self,
+        conf_frame: dpdata.System,
+        inputs,
+    ):
+        r"""Define how one Deepmd task is prepared.
+
+        Parameters
+        ----------
+        conf_frame : dpdata.System
+            One frame of configuration in the dpdata format.
+        inputs: str or dict
+            This parameter is useless in deepmd.
+        """
+        conf_frame.to("deepmd/npy", deepmd_input_path)
+
+
+class RunDeepmd(RunFp):
+    def input_files(self) -> List[str]:
+        r"""The mandatory input files to run a Deepmd task.
+
+        Returns
+        -------
+        files: List[str]
+            A list of madatory input files names.
+
+        """
+        return [deepmd_input_path]
+
+    def optional_input_files(self) -> List[str]:
+        r"""The optional input files to run a Deepmd task.
+
+        Returns
+        -------
+        files: List[str]
+            A list of optional input files names.
+
+        """
+        return []
+
+    def run_task(
+        self,
+        teacher_model_path: BinaryFileInput,
+        out: str,
+        log: str,
+    ) -> Tuple[str, str]:
+        r"""Defines how one FP task runs
+
+        Parameters
+        ----------
+        command: str
+            The command of running Deepmd task
+        out: str
+            The name of the output data file.
+
+        Returns
+        -------
+        out_name: str
+            The file name of the output data in the dpdata.LabeledSystem format.
+        log_name: str
+            The file name of the log.
+        """
+        log_name = log
+        out_name = out
+
+        dp, type_map_teacher = self._get_dp_model(teacher_model_path)
+
+        # Run deepmd
+        self._dp_infer(dp, type_map_teacher, out_name)
+
+        run_command(f'echo "job finished!" > {log_name}', shell=True)
+
+        return out_name, log_name
+
+    def _get_dp_model(self, teacher_model_path: BinaryFileInput):
+        from deepmd.infer import DeepPot  # type: ignore
+
+        teacher_model_path.save_as_file(deepmd_teacher_model)
+        dp = DeepPot(deepmd_teacher_model)
+
+        type_map_teacher = dp.get_type_map()
+
+        os.remove(deepmd_teacher_model)
+        return dp, type_map_teacher
+
+    def _prep_input(self, type_map_teacher):
+        ss = dpdata.System(deepmd_input_path, fmt="deepmd/npy")
+        conf_type_map = ss["atom_names"]
+
+        if not set(conf_type_map).issubset(set(type_map_teacher)):
+            err_message = (
+                f"the type map of system ({conf_type_map}) is not subset of "
+                + f"the type map of the teacher model ({type_map_teacher})."
+            )
+            raise FatalError("deepmd labeling failed\n", "err msg", err_message, "\n")
+
+        # make sure the order of elements in sys_type_map
+        # is the same as that in type_map_teacher
+        temp_type_map = [ele for ele in type_map_teacher if ele in set(conf_type_map)]
+
+        ss.apply_type_map(temp_type_map)
+        ss.to("deepmd/npy", deepmd_temp_path)
+        return conf_type_map, temp_type_map
+
+    def _dp_infer(self, dp, type_map_teacher, out_name):
+        conf_type_map, temp_type_map = self._prep_input(type_map_teacher)
+
+        ss = dpdata.System(
+            deepmd_temp_path, fmt="deepmd/npy", type_map=type_map_teacher
+        )
+
+        coord_npy_path_list = list(Path(deepmd_temp_path).glob("*/coord.npy"))
+        assert len(coord_npy_path_list) == 1, coord_npy_path_list
+        coord_npy_path = coord_npy_path_list[0]
+        energy_npy_path = coord_npy_path.parent / "energy.npy"
+        force_npy_path = coord_npy_path.parent / "force.npy"
+        virial_npy_path = coord_npy_path.parent / "virial.npy"
+
+        nframe = ss.get_nframes()
+        coord = ss["coords"]
+        cell = ss["cells"].reshape([nframe, -1])
+        atype = ss["atom_types"].tolist()
+
+        energy, force, virial_force = dp.eval(coord, cell, atype)
+
+        with open(energy_npy_path, "wb") as f:
+            np.save(f, energy)
+        with open(force_npy_path, "wb") as f:
+            np.save(f, force)
+        with open(virial_npy_path, "wb") as f:
+            np.save(f, virial_force)
+
+        ss = dpdata.LabeledSystem(
+            deepmd_temp_path, fmt="deepmd/npy", type_map=temp_type_map
+        )
+        ss.apply_type_map(conf_type_map)
+        ss.to("deepmd/npy", out_name)
+
+    @staticmethod
+    def args() -> List[dargs.Argument]:
+        r"""The argument definition of the `run_task` method.
+
+        Returns
+        -------
+        arguments: List[dargs.Argument]
+            List of dargs.Argument defines the arguments of `run_task` method.
+        """
+
+        doc_deepmd_teacher_model = (
+            "The path of teacher model, which can be loaded by deepmd.infer.DeepPot"
+        )
+        doc_deepmd_log = "The log file name of dp"
+        doc_deepmd_out = "The output dir name of labeled data. In `deepmd/npy` format provided by `dpdata`."
+        return [
+            Argument(
+                "teacher_model_path",
+                [str, BinaryFileInput],
+                optional=False,
+                doc=doc_deepmd_teacher_model,
+            ),
+            Argument(
+                "out",
+                str,
+                optional=True,
+                default=fp_default_out_data_name,
+                doc=doc_deepmd_out,
+            ),
+            Argument(
+                "log",
+                str,
+                optional=True,
+                default=fp_default_log_name,
+                doc=doc_deepmd_log,
+            ),
+        ]

dpgen2/fp/run_fp.py

@@ -174,12 +174,14 @@ def execute(
         with set_directory(work_dir):
             # link input files
             for ii in input_files:
-                if not os.path.isfile(ii):
-                    raise FatalError(f"cannot file file {ii}")
-                iname = ii.name
-                Path(iname).symlink_to(ii)
+                if os.path.isfile(ii) or os.path.isdir(ii):
+                    iname = ii.name
+                    Path(iname).symlink_to(ii)
+                else:
+                    raise FatalError(f"cannot find file {ii}")
+
             for ii in opt_input_files:
-                if os.path.isfile(ii):
+                if os.path.isfile(ii) or os.path.isdir(ii):
                     iname = ii.name
                     Path(iname).symlink_to(ii)
             out_name, log_name = self.run_task(**config)

dpgen2/fp/vasp_input.py

@@ -92,7 +92,7 @@ def args():
         doc_kspacing = "The spacing of k-point sampling. `ksapcing` will overwrite the incar template"
         doc_kgamma = "If the k-mesh includes the gamma point. `kgamma` will overwrite the incar template"
         return [
-            Argument("incar", str, optional=False, doc=doc_pp_files),
+            Argument("incar", str, optional=False, doc=doc_incar),
             Argument("pp_files", dict, optional=False, doc=doc_pp_files),
             Argument("kspacing", float, optional=False, doc=doc_kspacing),
             Argument("kgamma", bool, optional=True, default=True, doc=doc_kgamma),