add driver plugin system framework; migrate DP predict to DPDriver

dpdata/driver.py

@@ -0,0 +1,80 @@
+"""Driver plugin system."""
+from typing import Callable
+from .plugin import Plugin
+from abc import ABC, abstractmethod
+
+
+class Driver(ABC):
+    """The base class for a driver plugin. A driver can
+    label a pure System to generate the LabeledSystem.
+
+    See Also
+    --------
+    dpdata.plugins.deepmd.DPDriver : an example of Driver
+    """
+    __DriverPlugin = Plugin()
+
+    @staticmethod
+    def register(key: str) -> Callable:
+        """Register a driver plugin. Used as decorators.
+
+        Parameter
+        ---------
+        key: str
+            key of the plugin.
+
+        Returns
+        -------
+        Callable
+            decorator of a class
+
+        Examples
+        --------
+        >>> @Driver.register("some_driver")
+        ... class SomeDriver(Driver):
+        ...     pass
+        """
+        return Driver.__DriverPlugin.register(key)
+
+    @staticmethod
+    def get_driver(key: str) -> "Driver":
+        """Get a driver plugin.
+
+        Parameter
+        ---------
+        key: str
+            key of the plugin.
+
+        Returns
+        -------
+        Driver
+            the specific driver class
+
+        Raises
+        ------
+        RuntimeError
+            if the requested driver is not implemented
+        """
+        try:
+            return Driver.__DriverPlugin.plugins[key]
+        except KeyError as e:
+            raise RuntimeError('Unknown driver: ' + key) from e
+
+    def __init__(self, *args, **kwargs) -> None:
+        """Setup the driver."""
+
+    @abstractmethod
+    def label(self, data: dict) -> dict:
+        """Label a system data. Returns new data with energy, forces, and virials.
+
+        Parameters
+        ----------
+        data : dict
+            data with coordinates and atom types
+
+        Returns
+        -------
+        dict
+            labeled data with energies and forces
+        """
+        return NotImplemented

dpdata/plugins/deepmd.py

@@ -1,9 +1,11 @@
+import dpdata
 import dpdata.deepmd.raw
 import dpdata.deepmd.comp
 import dpdata.deepmd.hdf5
 import numpy as np
 import h5py
 from dpdata.format import Format
+from dpdata.driver import Driver
 
 
 @Format.register("deepmd")
@@ -102,3 +104,81 @@ def to_multi_systems(self,
                   directory,
                   **kwargs):
         return ["%s#%s" % (directory, ff) for ff in formulas]
+
+
+@Driver.register("dp")
+@Driver.register("deepmd")
+@Driver.register("deepmd-kit")
+class DPDriver(Driver):
+    """DeePMD-kit driver.
+
+    Parameters
+    ----------
+    dp : deepmd.DeepPot or str
+        The deepmd-kit potential class or the filename of the model.
+
+    Examples
+    --------
+    >>> DPDriver("frozen_model.pb")
+    """
+    def __init__(self, dp: str) -> None:
+        try:
+            # DP 1.x
+            import deepmd.DeepPot as DeepPot
+        except ModuleNotFoundError:
+            # DP 2.x
+            from deepmd.infer import DeepPot
+        if not isinstance(dp, DeepPot):
+            self.dp = DeepPot(dp)
+        else:
+            self.dp = dp
+        self.enable_auto_batch_size = 'auto_batch_size' in DeepPot.__init__.__code__.co_varnames
+
+    def label(self, data: dict) -> dict:
+        """Label a system data by deepmd-kit. Returns new data with energy, forces, and virials.
+
+        Parameters
+        ----------
+        data : dict
+            data with coordinates and atom types
+
+        Returns
+        -------
+        dict
+            labeled data with energies and forces
+        """
+        type_map = self.dp.get_type_map()
+
+        ori_sys = dpdata.System.from_dict({'data': data})
+        ori_sys.sort_atom_names(type_map=type_map)
+        atype = ori_sys['atom_types']
+
+        if not self.enable_auto_batch_size:
+            labeled_sys = dpdata.LabeledSystem()
+            for ss in ori_sys:
+                coord = ss['coords'].reshape((1, ss.get_natoms()*3))
+                if not ss.nopbc:
+                    cell = ss['cells'].reshape((1, 9))
+                else:
+                    cell = None
+                e, f, v = self.dp.eval(coord, cell, atype)
+                data = ss.data
+                data['energies'] = e.reshape((1, 1))
+                data['forces'] = f.reshape((1, ss.get_natoms(), 3))
+                data['virials'] = v.reshape((1, 3, 3))
+                this_sys = dpdata.LabeledSystem.from_dict({'data': data})
+                labeled_sys.append(this_sys)
+            data = labeled_sys.data
+        else:
+            # since v2.0.2, auto batch size is supported
+            coord = ori_sys.data['coords'].reshape((ori_sys.get_nframes(), ori_sys.get_natoms()*3))
+            if not ori_sys.nopbc:
+                cell = ori_sys.data['cells'].reshape((ori_sys.get_nframes(), 9))
+            else:
+                cell = None
+            e, f, v = self.dp.eval(coord, cell, atype)
+            data = ori_sys.data.copy()
+            data['energies'] = e.reshape((ori_sys.get_nframes(), 1))
+            data['forces'] = f.reshape((ori_sys.get_nframes(), ori_sys.get_natoms(), 3))
+            data['virials'] = v.reshape((ori_sys.get_nframes(), 3, 3))
+        return data

dpdata/system.py

@@ -5,6 +5,7 @@
 import numpy as np
 import dpdata.md.pbc
 from copy import deepcopy
+from typing import Any
 from monty.json import MSONable
 from monty.serialization import loadfn,dumpfn
 from dpdata.periodic_table import Element
@@ -15,6 +16,7 @@
 import dpdata.plugins
 from dpdata.plugin import Plugin
 from dpdata.format import Format
+from dpdata.driver import Driver
 
 from dpdata.utils import (
     elements_index_map,
@@ -207,6 +209,7 @@ def dump(self,filename,indent=4):
     def map_atom_types(self,type_map=None):
         """
         Map the atom types of the system
+
         Parameters
         ----------
         type_map :
@@ -640,63 +643,33 @@ def shuffle(self):
             self.data[ii] = self.data[ii][idx]
         return idx
 
-    def predict(self, dp):
+    def predict(self, *args: Any, driver: str="dp", **kwargs: Any) -> "LabeledSystem":
         """
-        Predict energies and forces by deepmd-kit.
+        Predict energies and forces by a driver.
 
         Parameters
         ----------
-        dp : deepmd.DeepPot or str
-            The deepmd-kit potential class or the filename of the model.
+        *args : iterable
+            Arguments passing to the driver
+        driver : str, default=dp
+            The assigned driver. For compatibility, default is dp
+        **kwargs : dict
+            Other arguments passing to the driver
 
         Returns
         -------
         labeled_sys : LabeledSystem
-            The labeled system.
+            A new labeled system.
+
+        Examples
+        --------
+        The default driver is DP:
+        >>> labeled_sys = ori_sys.predict("frozen_model_compressed.pb")
         """
-        try:
-            # DP 1.x
-            import deepmd.DeepPot as DeepPot
-        except ModuleNotFoundError:
-            # DP 2.x
-            from deepmd.infer import DeepPot
-        if not isinstance(dp, DeepPot):
-            dp = DeepPot(dp)
-        type_map = dp.get_type_map()
-        ori_sys = self.copy()
-        ori_sys.sort_atom_names(type_map=type_map)
-        atype = ori_sys['atom_types']
-
-        labeled_sys = LabeledSystem()
-
-        if 'auto_batch_size' not in DeepPot.__init__.__code__.co_varnames:
-            for ss in self:
-                coord = ss['coords'].reshape((1, ss.get_natoms()*3))
-                if not ss.nopbc:
-                    cell = ss['cells'].reshape((1, 9))
-                else:
-                    cell = None
-                e, f, v = dp.eval(coord, cell, atype)
-                data = ss.data
-                data['energies'] = e.reshape((1, 1))
-                data['forces'] = f.reshape((1, ss.get_natoms(), 3))
-                data['virials'] = v.reshape((1, 3, 3))
-                this_sys = LabeledSystem.from_dict({'data': data})
-                labeled_sys.append(this_sys)
-        else:
-            # since v2.0.2, auto batch size is supported
-            coord = self.data['coords'].reshape((self.get_nframes(), self.get_natoms()*3))
-            if not self.nopbc:
-                cell = self.data['cells'].reshape((self.get_nframes(), 9))
-            else:
-                cell = None
-            e, f, v = dp.eval(coord, cell, atype)
-            data = self.data.copy()
-            data['energies'] = e.reshape((self.get_nframes(), 1))
-            data['forces'] = f.reshape((self.get_nframes(), self.get_natoms(), 3))
-            data['virials'] = v.reshape((self.get_nframes(), 3, 3))
-            labeled_sys = LabeledSystem.from_dict({'data': data})
-        return labeled_sys
+        if not isinstance(driver, Driver):
+            driver = Driver.get_driver(driver)(*args, **kwargs)
+        data = driver.label(self.data.copy())
+        return LabeledSystem(data=data)
 
     def pick_atom_idx(self, idx, nopbc=None):
         """Pick atom index
@@ -1026,6 +999,7 @@ def correction(self, hl_sys):
         ----------
         hl_sys: LabeledSystem
             high-level LabeledSystem
+
         Returns
         ----------
         corrected_sys: LabeledSystem
@@ -1208,18 +1182,29 @@ def check_atom_names(self, system):
             system.add_atom_names(new_in_self)
         system.sort_atom_names(type_map=self.atom_names)
 
-    def predict(self, dp):
-        try:
-            # DP 1.x
-            import deepmd.DeepPot as DeepPot
-        except ModuleNotFoundError:
-            # DP 2.x
-            from deepmd.infer import DeepPot
-        if not isinstance(dp, DeepPot):
-            dp = DeepPot(dp)
+    def predict(self, *args: Any, driver="dp", **kwargs: Any) -> "MultiSystems":
+        """
+        Predict energies and forces by a driver.
+
+        Parameters
+        ----------
+        *args : iterable
+            Arguments passing to the driver
+        driver : str, default=dp
+            The assigned driver. For compatibility, default is dp
+        **kwargs : dict
+            Other arguments passing to the driver
+
+        Returns
+        -------
+        MultiSystems
+            A new labeled MultiSystems.
+        """
+        if not isinstance(driver, Driver):
+            driver = Driver.get_driver(driver)(*args, **kwargs)
         new_multisystems = dpdata.MultiSystems()
         for ss in self:
-            new_multisystems.append(ss.predict(dp))
+            new_multisystems.append(ss.predict(*args, driver=driver, **kwargs))
         return new_multisystems
 
     def pick_atom_idx(self, idx, nopbc=None):

tests/test_predict.py

@@ -0,0 +1,34 @@
+import unittest
+import numpy as np
+
+from comp_sys import CompLabeledSys
+from context import dpdata
+
+
+@dpdata.driver.Driver.register("zero")
+class ZeroDriver(dpdata.driver.Driver):
+    def label(self, data):
+        nframes = data['coords'].shape[0]
+        natoms = data['coords'].shape[1]
+        data['energies'] = np.zeros((nframes,))
+        data['forces'] = np.zeros((nframes, natoms, 3))
+        data['virials'] = np.zeros((nframes, 3, 3))
+        return data
+
+
+class TestPredict(unittest.TestCase, CompLabeledSys):
+    def setUp (self) :
+        ori_sys = dpdata.LabeledSystem('poscars/deepmd.h2o.md', 
+                                        fmt = 'deepmd/raw', 
+                                        type_map = ['O', 'H'])
+        self.system_1 = ori_sys.predict(driver="zero")
+        self.system_2 = dpdata.LabeledSystem('poscars/deepmd.h2o.md', 
+                                             fmt = 'deepmd/raw', 
+                                             type_map = ['O', 'H'])
+        for pp in ('energies', 'forces', 'virials'):
+            self.system_2.data[pp][:] = 0.
+
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6