diff --git a/dpdata/driver.py b/dpdata/driver.py new file mode 100644 index 000000000..84419cfe6 --- /dev/null +++ b/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 diff --git a/dpdata/plugins/deepmd.py b/dpdata/plugins/deepmd.py index 9379e9ecd..2798d78ae 100644 --- a/dpdata/plugins/deepmd.py +++ b/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 diff --git a/dpdata/system.py b/dpdata/system.py index b003461cd..df305a6e0 100644 --- a/dpdata/system.py +++ b/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): diff --git a/tests/test_predict.py b/tests/test_predict.py new file mode 100644 index 000000000..1cb816cce --- /dev/null +++ b/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