From e7139103da5039f18060d1d99330a4c54dc06a46 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Feb 2024 08:18:57 -0500 Subject: [PATCH 01/20] refactor DeepEval Signed-off-by: Jinzhe Zeng --- deepmd/__init__.py | 10 + deepmd/infer/deep_dipole.py | 10 + deepmd/infer/deep_dos.py | 138 ++++ deepmd/infer/deep_eval.py | 467 +++++++++++++ deepmd/infer/deep_polar.py | 77 +++ deepmd/infer/deep_pot.py | 146 ++-- deepmd/infer/deep_tensor.py | 222 ++++++ deepmd/infer/deep_wfc.py | 10 + deepmd/pt/infer/deep_eval.py | 222 +++--- deepmd/tf/entrypoints/test.py | 40 +- deepmd/tf/infer/__init__.py | 101 +-- deepmd/tf/infer/data_modifier.py | 4 +- deepmd/tf/infer/deep_dipole.py | 17 +- deepmd/tf/infer/deep_dos.py | 508 +------------- deepmd/tf/infer/deep_eval.py | 1074 +++++++++++++++++++++++++++++- deepmd/tf/infer/deep_polar.py | 169 +---- deepmd/tf/infer/deep_pot.py | 692 +------------------ deepmd/tf/infer/deep_tensor.py | 4 +- deepmd/tf/infer/deep_wfc.py | 70 +- 19 files changed, 2300 insertions(+), 1681 deletions(-) create mode 100644 deepmd/infer/deep_dipole.py create mode 100644 deepmd/infer/deep_dos.py create mode 100644 deepmd/infer/deep_eval.py create mode 100644 deepmd/infer/deep_polar.py create mode 100644 deepmd/infer/deep_tensor.py create mode 100644 deepmd/infer/deep_wfc.py diff --git a/deepmd/__init__.py b/deepmd/__init__.py index f95536db50..4df06e2b9e 100644 --- a/deepmd/__init__.py +++ b/deepmd/__init__.py @@ -14,6 +14,16 @@ __version__, ) + +def DeepPotential(*args, **kwargs): + from deepmd.infer.deep_eval import ( + DeepEval, + ) + + return DeepEval(*args, **kwargs) + + __all__ = [ "__version__", + "DeepPotential", ] diff --git a/deepmd/infer/deep_dipole.py b/deepmd/infer/deep_dipole.py new file mode 100644 index 0000000000..fe51202a43 --- /dev/null +++ b/deepmd/infer/deep_dipole.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.infer.deep_tensor import ( + DeepTensor, +) + + +class DeepDipole(DeepTensor): + @property + def output_tensor_name(self) -> str: + return "dipole" diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py new file mode 100644 index 0000000000..7ece365d95 --- /dev/null +++ b/deepmd/infer/deep_dos.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + Dict, + Optional, + Tuple, +) + +import numpy as np + +from deepmd.model_format.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) + +from .deep_eval import ( + DeepEval, +) + + +class DeepDOS(DeepEval): + """Potential energy model. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + """ + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + "dos", + shape=[-1], + reduciable=True, + atomic=True, + ), + ] + ) + ) + + def eval( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: np.ndarray, + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: Dict[str, Any], + ) -> Tuple[np.ndarray, ...]: + """Evaluate energy, force, and virial. If atomic is True, + also return atomic energy and atomic virial. + + Parameters + ---------- + coords : np.ndarray + The coordinates of the atoms, in shape (nframes, natoms, 3). + cells : np.ndarray + The cell vectors of the system, in shape (nframes, 9). If the system + is not periodic, set it to None. + atom_types : List[int] + The types of the atoms. If mixed_type is False, the shape is (natoms,); + otherwise, the shape is (nframes, natoms). + atomic : bool, optional + Whether to return atomic energy and atomic virial, by default False. + fparam : np.ndarray, optional + The frame parameters, by default None. + aparam : np.ndarray, optional + The atomic parameters, by default None. + mixed_type : bool, optional + Whether the atom_types is mixed type, by default False. + **kwargs : Dict[str, Any] + Keyword arguments. + + Returns + ------- + energy + The energy of the system, in shape (nframes,). + force + The force of the system, in shape (nframes, natoms, 3). + virial + The virial of the system, in shape (nframes, 9). + atomic_energy + The atomic energy of the system, in shape (nframes, natoms). Only returned + when atomic is True. + atomic_virial + The atomic virial of the system, in shape (nframes, natoms, 9). Only returned + when atomic is True. + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + results = self.deep_eval.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + # energy = results["dos_redu"].reshape(nframes, self.get_numb_dos()) + atomic_energy = results["dos"].reshape(nframes, natoms, self.get_numb_dos()) + # not same as dos_redu... why? + energy = np.sum(atomic_energy, axis=1) + + if atomic: + return ( + energy, + atomic_energy, + ) + else: + return (energy,) + + def get_numb_dos(self) -> int: + return self.deep_eval.get_numb_dos() + + +__all__ = ["DeepDOS"] diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py new file mode 100644 index 0000000000..92b449eb36 --- /dev/null +++ b/deepmd/infer/deep_eval.py @@ -0,0 +1,467 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, + abstractmethod, +) +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Optional, + Tuple, + Union, +) + +import numpy as np + +from deepmd.model_format.output_def import ( + FittingOutputDef, + ModelOutputDef, +) +from deepmd.utils.batch_size import ( + AutoBatchSize, +) + +from .backend import ( + DPBackend, + detect_backend, +) + +if TYPE_CHECKING: + import ase.neighborlist + + +class DeepEvalBase(ABC): + """Low-level Deep Evaluator interface. + + Backends should inherbit implement this interface. High-level interface + will be built on top of this. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + + @abstractmethod + def __init__( + self, + model_file: str, + output_def: ModelOutputDef, + *args: List[Any], + auto_batch_size: Union[bool, int, AutoBatchSize] = True, + neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, + **kwargs: Dict[str, Any], + ) -> None: + pass + + def __new__(cls, model_file: str, *args, **kwargs): + if cls is DeepEvalBase: + backend = detect_backend(model_file) + if backend == DPBackend.TensorFlow: + from deepmd.tf.infer.deep_eval import DeepEval as DeepEvalTF + + return super().__new__(DeepEvalTF) + elif backend == DPBackend.PyTorch: + from deepmd.pt.infer.deep_eval import DeepEval as DeepEvalPT + + return super().__new__(DeepEvalPT) + else: + raise NotImplementedError("Unsupported backend: " + str(backend)) + return super().__new__(cls) + + @abstractmethod + def eval( + self, + coords: np.ndarray, + cells: np.ndarray, + atom_types: List[int], + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + **kwargs: Dict[str, Any], + ) -> Dict[str, np.ndarray]: + """Evaluate the energy, force and virial by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Calculate the atomic energy and virial + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + **kwargs + Other parameters + + Returns + ------- + output_dict : dict + The output of the evaluation. The keys are the names of the output + variables, and the values are the corresponding output arrays. + """ + + @abstractmethod + def get_rcut(self) -> float: + """Get the cutoff radius of this model.""" + + @abstractmethod + def get_ntypes(self) -> int: + """Get the number of atom types of this model.""" + + @abstractmethod + def get_type_map(self) -> List[str]: + """Get the type map (element name of the atom types) of this model.""" + + @abstractmethod + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this DP.""" + + @abstractmethod + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this DP.""" + + def eval_descriptor( + self, + coords: np.ndarray, + cells: np.ndarray, + atom_types: List[int], + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + efield: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: Dict[str, Any], + ) -> np.ndarray: + """Evaluate descriptors by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + efield + The external field on atoms. + The array should be of size nframes x natoms x 3 + mixed_type + Whether to perform the mixed_type mode. + If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), + in which frames in a system may have different natoms_vec(s), with the same nloc. + + Returns + ------- + descriptor + Descriptors. + """ + raise NotImplementedError + + def eval_typeebd(self) -> np.ndarray: + """Evaluate output of type embedding network by using this model. + + Returns + ------- + np.ndarray + The output of type embedding network. The shape is [ntypes, o_size], + where ntypes is the number of types, and o_size is the number of nodes + in the output layer. + + Raises + ------ + KeyError + If the model does not enable type embedding. + """ + raise NotImplementedError + + def _check_distinguished_types(self, atom_types: np.ndarray) -> bool: + """Check if atom types of each frame.""" + return np.all(np.equal(atom_types, atom_types[0])) + + @property + @abstractmethod + def model_type(self) -> "DeepEval": + """The the evaluator of the model type.""" + + @abstractmethod + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model.""" + + def get_numb_dos(self) -> int: + """Get the number of DOS.""" + raise NotImplementedError + + def get_has_efield(self): + """Check if the model has efield.""" + return False + + @abstractmethod + def get_ntypes_spin(self): + """Get the number of spin atom types of this model.""" + + +class DeepEval(ABC): + """High-level Deep Evaluator interface. + + The specific DeepEval, such as DeepPot and DeepTensor, should inherit + from this class. This class provides a high-level interface on the top + of the low-level interface. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + + def __new__(cls, model_file: str, *args, **kwargs): + if cls is DeepEval: + deep_eval = DeepEvalBase( + model_file, + ModelOutputDef(FittingOutputDef([])), + *args, + **kwargs, + ) + return super().__new__(deep_eval.model_type) + return super().__new__(cls) + + def __init__( + self, + model_file: str, + *args: List[Any], + auto_batch_size: Union[bool, int, AutoBatchSize] = True, + neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, + **kwargs: Dict[str, Any], + ) -> None: + self.deep_eval = DeepEvalBase( + model_file, + self.output_def, + *args, + auto_batch_size=auto_batch_size, + neighbor_list=neighbor_list, + **kwargs, + ) + + @property + @abstractmethod + def output_def(self) -> ModelOutputDef: + """Returns the output variable definitions.""" + + def get_rcut(self) -> float: + """Get the cutoff radius of this model.""" + return self.deep_eval.get_rcut() + + def get_ntypes(self) -> int: + """Get the number of atom types of this model.""" + return self.deep_eval.get_ntypes() + + def get_type_map(self) -> List[str]: + """Get the type map (element name of the atom types) of this model.""" + return self.deep_eval.get_type_map() + + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this DP.""" + return self.deep_eval.get_dim_fparam() + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this DP.""" + return self.deep_eval.get_dim_aparam() + + def _get_natoms_and_nframes( + self, + coords: np.ndarray, + atom_types: Union[List[int], np.ndarray], + mixed_type: bool = False, + ) -> Tuple[int, int]: + if mixed_type: + natoms = len(atom_types[0]) + else: + natoms = len(atom_types) + if natoms == 0: + assert coords.size == 0 + else: + coords = np.reshape(np.array(coords), [-1, natoms * 3]) + nframes = coords.shape[0] + return natoms, nframes + + def _expande_atype(self, atype: np.ndarray, nframes: int, mixed_type: bool): + if not mixed_type: + atype = np.tile(atype.reshape(1, -1), (nframes, 1)) + return atype + + def eval_descriptor( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: np.ndarray, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: Dict[str, Any], + ) -> np.ndarray: + """Evaluate descriptors by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + efield + The external field on atoms. + The array should be of size nframes x natoms x 3 + mixed_type + Whether to perform the mixed_type mode. + If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), + in which frames in a system may have different natoms_vec(s), with the same nloc. + + Returns + ------- + descriptor + Descriptors. + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + descriptor = self.deep_eval.eval_descriptor( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + return descriptor + + def eval_typeebd(self) -> np.ndarray: + """Evaluate output of type embedding network by using this model. + + Returns + ------- + np.ndarray + The output of type embedding network. The shape is [ntypes, o_size], + where ntypes is the number of types, and o_size is the number of nodes + in the output layer. + + Raises + ------ + KeyError + If the model does not enable type embedding. + + See Also + -------- + deepmd.tf.utils.type_embed.TypeEmbedNet : The type embedding network. + + Examples + -------- + Get the output of type embedding network of `graph.pb`: + + >>> from deepmd.infer import DeepPotential + >>> dp = DeepPotential('graph.pb') + >>> dp.eval_typeebd() + """ + return self.deep_eval.eval_typeebd() + + def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type): + coords = np.array(coords) + if cells is not None: + cells = np.array(cells) + atom_types = np.array(atom_types) + if fparam is not None: + fparam = np.array(fparam) + if aparam is not None: + aparam = np.array(aparam) + natoms, nframes = self._get_natoms_and_nframes(coords, atom_types, mixed_type) + atom_types = self._expande_atype(atom_types, nframes, mixed_type) + return coords, cells, atom_types, fparam, aparam, nframes, natoms + + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model.""" + return self.deep_eval.get_sel_type() + + def _get_sel_natoms(self, atype) -> int: + return np.sum(np.isin(atype, self.get_sel_type()).astype(int)) + + @property + def has_efield(self): + """Check if the model has efield.""" + return self.deep_eval.get_has_efield() + + def get_ntypes_spin(self): + """Get the number of spin atom types of this model.""" + return self.deep_eval.get_ntypes_spin() diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py new file mode 100644 index 0000000000..96ad70dcc0 --- /dev/null +++ b/deepmd/infer/deep_polar.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Optional, +) + +import numpy as np + +from deepmd.infer.deep_tensor import ( + DeepTensor, +) + + +class DeepPolar(DeepTensor): + @property + def output_tensor_name(self) -> str: + return "polar" + + +class DeepGlobalPolar(DeepTensor): + @property + def output_tensor_name(self) -> str: + return "global_polar" + + def eval( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: np.ndarray, + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: dict, + ) -> np.ndarray: + """Evaluate the model. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + If True (default), return the atomic tensor + Otherwise return the global tensor + fparam + Not used in this model + aparam + Not used in this model + mixed_type + Whether to perform the mixed_type mode. + If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), + in which frames in a system may have different natoms_vec(s), with the same nloc. + + Returns + ------- + tensor + The returned tensor + If atomic == False then of size nframes x output_dim + else of size nframes x natoms x output_dim + """ + return super().eval( + coords, + cells, + atom_types, + atomic=atomic, + fparam=fparam, + aparam=aparam, + mixed_type=mixed_type, + **kwargs, + ) diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 546c0f3c7e..3f207d03f8 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -1,28 +1,25 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from abc import ( - ABC, - abstractmethod, -) from typing import ( - List, + Any, + Dict, Optional, Tuple, - Union, ) import numpy as np -from deepmd.utils.batch_size import ( - AutoBatchSize, +from deepmd.model_format.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, ) -from .backend import ( - DPBackend, - detect_backend, +from .deep_eval import ( + DeepEval, ) -class DeepPot(ABC): +class DeepPot(DeepEval): """Potential energy model. Parameters @@ -35,45 +32,52 @@ class DeepPot(ABC): neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model. - """ - - @abstractmethod - def __init__( - self, - model_file, - *args, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, - neighbor_list=None, - **kwargs, - ) -> None: - pass - - def __new__(cls, model_file: str, *args, **kwargs): - if cls is DeepPot: - backend = detect_backend(model_file) - if backend == DPBackend.TensorFlow: - from deepmd.tf.infer.deep_pot import DeepPot as DeepPotTF - return super().__new__(DeepPotTF) - elif backend == DPBackend.PyTorch: - from deepmd.pt.infer.deep_eval import DeepPot as DeepPotPT + Examples + -------- + >>> from deepmd.infer import DeepPot + >>> import numpy as np + >>> dp = DeepPot('graph.pb') + >>> coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1]) + >>> cell = np.diag(10 * np.ones(3)).reshape([1, -1]) + >>> atype = [1,0,1] + >>> e, f, v = dp.eval(coord, cell, atype) + + where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively. + """ - return super().__new__(DeepPotPT) - else: - raise NotImplementedError("Unsupported backend: " + str(backend)) - return super().__new__(cls) + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + "energy", + shape=[1], + reduciable=True, + differentiable=True, + atomic=True, + ), + OutputVariableDef( + # ugly... + "energy_derv_c_redu", + shape=[1, 3, 3], + ), + ] + ) + ) - @abstractmethod def eval( self, coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], + cells: Optional[np.ndarray], + atom_types: np.ndarray, atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, mixed_type: bool = False, + **kwargs: Dict[str, Any], ) -> Tuple[np.ndarray, ...]: """Evaluate energy, force, and virial. If atomic is True, also return atomic energy and atomic virial. @@ -94,10 +98,10 @@ def eval( The frame parameters, by default None. aparam : np.ndarray, optional The atomic parameters, by default None. - efield : np.ndarray, optional - The electric field, by default None. mixed_type : bool, optional - Whether the system contains mixed atom types, by default False. + Whether the atom_types is mixed type, by default False. + **kwargs : Dict[str, Any] + Keyword arguments. Returns ------- @@ -121,22 +125,44 @@ def eval( # finetune: +mixed_type # dpdata # ase - - @abstractmethod - def get_ntypes(self) -> int: - """Get the number of atom types of this model.""" - - @abstractmethod - def get_type_map(self) -> List[str]: - """Get the type map (element name of the atom types) of this model.""" - - @abstractmethod - def get_dim_fparam(self) -> int: - """Get the number (dimension) of frame parameters of this DP.""" - - @abstractmethod - def get_dim_aparam(self) -> int: - """Get the number (dimension) of atomic parameters of this DP.""" + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + results = self.deep_eval.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + energy = results["energy_redu"].reshape(nframes, 1) + force = results["energy_derv_r"].reshape(nframes, natoms, 3) + virial = results["energy_derv_c_redu"].reshape(nframes, 9) + + if atomic: + atomic_energy = results["energy"].reshape(nframes, natoms, 1) + atomic_virial = results["energy_derv_c"].reshape(nframes, natoms, 9) + return ( + energy, + force, + virial, + atomic_energy, + atomic_virial, + ) + else: + return ( + energy, + force, + virial, + ) __all__ = ["DeepPot"] diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py new file mode 100644 index 0000000000..64c2fc1130 --- /dev/null +++ b/deepmd/infer/deep_tensor.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + abstractmethod, +) +from typing import ( + Optional, + Tuple, +) + +import numpy as np + +from deepmd.infer.deep_eval import ( + DeepEval, +) +from deepmd.model_format.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) + + +class DeepTensor(DeepEval): + def eval( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: np.ndarray, + atomic: bool = True, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: dict, + ) -> np.ndarray: + """Evaluate the model. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + If True (default), return the atomic tensor + Otherwise return the global tensor + fparam + Not used in this model + aparam + Not used in this model + efield + Not used in this model + mixed_type + Whether to perform the mixed_type mode. + If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), + in which frames in a system may have different natoms_vec(s), with the same nloc. + + Returns + ------- + tensor + The returned tensor + If atomic == False then of size nframes x output_dim + else of size nframes x natoms x output_dim + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + results = self.deep_eval.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + sel_natoms = self._get_sel_natoms(atom_types[0]) + if atomic: + return results[self.output_tensor_name].reshape(nframes, sel_natoms, -1) + else: + return results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1) + + def eval_full( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: np.ndarray, + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: dict, + ) -> Tuple[np.ndarray, ...]: + """Evaluate the model with interface similar to the energy model. + Will return global tensor, component-wise force and virial + and optionally atomic tensor and atomic virial. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Whether to calculate atomic tensor and virial + fparam + Not used in this model + aparam + Not used in this model + mixed_type + Whether to perform the mixed_type mode. + If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), + in which frames in a system may have different natoms_vec(s), with the same nloc. + + Returns + ------- + tensor + The global tensor. + shape: [nframes x nout] + force + The component-wise force (negative derivative) on each atom. + shape: [nframes x nout x natoms x 3] + virial + The component-wise virial of the tensor. + shape: [nframes x nout x 9] + atom_tensor + The atomic tensor. Only returned when atomic == True + shape: [nframes x natoms x nout] + atom_virial + The atomic virial. Only returned when atomic == True + shape: [nframes x nout x natoms x 9] + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + results = self.deep_eval.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + sel_natoms = self._get_sel_natoms(atom_types[0]) + energy = results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1) + force = results[f"{self.output_tensor_name}_derv_r"].reshape( + nframes, -1, natoms, 3 + ) + virial = results[f"{self.output_tensor_name}_derv_c_redu"].reshape( + nframes, -1, 9 + ) + atomic_energy = results[self.output_tensor_name].reshape( + nframes, sel_natoms, -1 + ) + atomic_virial = results[f"{self.output_tensor_name}_derv_c"].reshape( + nframes, -1, natoms, 9 + ) + + if atomic: + return ( + energy, + force, + virial, + atomic_energy, + atomic_virial, + ) + else: + return ( + energy, + force, + virial, + ) + + @property + @abstractmethod + def output_tensor_name(self) -> str: + """The name of the tensor.""" + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + self.output_tensor_name, + shape=[-1], + reduciable=True, + differentiable=True, + atomic=True, + ), + OutputVariableDef( + # ugly... + f"{self.output_tensor_name}_derv_c_redu", + shape=[-1, 3, 3], + ), + ] + ) + ) diff --git a/deepmd/infer/deep_wfc.py b/deepmd/infer/deep_wfc.py new file mode 100644 index 0000000000..d1c7a92ad2 --- /dev/null +++ b/deepmd/infer/deep_wfc.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.infer.deep_tensor import ( + DeepTensor, +) + + +class DeepWFC(DeepTensor): + @property + def output_tensor_name(self) -> str: + return "wfc" diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index b5d596ed0f..31f61f68ca 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from pathlib import ( - Path, -) from typing import ( + TYPE_CHECKING, + Any, Callable, + ClassVar, + Dict, List, Optional, Tuple, @@ -13,7 +14,15 @@ import numpy as np import torch -from deepmd.infer.deep_pot import DeepPot as DeepPotBase +from deepmd.infer.deep_eval import ( + DeepEvalBase, +) +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.model_format.output_def import ( + ModelOutputDef, +) from deepmd.pt.model.model import ( get_model, ) @@ -31,13 +40,21 @@ GLOBAL_PT_FLOAT_PRECISION, ) +if TYPE_CHECKING: + import ase.neighborlist + -class DeepEval: +class DeepEval(DeepEvalBase): def __init__( self, - model_file: "Path", + model_file: str, + output_def: ModelOutputDef, + *args: List[Any], auto_batch_size: Union[bool, int, AutoBatchSize] = True, + neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, + **kwargs: Dict[str, Any], ): + self.output_def = output_def self.model_path = model_file state_dict = torch.load(model_file, map_location=env.DEVICE) if "model" in state_dict: @@ -64,29 +81,54 @@ def __init__( else: raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") - def eval( - self, - coords: Union[np.ndarray, torch.Tensor], - cells: Optional[Union[np.ndarray, torch.Tensor]], - atom_types: Union[np.ndarray, torch.Tensor, List[int]], - atomic: bool = False, - ): - raise NotImplementedError + def get_rcut(self) -> float: + """Get the cutoff radius of this model.""" + return self.rcut + def get_ntypes(self) -> int: + """Get the number of atom types of this model.""" + return len(self.type_map) -class DeepPot(DeepEval, DeepPotBase): - def __init__( - self, - model_file: "Path", - auto_batch_size: Union[bool, int, AutoBatchSize] = True, - neighbor_list=None, - ): - if neighbor_list is not None: - raise NotImplementedError - super().__init__( - model_file, - auto_batch_size=auto_batch_size, - ) + def get_type_map(self) -> List[str]: + """Get the type map (element name of the atom types) of this model.""" + return self.type_map + + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this DP.""" + return 0 + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this DP.""" + return 0 + + @property + def model_type(self) -> "DeepEval": + """The the evaluator of the model type.""" + return DeepPot + + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model.""" + return [] + + def get_numb_dos(self) -> int: + """Get the number of DOS.""" + raise 0 + + def get_has_efield(self): + """Check if the model has efield.""" + return False + + def get_ntypes_spin(self): + """Get the number of spin atom types of this model.""" + return 0 + + _OUTDEF_DP2PT: ClassVar[dict] = { + "energy": "atom_energy", + "energy_redu": "energy", + "energy_derv_r": "force", + "energy_derv_c": "atom_virial", + "energy_derv_c_redu": "virial", + } def eval( self, @@ -96,10 +138,45 @@ def eval( atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - mixed_type: bool = False, - ): - if fparam is not None or aparam is not None or efield is not None: + **kwargs: Dict[str, Any], + ) -> Dict[str, np.ndarray]: + """Evaluate the energy, force and virial by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Calculate the atomic energy and virial + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + **kwargs + Other parameters + + Returns + ------- + output_dict : dict + The output of the evaluation. The keys are the names of the output + variables, and the values are the corresponding output arrays. + """ + if fparam is not None or aparam is not None: raise NotImplementedError # convert all of the input to numpy array atom_types = np.array(atom_types, dtype=np.int32) @@ -109,9 +186,19 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) - return self._eval_func(self._eval_model, numb_test, natoms)( + out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, atomic ) + return dict( + zip( + [ + x.name + for x in self.output_def.var_defs.values() + if atomic or "_redu" in x.name or "_derv_r" in x.name + ], + out, + ) + ) def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: """Wrapper method with auto batch size. @@ -195,54 +282,33 @@ def _eval_model( ) if isinstance(batch_output, tuple): batch_output = batch_output[0] - energy_out = batch_output["energy"].reshape(nframes, 1).detach().cpu().numpy() - if "atom_energy" in batch_output: - atomic_energy_out = ( - batch_output["atom_energy"] - .reshape(nframes, natoms, 1) - .detach() - .cpu() - .numpy() - ) - force_out = ( - batch_output["force"].reshape(nframes, natoms, 3).detach().cpu().numpy() - ) - virial_out = batch_output["virial"].reshape(nframes, 9).detach().cpu().numpy() - if "atomic_virial" in batch_output: - atomic_virial_out = ( - batch_output["atomic_virial"] - .reshape(nframes, natoms, 9) - .detach() - .cpu() - .numpy() - ) - if not atomic: - return energy_out, force_out, virial_out - else: - return ( - energy_out, - force_out, - virial_out, - atomic_energy_out, - atomic_virial_out, - ) - - def get_ntypes(self) -> int: - """Get the number of atom types of this model.""" - return len(self.type_map) + results = [] + for ii, odef in enumerate(self.output_def.var_defs.values()): + if not atomic and "_redu" not in odef.name and "_derv_r" not in odef.name: + continue + pt_name = self._OUTDEF_DP2PT[odef.name] + if pt_name in batch_output: + shape = self._get_output_shape(odef.name, nframes, natoms, odef.shape) + out = batch_output[pt_name].reshape(shape).detach().cpu().numpy() + results.append(out) + return tuple(results) - def get_type_map(self) -> List[str]: - """Get the type map (element name of the atom types) of this model.""" - return self.type_map - - def get_dim_fparam(self) -> int: - """Get the number (dimension) of frame parameters of this DP.""" - return 0 - - def get_dim_aparam(self) -> int: - """Get the number (dimension) of atomic parameters of this DP.""" - return 0 + def _get_output_shape(self, name, nframes, natoms, shape): + if "_redu" in name: + if "_derv_c" in name: + return [nframes, *shape[:-2], 9] + else: + return [nframes, *shape, 1] + else: + if "_derv_c" in name: + return [nframes, *shape[:-2], natoms, 9] + elif "_derv_r" in name: + return [nframes, *shape[:-1], natoms, 3] + else: + # Something wrong here? + # return [nframes, *shape, natoms, 1] + return [nframes, natoms, *shape, 1] # For tests only diff --git a/deepmd/tf/entrypoints/test.py b/deepmd/tf/entrypoints/test.py index 1b917bc276..e1a2e3d46b 100644 --- a/deepmd/tf/entrypoints/test.py +++ b/deepmd/tf/entrypoints/test.py @@ -14,6 +14,22 @@ import numpy as np +from deepmd.infer.deep_dipole import ( + DeepDipole, +) +from deepmd.infer.deep_dos import ( + DeepDOS, +) +from deepmd.infer.deep_polar import ( + DeepGlobalPolar, + DeepPolar, +) +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.infer.deep_wfc import ( + DeepWFC, +) from deepmd.tf import ( DeepPotential, ) @@ -115,7 +131,7 @@ def test( log.info(f"# testing system : {system}") # create data class - tmap = dp.get_type_map() if dp.model_type == "ener" else None + tmap = dp.get_type_map() if isinstance(dp, DeepPot) else None data = DeepmdData( system, set_prefix, @@ -124,7 +140,7 @@ def test( sort_atoms=False, ) - if dp.model_type == "ener": + if isinstance(dp, DeepPot): err = test_ener( dp, data, @@ -134,7 +150,7 @@ def test( atomic, append_detail=(cc != 0), ) - elif dp.model_type == "dos": + elif isinstance(dp, DeepDOS): err = test_dos( dp, data, @@ -144,11 +160,11 @@ def test( atomic, append_detail=(cc != 0), ) - elif dp.model_type == "dipole": + elif isinstance(dp, DeepDipole): err = test_dipole(dp, data, numb_test, detail_file, atomic) - elif dp.model_type == "polar": + elif isinstance(dp, DeepPolar): err = test_polar(dp, data, numb_test, detail_file, atomic=atomic) - elif dp.model_type == "global_polar": # should not appear in this new version + elif isinstance(dp, DeepGlobalPolar): # should not appear in this new version log.warning( "Global polar model is not currently supported. Please directly use the polar mode and change loss parameters." ) @@ -166,17 +182,17 @@ def test( if len(all_sys) > 1: log.info("# ----------weighted average of errors----------- ") log.info(f"# number of systems : {len(all_sys)}") - if dp.model_type == "ener": + if dp == "ener": print_ener_sys_avg(avg_err) - elif dp.model_type == "dos": + elif dp == "dos": print_dos_sys_avg(avg_err) - elif dp.model_type == "dipole": + elif dp == "dipole": print_dipole_sys_avg(avg_err) - elif dp.model_type == "polar": + elif dp == "polar": print_polar_sys_avg(avg_err) - elif dp.model_type == "global_polar": + elif dp == "global_polar": print_polar_sys_avg(avg_err) - elif dp.model_type == "wfc": + elif dp == "wfc": print_wfc_sys_avg(avg_err) log.info("# ----------------------------------------------- ") diff --git a/deepmd/tf/infer/__init__.py b/deepmd/tf/infer/__init__.py index c1071af35c..f0cd86a885 100644 --- a/deepmd/tf/infer/__init__.py +++ b/deepmd/tf/infer/__init__.py @@ -1,12 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Submodule containing all the implemented potentials.""" -from pathlib import ( - Path, -) -from typing import ( - Optional, - Union, +from deepmd import ( + DeepPotential, ) from .data_modifier import ( @@ -51,96 +47,3 @@ "EwaldRecp", "calc_model_devi", ] - - -def DeepPotential( - model_file: Union[str, Path], - load_prefix: str = "load", - default_tf_graph: bool = False, - input_map: Optional[dict] = None, - neighbor_list=None, -) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepDOS, DeepWFC]: - """Factory function that will inialize appropriate potential read from `model_file`. - - Parameters - ---------- - model_file : str - The name of the frozen model file. - load_prefix : str - The prefix in the load computational graph - default_tf_graph : bool - If uses the default tf graph, otherwise build a new tf graph for evaluation - input_map : dict, optional - The input map for tf.import_graph_def. Only work with default tf graph - neighbor_list : ase.neighborlist.NeighborList, optional - The neighbor list object. If None, then build the native neighbor list. - - Returns - ------- - Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC] - one of the available potentials - - Raises - ------ - RuntimeError - if model file does not correspond to any implementd potential - """ - mf = Path(model_file) - - model_type = DeepEval( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - ).model_type - - if model_type == "ener": - dp = DeepPot( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - neighbor_list=neighbor_list, - ) - elif model_type == "dos": - dp = DeepDOS( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - ) - elif model_type == "dipole": - dp = DeepDipole( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - neighbor_list=neighbor_list, - ) - elif model_type == "polar": - dp = DeepPolar( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - neighbor_list=neighbor_list, - ) - elif model_type == "global_polar": - dp = DeepGlobalPolar( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - neighbor_list=neighbor_list, - ) - elif model_type == "wfc": - dp = DeepWFC( - mf, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - ) - else: - raise RuntimeError(f"unknown model type {model_type}") - - return dp diff --git a/deepmd/tf/infer/data_modifier.py b/deepmd/tf/infer/data_modifier.py index e53151d80a..ccd072673d 100644 --- a/deepmd/tf/infer/data_modifier.py +++ b/deepmd/tf/infer/data_modifier.py @@ -17,9 +17,7 @@ op_module, tf, ) -from deepmd.tf.infer.deep_dipole import ( - DeepDipole, -) +from deepmd.tf.infer.deep_dipole import DeepDipoleOld as DeepDipole from deepmd.tf.infer.ewald_recp import ( EwaldRecp, ) diff --git a/deepmd/tf/infer/deep_dipole.py b/deepmd/tf/infer/deep_dipole.py index a0a4a0f78e..e10d09564d 100644 --- a/deepmd/tf/infer/deep_dipole.py +++ b/deepmd/tf/infer/deep_dipole.py @@ -1,20 +1,25 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from pathlib import ( + Path, +) from typing import ( - TYPE_CHECKING, Optional, ) +from deepmd.infer.deep_dipole import ( + DeepDipole, +) from deepmd.tf.infer.deep_tensor import ( DeepTensor, ) -if TYPE_CHECKING: - from pathlib import ( - Path, - ) +__all__ = [ + "DeepDipole", +] -class DeepDipole(DeepTensor): +class DeepDipoleOld(DeepTensor): + # used for DipoleChargeModifier only """Constructor. Parameters diff --git a/deepmd/tf/infer/deep_dos.py b/deepmd/tf/infer/deep_dos.py index 9b2830161d..7a9f9b781c 100644 --- a/deepmd/tf/infer/deep_dos.py +++ b/deepmd/tf/infer/deep_dos.py @@ -1,506 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import logging -from typing import ( - TYPE_CHECKING, - Callable, - List, - Optional, - Tuple, - Union, +from deepmd.infer.deep_dos import ( + DeepDOS, ) -import numpy as np - -from deepmd.tf.common import ( - make_default_mesh, -) -from deepmd.tf.infer.deep_eval import ( - DeepEval, -) -from deepmd.tf.utils.batch_size import ( - AutoBatchSize, -) -from deepmd.tf.utils.sess import ( - run_sess, -) - -if TYPE_CHECKING: - from pathlib import ( - Path, - ) - -log = logging.getLogger(__name__) - - -class DeepDOS(DeepEval): - """Constructor. - - Parameters - ---------- - model_file : Path - The name of the frozen model file. - load_prefix: str - The prefix in the load computational graph - default_tf_graph : bool - If uses the default tf graph, otherwise build a new tf graph for evaluation - auto_batch_size : bool or int or AutomaticBatchSize, default: True - If True, automatic batch size will be used. If int, it will be used - as the initial batch size. - input_map : dict, optional - The input map for tf.import_graph_def. Only work with default tf graph - - Warnings - -------- - For developers: `DeepTensor` initializer must be called at the end after - `self.tensors` are modified because it uses the data in `self.tensors` dict. - Do not chanage the order! - """ - - def __init__( - self, - model_file: "Path", - load_prefix: str = "load", - default_tf_graph: bool = False, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, - input_map: Optional[dict] = None, - ) -> None: - # add these tensors on top of what is defined by DeepTensor Class - # use this in favor of dict update to move attribute from class to - # instance namespace - self.tensors = { - # descrpt attrs - "t_ntypes": "descrpt_attr/ntypes:0", - "t_rcut": "descrpt_attr/rcut:0", - # fitting attrs - "t_dfparam": "fitting_attr/dfparam:0", - "t_daparam": "fitting_attr/daparam:0", - "t_numb_dos": "fitting_attr/numb_dos:0", - # model attrs - "t_tmap": "model_attr/tmap:0", - # inputs - "t_coord": "t_coord:0", - "t_type": "t_type:0", - "t_natoms": "t_natoms:0", - "t_box": "t_box:0", - "t_mesh": "t_mesh:0", - # add output tensors - "t_dos": "o_dos:0", - "t_atom_dos": "o_atom_dos:0", - "t_descriptor": "o_descriptor:0", - } - DeepEval.__init__( - self, - model_file, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - auto_batch_size=auto_batch_size, - input_map=input_map, - ) - - # load optional tensors - operations = [op.name for op in self.graph.get_operations()] - # check if the graph has these operations: - # if yes add them - if "load/t_fparam" in operations: - self.tensors.update({"t_fparam": "t_fparam:0"}) - self.has_fparam = True - else: - log.debug("Could not get tensor 't_fparam:0'") - self.t_fparam = None - self.has_fparam = False - - if "load/t_aparam" in operations: - self.tensors.update({"t_aparam": "t_aparam:0"}) - self.has_aparam = True - else: - log.debug("Could not get tensor 't_aparam:0'") - self.t_aparam = None - self.has_aparam = False - - # now load tensors to object attributes - for attr_name, tensor_name in self.tensors.items(): - try: - self._get_tensor(tensor_name, attr_name) - except KeyError: - if attr_name != "t_descriptor": - raise - - self._run_default_sess() - self.tmap = self.tmap.decode("UTF-8").split() - - # setup modifier - try: - t_modifier_type = self._get_tensor("modifier_attr/type:0") - self.modifier_type = run_sess(self.sess, t_modifier_type).decode("UTF-8") - except (ValueError, KeyError): - self.modifier_type = None - - def _run_default_sess(self): - [ - self.ntypes, - self.rcut, - self.numb_dos, - self.dfparam, - self.daparam, - self.tmap, - ] = run_sess( - self.sess, - [ - self.t_ntypes, - self.t_rcut, - self.t_numb_dos, - self.t_dfparam, - self.t_daparam, - self.t_tmap, - ], - ) - - def get_ntypes(self) -> int: - """Get the number of atom types of this model.""" - return self.ntypes - - def get_rcut(self) -> float: - """Get the cut-off radius of this model.""" - return self.rcut - - def get_numb_dos(self) -> int: - """Get the length of DOS output of this DP model.""" - return self.numb_dos - - def get_type_map(self) -> List[str]: - """Get the type map (element name of the atom types) of this model.""" - return self.tmap - - def get_sel_type(self) -> List[int]: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") - - def get_dim_fparam(self) -> int: - """Get the number (dimension) of frame parameters of this DP.""" - return self.dfparam - - def get_dim_aparam(self) -> int: - """Get the number (dimension) of atomic parameters of this DP.""" - return self.daparam - - def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: - """Wrapper method with auto batch size. - - Parameters - ---------- - inner_func : Callable - the method to be wrapped - numb_test : int - number of tests - natoms : int - number of atoms - - Returns - ------- - Callable - the wrapper - """ - if self.auto_batch_size is not None: - - def eval_func(*args, **kwargs): - return self.auto_batch_size.execute_all( - inner_func, numb_test, natoms, *args, **kwargs - ) - - else: - eval_func = inner_func - return eval_func - - def _get_natoms_and_nframes( - self, - coords: np.ndarray, - atom_types: Union[List[int], np.ndarray], - mixed_type: bool = False, - ) -> Tuple[int, int]: - if mixed_type: - natoms = len(atom_types[0]) - else: - natoms = len(atom_types) - coords = np.reshape(np.array(coords), [-1, natoms * 3]) - nframes = coords.shape[0] - return natoms, nframes - - def eval( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - mixed_type: bool = False, - ) -> Tuple[np.ndarray, ...]: - """Evaluate the dos, atom_dos by using this model. - - Parameters - ---------- - coords - The coordinates of atoms. - The array should be of size nframes x natoms x 3 - cells - The cell of the region. - If None then non-PBC is assumed, otherwise using PBC. - The array should be of size nframes x 9 - atom_types - The atom types - The list should contain natoms ints - atomic - Calculate the atomic energy and virial - fparam - The frame parameter. - The array can be of size : - - nframes x dim_fparam. - - dim_fparam. Then all frames are assumed to be provided with the same fparam. - aparam - The atomic parameter - The array can be of size : - - nframes x natoms x dim_aparam. - - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - - dim_aparam. Then all frames and atoms are provided with the same aparam. - mixed_type - Whether to perform the mixed_type mode. - If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), - in which frames in a system may have different natoms_vec(s), with the same nloc. - - Returns - ------- - dos - The electron density of state. - atom_dos - The atom-sited density of state. Only returned when atomic == True - """ - # reshape coords before getting shape - natoms, numb_test = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - output = self._eval_func(self._eval_inner, numb_test, natoms)( - coords, - cells, - atom_types, - fparam=fparam, - aparam=aparam, - atomic=atomic, - mixed_type=mixed_type, - ) - - return output - - def _prepare_feed_dict( - self, - coords, - cells, - atom_types, - fparam=None, - aparam=None, - atomic=False, - mixed_type=False, - ): - # standarize the shape of inputs - natoms, nframes = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - if mixed_type: - atom_types = np.array(atom_types, dtype=int).reshape([-1, natoms]) - else: - atom_types = np.array(atom_types, dtype=int).reshape([-1]) - coords = np.reshape(np.array(coords), [-1, natoms * 3]) - if cells is None: - pbc = False - # make cells to work around the requirement of pbc - cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9]) - else: - pbc = True - cells = np.array(cells).reshape([nframes, 9]) - - if self.has_fparam: - assert fparam is not None - fparam = np.array(fparam) - if self.has_aparam: - assert aparam is not None - aparam = np.array(aparam) - - # reshape the inputs - if self.has_fparam: - fdim = self.get_dim_fparam() - if fparam.size == nframes * fdim: - fparam = np.reshape(fparam, [nframes, fdim]) - elif fparam.size == fdim: - fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) - else: - raise RuntimeError( - "got wrong size of frame param, should be either %d x %d or %d" - % (nframes, fdim, fdim) - ) - if self.has_aparam: - fdim = self.get_dim_aparam() - if aparam.size == nframes * natoms * fdim: - aparam = np.reshape(aparam, [nframes, natoms * fdim]) - elif aparam.size == natoms * fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) - elif aparam.size == fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) - else: - raise RuntimeError( - "got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d" - % (nframes, natoms, fdim, natoms, fdim, fdim) - ) - - # sort inputs - coords, atom_types, imap = self.sort_input( - coords, atom_types, mixed_type=mixed_type - ) - - # make natoms_vec and default_mesh - natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) - assert natoms_vec[0] == natoms - - # evaluate - feed_dict_test = {} - feed_dict_test[self.t_natoms] = natoms_vec - if mixed_type: - feed_dict_test[self.t_type] = atom_types.reshape([-1]) - else: - feed_dict_test[self.t_type] = np.tile(atom_types, [nframes, 1]).reshape( - [-1] - ) - feed_dict_test[self.t_coord] = np.reshape(coords, [-1]) - - if len(self.t_box.shape) == 1: - feed_dict_test[self.t_box] = np.reshape(cells, [-1]) - elif len(self.t_box.shape) == 2: - feed_dict_test[self.t_box] = cells - else: - raise RuntimeError - feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type) - if self.has_fparam: - feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1]) - if self.has_aparam: - feed_dict_test[self.t_aparam] = np.reshape(aparam, [-1]) - return feed_dict_test, imap - - def _eval_inner( - self, - coords, - cells, - atom_types, - fparam=None, - aparam=None, - atomic=False, - mixed_type=False, - ): - natoms, nframes = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - feed_dict_test, imap = self._prepare_feed_dict( - coords, cells, atom_types, fparam, aparam, mixed_type=mixed_type - ) - t_out = [self.t_dos] - if atomic: - t_out += [self.t_atom_dos] - - v_out = run_sess(self.sess, t_out, feed_dict=feed_dict_test) - dos = v_out[0] - if atomic: - atom_dos = v_out[1] - - # reverse map of the outputs - if atomic: - atom_dos = self.reverse_map( - np.reshape(atom_dos, [nframes, -1, self.numb_dos]), imap - ) - dos = np.sum(atom_dos, axis=1) - - dos = np.reshape(dos, [nframes, self.numb_dos]) - if atomic: - atom_dos = np.reshape(atom_dos, [nframes, natoms, self.numb_dos]) - return dos, atom_dos - else: - return dos - - def eval_descriptor( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - mixed_type: bool = False, - ) -> np.array: - """Evaluate descriptors by using this DP. - - Parameters - ---------- - coords - The coordinates of atoms. - The array should be of size nframes x natoms x 3 - cells - The cell of the region. - If None then non-PBC is assumed, otherwise using PBC. - The array should be of size nframes x 9 - atom_types - The atom types - The list should contain natoms ints - fparam - The frame parameter. - The array can be of size : - - nframes x dim_fparam. - - dim_fparam. Then all frames are assumed to be provided with the same fparam. - aparam - The atomic parameter - The array can be of size : - - nframes x natoms x dim_aparam. - - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - - dim_aparam. Then all frames and atoms are provided with the same aparam. - efield - The external field on atoms. - The array should be of size nframes x natoms x 3 - mixed_type - Whether to perform the mixed_type mode. - If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), - in which frames in a system may have different natoms_vec(s), with the same nloc. - - Returns - ------- - descriptor - Descriptors. - """ - natoms, numb_test = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)( - coords, - cells, - atom_types, - fparam=fparam, - aparam=aparam, - efield=efield, - mixed_type=mixed_type, - ) - return descriptor - - def _eval_descriptor_inner( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - mixed_type: bool = False, - ) -> np.array: - natoms, nframes = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - feed_dict_test, imap = self._prepare_feed_dict( - coords, cells, atom_types, fparam, aparam, efield, mixed_type=mixed_type - ) - (descriptor,) = run_sess( - self.sess, [self.t_descriptor], feed_dict=feed_dict_test - ) - return self.reverse_map(np.reshape(descriptor, [nframes, natoms, -1]), imap) +__all__ = [ + "DeepDOS", +] diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 9e3106f4ad..1a917d4f64 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -4,18 +4,50 @@ ) from typing import ( TYPE_CHECKING, + Callable, + ClassVar, + Dict, List, Optional, + Tuple, Union, ) import numpy as np +from deepmd.common import ( + make_default_mesh, +) +from deepmd.infer.deep_dipole import ( + DeepDipole, +) +from deepmd.infer.deep_dos import ( + DeepDOS, +) +from deepmd.infer.deep_eval import ( + DeepEvalBase, +) +from deepmd.infer.deep_polar import ( + DeepGlobalPolar, + DeepPolar, +) +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.infer.deep_wfc import ( + DeepWFC, +) +from deepmd.model_format.output_def import ( + ModelOutputDef, +) from deepmd.tf.env import ( MODEL_VERSION, default_tf_session_config, tf, ) +from deepmd.tf.infer.data_modifier import ( + DipoleChargeModifier, +) from deepmd.tf.utils.batch_size import ( AutoBatchSize, ) @@ -29,7 +61,1047 @@ ) -class DeepEval: +class DeepEval(DeepEvalBase): + """Common methods for DeepPot, DeepWFC, DeepPolar, ... + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + load_prefix: str + The prefix in the load computational graph + default_tf_graph : bool + If uses the default tf graph, otherwise build a new tf graph for evaluation + auto_batch_size : bool or int or AutomaticBatchSize, default: False + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + """ + + load_prefix: str # set by subclass + + def __init__( + self, + model_file: "Path", + output_def: ModelOutputDef, + load_prefix: str = "load", + default_tf_graph: bool = False, + auto_batch_size: Union[bool, int, AutoBatchSize] = False, + input_map: Optional[dict] = None, + neighbor_list=None, + ): + self.graph = self._load_graph( + model_file, + prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, + ) + self.load_prefix = load_prefix + + # graph_compatable should be called after graph and prefix are set + if not self._graph_compatable(): + raise RuntimeError( + f"model in graph (version {self.model_version}) is incompatible" + f"with the model (version {MODEL_VERSION}) supported by the current code." + "See https://deepmd.rtfd.io/compatability/ for details." + ) + + # set default to False, as subclasses may not support + if isinstance(auto_batch_size, bool): + if auto_batch_size: + self.auto_batch_size = AutoBatchSize() + else: + self.auto_batch_size = None + elif isinstance(auto_batch_size, int): + self.auto_batch_size = AutoBatchSize(auto_batch_size) + elif isinstance(auto_batch_size, AutoBatchSize): + self.auto_batch_size = auto_batch_size + else: + raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") + + self.neighbor_list = neighbor_list + + self.output_def = output_def + self._init_tensors() + self._init_attr() + self.has_efield = self.tensors["efield"] is not None + self.has_fparam = self.tensors["fparam"] is not None + self.has_aparam = self.tensors["aparam"] is not None + self.has_spin = self.ntypes_spin > 0 + self.modifier_type = None + + # looks ugly... + if self.modifier_type == "dipole_charge": + t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0") + t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0") + t_sys_charge_map = self._get_tensor("modifier_attr/sys_charge_map:0") + t_ewald_h = self._get_tensor("modifier_attr/ewald_h:0") + t_ewald_beta = self._get_tensor("modifier_attr/ewald_beta:0") + [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = run_sess( + self.sess, + [ + t_mdl_name, + t_mdl_charge_map, + t_sys_charge_map, + t_ewald_h, + t_ewald_beta, + ], + ) + mdl_name = mdl_name.decode("UTF-8") + mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode("UTF-8").split()] + sys_charge_map = [int(ii) for ii in sys_charge_map.decode("UTF-8").split()] + self.dm = DipoleChargeModifier( + mdl_name, + mdl_charge_map, + sys_charge_map, + ewald_h=ewald_h, + ewald_beta=ewald_beta, + ) + + _OUTDEF_DP2TF: ClassVar[dict] = { + "energy": "atom_energy", + "energy_redu": "energy", + "energy_derv_r": "force", + "energy_derv_c": "atom_virial", + "energy_derv_c_redu": "virial", + "polar": "polar", + "polar_redu": "global_polar", + "polar_derv_r": "force", + "polar_derv_c": "atom_virial", + "polar_derv_c_redu": "virial", + "dipole": "dipole", + "dipole_redu": "global_dipole", + "dipole_derv_r": "force", + "dipole_derv_c": "atom_virial", + "dipole_derv_c_redu": "virial", + "dos": "atom_dos", + "dos_redu": "dos", + } + + def _init_tensors(self): + tensor_names = { + # descrpt attrs + "ntypes": "descrpt_attr/ntypes:0", + "rcut": "descrpt_attr/rcut:0", + # model attrs + "tmap": "model_attr/tmap:0", + # inputs + "coord": "t_coord:0", + "type": "t_type:0", + "natoms": "t_natoms:0", + "box": "t_box:0", + "mesh": "t_mesh:0", + } + optional_tensor_names = { + # fitting attrs + "dfparam": "fitting_attr/dfparam:0", + "daparam": "fitting_attr/daparam:0", + "numb_dos": "fitting_attr/numb_dos:0", + # model attrs + "sel_type": "model_attr/sel_type:0", + # additonal inputs + "efield": "t_efield:0", + "fparam": "t_fparam:0", + "aparam": "t_aparam:0", + "ntypes_spin": "descrpt_attr/ntypes_spin:0", + # descriptor + "descriptor": "o_descriptor:0", + } + # output tensors + output_tensor_names = {} + for vv in self.output_def.var_defs: + output_tensor_names[vv] = f"o_{self._OUTDEF_DP2TF[vv]}:0" + + self.tensors = {} + for tensor_key, tensor_name in tensor_names.items(): + self.tensors[tensor_key] = self._get_tensor(tensor_name) + for tensor_key, tensor_name in optional_tensor_names.items(): + try: + self.tensors[tensor_key] = self._get_tensor(tensor_name) + except KeyError: + self.tensors[tensor_key] = None + self.output_tensors = {} + removed_defs = [] + for ii, (tensor_key, tensor_name) in enumerate(output_tensor_names.items()): + try: + self.output_tensors[tensor_key] = self._get_tensor(tensor_name) + except KeyError: + # do not output + removed_defs.append(ii) + for ii in sorted(removed_defs, reverse=True): + del self.output_def.var_defs[list(self.output_def.var_defs.keys())[ii]] + + def _init_attr(self): + [ + self.ntypes, + self.rcut, + tmap, + ] = run_sess( + self.sess, + [ + self.tensors["ntypes"], + self.tensors["rcut"], + self.tensors["tmap"], + ], + ) + if self.tensors["ntypes_spin"] is not None: + self.ntypes_spin = run_sess(self.sess, [self.tensors["ntypes_spin"]])[0] + else: + self.ntypes_spin = 0 + if self.tensors["dfparam"] is not None: + self.dfparam = run_sess(self.sess, [self.tensors["dfparam"]])[0] + else: + self.dfparam = 0 + if self.tensors["daparam"] is not None: + self.daparam = run_sess(self.sess, [self.tensors["daparam"]])[0] + else: + self.daparam = 0 + if self.tensors["sel_type"] is not None: + self.sel_type = run_sess(self.sess, [self.tensors["sel_type"]])[0] + else: + self.sel_type = None + if self.tensors["numb_dos"] is not None: + self.numb_dos = run_sess(self.sess, [self.tensors["numb_dos"]])[0] + else: + self.numb_dos = 0 + self.tmap = tmap.decode("utf-8").split() + + @property + @lru_cache(maxsize=None) + def model_type(self) -> str: + """Get type of model. + + :type:str + """ + t_mt = self._get_tensor("model_attr/model_type:0") + [mt] = run_sess(self.sess, [t_mt], feed_dict={}) + model_type = mt.decode("utf-8") + if model_type == "ener": + return DeepPot + elif model_type == "dos": + return DeepDOS + elif model_type == "dipole": + return DeepDipole + elif model_type == "polar": + return DeepPolar + elif model_type == "global_polar": + return DeepGlobalPolar + elif model_type == "wfc": + return DeepWFC + else: + raise RuntimeError(f"unknown model type {model_type}") + + @property + @lru_cache(maxsize=None) + def model_version(self) -> str: + """Get version of model. + + Returns + ------- + str + version of model + """ + try: + t_mt = self._get_tensor("model_attr/model_version:0") + except KeyError: + # For deepmd-kit version 0.x - 1.x, set model version to 0.0 + return "0.0" + else: + [mt] = run_sess(self.sess, [t_mt], feed_dict={}) + return mt.decode("utf-8") + + @property + @lru_cache(maxsize=None) + def sess(self) -> tf.Session: + """Get TF session.""" + # start a tf session associated to the graph + return tf.Session(graph=self.graph, config=default_tf_session_config) + + def _graph_compatable(self) -> bool: + """Check the model compatability. + + Returns + ------- + bool + If the model stored in the graph file is compatable with the current code + """ + model_version_major = int(self.model_version.split(".")[0]) + model_version_minor = int(self.model_version.split(".")[1]) + MODEL_VERSION_MAJOR = int(MODEL_VERSION.split(".")[0]) + MODEL_VERSION_MINOR = int(MODEL_VERSION.split(".")[1]) + if (model_version_major != MODEL_VERSION_MAJOR) or ( + model_version_minor > MODEL_VERSION_MINOR + ): + return False + else: + return True + + def _get_tensor( + self, + tensor_name: str, + ) -> tf.Tensor: + """Get TF graph tensor. + + Parameters + ---------- + tensor_name : str + name of tensor to get + + Returns + ------- + tf.Tensor + loaded tensor + """ + # do not use os.path.join as it doesn't work on Windows + tensor_path = "/".join((self.load_prefix, tensor_name)) + tensor = self.graph.get_tensor_by_name(tensor_path) + return tensor + + @staticmethod + def _load_graph( + frozen_graph_filename: "Path", + prefix: str = "load", + default_tf_graph: bool = False, + input_map: Optional[dict] = None, + ): + # We load the protobuf file from the disk and parse it to retrieve the + # unserialized graph_def + with tf.gfile.GFile(str(frozen_graph_filename), "rb") as f: + graph_def = tf.GraphDef() + graph_def.ParseFromString(f.read()) + + if default_tf_graph: + tf.import_graph_def( + graph_def, + input_map=input_map, + return_elements=None, + name=prefix, + producer_op_list=None, + ) + graph = tf.get_default_graph() + else: + # Then, we can use again a convenient built-in function to import + # a graph_def into the current default Graph + with tf.Graph().as_default() as graph: + tf.import_graph_def( + graph_def, + input_map=None, + return_elements=None, + name=prefix, + producer_op_list=None, + ) + + return graph + + @staticmethod + def sort_input( + coord: np.ndarray, + atom_type: np.ndarray, + sel_atoms: Optional[List[int]] = None, + ): + """Sort atoms in the system according their types. + + Parameters + ---------- + coord + The coordinates of atoms. + Should be of shape [nframes, natoms, 3] + atom_type + The type of atoms + Should be of shape [natoms] + sel_atoms + The selected atoms by type + + Returns + ------- + coord_out + The coordinates after sorting + atom_type_out + The atom types after sorting + idx_map + The index mapping from the input to the output. + For example coord_out = coord[:,idx_map,:] + sel_atom_type + Only output if sel_atoms is not None + The sorted selected atom types + sel_idx_map + Only output if sel_atoms is not None + The index mapping from the selected atoms to sorted selected atoms. + """ + natoms = atom_type.shape[1] + if sel_atoms is not None: + selection = [False] * natoms + for ii in sel_atoms: + selection += atom_type[0] == ii + sel_atom_type = atom_type[:, selection] + idx = np.arange(natoms) + idx_map = np.lexsort((idx, atom_type[0])) + nframes = coord.shape[0] + coord = coord.reshape([nframes, -1, 3]) + coord = np.reshape(coord[:, idx_map, :], [nframes, -1]) + atom_type = atom_type[:, idx_map] + if sel_atoms is not None: + sel_natoms = sel_atom_type.shape[1] + sel_idx = np.arange(sel_natoms) + sel_idx_map = np.lexsort((sel_idx, sel_atom_type[0])) + sel_atom_type = sel_atom_type[:, sel_idx_map] + return coord, atom_type, idx_map, sel_atom_type, sel_idx_map + else: + return coord, atom_type, idx_map, atom_type, idx_map + + @staticmethod + def reverse_map(vec: np.ndarray, imap: List[int]) -> np.ndarray: + """Reverse mapping of a vector according to the index map. + + Parameters + ---------- + vec + Input vector. Be of shape [nframes, natoms, -1] + imap + Index map. Be of shape [natoms] + + Returns + ------- + vec_out + Reverse mapped vector. + """ + ret = np.zeros(vec.shape) + # for idx,ii in enumerate(imap) : + # ret[:,ii,:] = vec[:,idx,:] + ret[:, imap, :] = vec + return ret + + def make_natoms_vec( + self, + atom_types: np.ndarray, + ) -> np.ndarray: + """Make the natom vector used by deepmd-kit. + + Parameters + ---------- + atom_types + The type of atoms + + Returns + ------- + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + + """ + natoms_vec = np.zeros(self.ntypes + 2).astype(int) + natoms = atom_types[0].size + natoms_vec[0] = natoms + natoms_vec[1] = natoms + for ii in range(self.ntypes): + natoms_vec[ii + 2] = np.count_nonzero(atom_types[0] == ii) + return natoms_vec + + def eval_typeebd(self) -> np.ndarray: + """Evaluate output of type embedding network by using this model. + + Returns + ------- + np.ndarray + The output of type embedding network. The shape is [ntypes, o_size], + where ntypes is the number of types, and o_size is the number of nodes + in the output layer. + + Raises + ------ + KeyError + If the model does not enable type embedding. + + See Also + -------- + deepmd.tf.utils.type_embed.TypeEmbedNet : The type embedding network. + + Examples + -------- + Get the output of type embedding network of `graph.pb`: + + >>> from deepmd.tf.infer import DeepPotential + >>> dp = DeepPotential('graph.pb') + >>> dp.eval_typeebd() + """ + t_typeebd = self._get_tensor("t_typeebd:0") + [typeebd] = run_sess(self.sess, [t_typeebd], feed_dict={}) + return typeebd + + def build_neighbor_list( + self, + coords: np.ndarray, + cell: Optional[np.ndarray], + atype: np.ndarray, + imap: np.ndarray, + neighbor_list, + ): + """Make the mesh with neighbor list for a single frame. + + Parameters + ---------- + coords : np.ndarray + The coordinates of atoms. Should be of shape [natoms, 3] + cell : Optional[np.ndarray] + The cell of the system. Should be of shape [3, 3] + atype : np.ndarray + The type of atoms. Should be of shape [natoms] + imap : np.ndarray + The index map of atoms. Should be of shape [natoms] + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList + ASE neighbor list. The following method or attribute will be + used/set: bothways, self_interaction, update, build, first_neigh, + pair_second, offset_vec. + + Returns + ------- + natoms_vec : np.ndarray + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: nloc + natoms[1]: nall + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms for nloc + coords : np.ndarray + The coordinates of atoms, including ghost atoms. Should be of + shape [nframes, nall, 3] + atype : np.ndarray + The type of atoms, including ghost atoms. Should be of shape [nall] + mesh : np.ndarray + The mesh in nei_mode=4. + imap : np.ndarray + The index map of atoms. Should be of shape [nall] + ghost_map : np.ndarray + The index map of ghost atoms. Should be of shape [nghost] + """ + pbc = np.repeat(cell is not None, 3) + cell = cell.reshape(3, 3) + positions = coords.reshape(-1, 3) + neighbor_list.bothways = True + neighbor_list.self_interaction = False + if neighbor_list.update(pbc, cell, positions): + neighbor_list.build(pbc, cell, positions) + first_neigh = neighbor_list.first_neigh.copy() + pair_second = neighbor_list.pair_second.copy() + offset_vec = neighbor_list.offset_vec.copy() + # get out-of-box neighbors + out_mask = np.any(offset_vec != 0, axis=1) + out_idx = pair_second[out_mask] + out_offset = offset_vec[out_mask] + out_coords = positions[out_idx] + out_offset.dot(cell) + atype = np.array(atype, dtype=int).reshape(-1) + out_atype = atype[out_idx] + + nloc = positions.shape[0] + nghost = out_idx.size + all_coords = np.concatenate((positions, out_coords), axis=0) + all_atype = np.concatenate((atype, out_atype), axis=0) + # convert neighbor indexes + ghost_map = pair_second[out_mask] + pair_second[out_mask] = np.arange(nloc, nloc + nghost) + # get the mesh + mesh = np.zeros(16 + nloc * 2 + pair_second.size, dtype=int) + mesh[0] = nloc + # ilist + mesh[16 : 16 + nloc] = np.arange(nloc) + # numnei + mesh[16 + nloc : 16 + nloc * 2] = first_neigh[1:] - first_neigh[:-1] + # jlist + mesh[16 + nloc * 2 :] = pair_second + + # natoms_vec + natoms_vec = np.zeros(self.ntypes + 2).astype(int) + natoms_vec[0] = nloc + natoms_vec[1] = nloc + nghost + for ii in range(self.ntypes): + natoms_vec[ii + 2] = np.count_nonzero(atype == ii) + # imap append ghost atoms + imap = np.concatenate((imap, np.arange(nloc, nloc + nghost))) + return natoms_vec, all_coords, all_atype, mesh, imap, ghost_map + + def get_ntypes(self) -> int: + """Get the number of atom types of this model.""" + return self.ntypes + + def get_ntypes_spin(self): + """Get the number of spin atom types of this model.""" + return self.ntypes_spin + + def get_rcut(self) -> float: + """Get the cut-off radius of this model.""" + return self.rcut + + def get_type_map(self) -> List[str]: + """Get the type map (element name of the atom types) of this model.""" + return self.tmap + + def get_sel_type(self) -> Optional[np.ndarray]: + """Get selection type.""" + return np.array(self.sel_type).ravel() + + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this DP.""" + return self.dfparam + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this DP.""" + return self.daparam + + def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: + """Wrapper method with auto batch size. + + Parameters + ---------- + inner_func : Callable + the method to be wrapped + numb_test : int + number of tests + natoms : int + number of atoms + + Returns + ------- + Callable + the wrapper + """ + if self.auto_batch_size is not None: + + def eval_func(*args, **kwargs): + return self.auto_batch_size.execute_all( + inner_func, numb_test, natoms, *args, **kwargs + ) + + else: + eval_func = inner_func + return eval_func + + def _get_natoms_and_nframes( + self, + coords: np.ndarray, + atom_types: Union[List[int], np.ndarray], + ) -> Tuple[int, int]: + natoms = len(atom_types[0]) + if natoms == 0: + assert coords.size == 0 + else: + coords = np.reshape(np.array(coords), [-1, natoms * 3]) + nframes = coords.shape[0] + return natoms, nframes + + def eval( + self, + coords: np.ndarray, + cells: np.ndarray, + atom_types: List[int], + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + efield: Optional[np.ndarray] = None, + ) -> Dict[str, np.ndarray]: + """Evaluate the energy, force and virial by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Calculate the atomic energy and virial + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + efield + The external field on atoms. + The array should be of size nframes x natoms x 3 + + Returns + ------- + output_dict : dict + The output of the evaluation. The keys are the names of the output + variables, and the values are the corresponding output arrays. + """ + # reshape coords before getting shape + natoms, numb_test = self._get_natoms_and_nframes( + coords, + atom_types, + ) + output = self._eval_func(self._eval_inner, numb_test, natoms)( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + atomic=atomic, + efield=efield, + ) + if not isinstance(output, tuple): + output = (output,) + + output_dict = { + odef.name: oo for oo, odef in zip(output, self.output_def.var_defs.values()) + } + # ugly!! + if self.modifier_type is not None and isinstance(self.model_type, DeepPot): + if atomic: + raise RuntimeError("modifier does not support atomic modification") + me, mf, mv = self.dm.eval(coords, cells, atom_types) + output = list(output) # tuple to list + e, f, v = output[:3] + output_dict["energy_redu"] += me.reshape(e.shape) + output_dict["energy_deri_r"] += mf.reshape(f.shape) + output_dict["energy_deri_c_redu"] += mv.reshape(v.shape) + output = tuple(output) + return output_dict + + def _prepare_feed_dict( + self, + coords, + cells, + atom_types, + fparam=None, + aparam=None, + efield=None, + ): + # standarize the shape of inputs + natoms, nframes = self._get_natoms_and_nframes( + coords, + atom_types, + ) + atom_types = np.array(atom_types, dtype=int).reshape([-1, natoms]) + coords = np.reshape(np.array(coords), [nframes, natoms * 3]) + if cells is None: + pbc = False + # make cells to work around the requirement of pbc + cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9]) + else: + pbc = True + cells = np.array(cells).reshape([nframes, 9]) + + if self.has_fparam: + assert fparam is not None + fparam = np.array(fparam) + if self.has_aparam: + assert aparam is not None + aparam = np.array(aparam) + if self.has_efield: + assert ( + efield is not None + ), "you are using a model with external field, parameter efield should be provided" + efield = np.array(efield) + + # reshape the inputs + if self.has_fparam: + fdim = self.get_dim_fparam() + if fparam.size == nframes * fdim: + fparam = np.reshape(fparam, [nframes, fdim]) + elif fparam.size == fdim: + fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) + else: + raise RuntimeError( + "got wrong size of frame param, should be either %d x %d or %d" + % (nframes, fdim, fdim) + ) + if self.has_aparam: + fdim = self.get_dim_aparam() + if aparam.size == nframes * natoms * fdim: + aparam = np.reshape(aparam, [nframes, natoms * fdim]) + elif aparam.size == natoms * fdim: + aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) + elif aparam.size == fdim: + aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) + else: + raise RuntimeError( + "got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d" + % (nframes, natoms, fdim, natoms, fdim, fdim) + ) + + # sort inputs + coords, atom_types, imap, sel_at, sel_imap = self.sort_input( + coords, atom_types, sel_atoms=self.get_sel_type() + ) + if self.has_efield: + efield = np.reshape(efield, [nframes, natoms, 3]) + efield = efield[:, imap, :] + efield = np.reshape(efield, [nframes, natoms * 3]) + if self.has_aparam: + aparam = np.reshape(aparam, [nframes, natoms, fdim]) + aparam = aparam[:, imap, :] + aparam = np.reshape(aparam, [nframes, natoms * fdim]) + + # make natoms_vec and default_mesh + if self.neighbor_list is None: + natoms_vec = self.make_natoms_vec(atom_types) + assert natoms_vec[0] == natoms + mesh = make_default_mesh( + pbc, not self._check_distinguished_types(atom_types) + ) + ghost_map = None + else: + if nframes > 1: + raise NotImplementedError( + "neighbor_list does not support multiple frames" + ) + ( + natoms_vec, + coords, + atom_types, + mesh, + imap, + ghost_map, + ) = self.build_neighbor_list( + coords, + cells if cells is not None else None, + atom_types, + imap, + self.neighbor_list, + ) + + # evaluate + feed_dict_test = {} + feed_dict_test[self.tensors["natoms"]] = natoms_vec + feed_dict_test[self.tensors["type"]] = atom_types.reshape([-1]) + feed_dict_test[self.tensors["coord"]] = np.reshape(coords, [-1]) + + if len(self.tensors["box"].shape) == 1: + feed_dict_test[self.tensors["box"]] = np.reshape(cells, [-1]) + elif len(self.tensors["box"].shape) == 2: + feed_dict_test[self.tensors["box"]] = cells + else: + raise RuntimeError + if self.has_efield: + feed_dict_test[self.tensors["efield"]] = np.reshape(efield, [-1]) + feed_dict_test[self.tensors["mesh"]] = mesh + if self.has_fparam: + feed_dict_test[self.tensors["fparam"]] = np.reshape(fparam, [-1]) + if self.has_aparam: + feed_dict_test[self.tensors["aparam"]] = np.reshape(aparam, [-1]) + return feed_dict_test, imap, natoms_vec, ghost_map, sel_at, sel_imap + + def _eval_inner( + self, + coords, + cells, + atom_types, + fparam=None, + aparam=None, + efield=None, + **kwargs, + ): + natoms, nframes = self._get_natoms_and_nframes( + coords, + atom_types, + ) + ( + feed_dict_test, + imap, + natoms_vec, + ghost_map, + sel_at, + sel_imap, + ) = self._prepare_feed_dict( + coords, + cells, + atom_types, + fparam, + aparam, + efield, + ) + + nloc = natoms_vec[0] + nloc_sel = sel_at.shape[1] + nall = natoms_vec[1] + + t_out = list(self.output_tensors.values()) + + v_out = run_sess(self.sess, t_out, feed_dict=feed_dict_test) + + if nloc_sel == 0: + nloc_sel = nloc + sel_imap = imap + if self.has_spin: + ntypes_real = self.ntypes - self.ntypes_spin + natoms_real = sum( + [ + np.count_nonzero(np.array(atom_types) == ii) + for ii in range(ntypes_real) + ] + ) + else: + natoms_real = nloc_sel + if ghost_map is not None: + # add the value of ghost atoms to real atoms + for ii, odef in enumerate(self.output_def.var_defs.values()): + # when the shape is nall + if "_redu" not in odef.name and "_derv" in odef.name: + odef_shape = self._get_output_shape( + odef.name, nframes, nall, odef.shape + ) + tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] + v_out[ii] = np.reshape(v_out[ii], tmp_shape) + for jj in range(v_out[ii].shape[0]): + np.add.at(v_out[ii][jj], ghost_map, v_out[ii][jj, nloc:]) + + for ii, odef in enumerate(self.output_def.var_defs.values()): + if "_redu" not in odef.name: + if "_derv" in odef.name: + odef_shape = self._get_output_shape( + odef.name, nframes, nall, odef.shape + ) + tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] + # reverse map of the outputs + v_out[ii] = self.reverse_map(np.reshape(v_out[ii], tmp_shape), imap) + v_out[ii] = np.reshape(v_out[ii], odef_shape) + if nloc < nall: + v_out[ii] = v_out[ii][:, :, :nloc] + else: + odef_shape = self._get_output_shape( + odef.name, nframes, natoms_real, odef.shape + ) + # tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] + v_out[ii] = self.reverse_map( + np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real] + ) + v_out[ii] = np.reshape(v_out[ii], odef_shape) + else: + odef_shape = self._get_output_shape(odef.name, nframes, 0, odef.shape) + v_out[ii] = np.reshape(v_out[ii], odef_shape) + return tuple(v_out) + + def _get_output_shape(self, name, nframes, natoms, shape): + if "_redu" in name: + if "_derv_c" in name: + return [nframes, *shape[:-2], 9] + else: + return [nframes, *shape, 1] + else: + if "_derv_c" in name: + return [nframes, *shape[:-2], natoms, 9] + elif "_derv_r" in name: + return [nframes, *shape[:-1], natoms, 3] + else: + # Something wrong here? + # return [nframes, *shape, natoms, 1] + return [nframes, natoms, *shape, 1] + + def eval_descriptor( + self, + coords: np.ndarray, + cells: np.ndarray, + atom_types: List[int], + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + efield: Optional[np.ndarray] = None, + ) -> np.ndarray: + """Evaluate descriptors by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + efield + The external field on atoms. + The array should be of size nframes x natoms x 3 + + Returns + ------- + descriptor + Descriptors. + """ + natoms, numb_test = self._get_natoms_and_nframes( + coords, + atom_types, + ) + descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + efield=efield, + ) + return descriptor + + def _eval_descriptor_inner( + self, + coords: np.ndarray, + cells: np.ndarray, + atom_types: List[int], + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + efield: Optional[np.ndarray] = None, + ) -> np.ndarray: + natoms, nframes = self._get_natoms_and_nframes( + coords, + atom_types, + ) + ( + feed_dict_test, + imap, + natoms_vec, + ghost_map, + sel_at, + sel_imap, + ) = self._prepare_feed_dict( + coords, + cells, + atom_types, + fparam, + aparam, + efield, + ) + (descriptor,) = run_sess( + self.sess, [self.tensors["descriptor"]], feed_dict=feed_dict_test + ) + imap = imap[:natoms] + return self.reverse_map(np.reshape(descriptor, [nframes, natoms, -1]), imap) + + def get_numb_dos(self) -> int: + return self.numb_dos + + def get_has_efield(self) -> bool: + return self.has_efield + + +class DeepEvalOld: + # old class for DipoleChargeModifier only """Common methods for DeepPot, DeepWFC, DeepPolar, ... Parameters diff --git a/deepmd/tf/infer/deep_polar.py b/deepmd/tf/infer/deep_polar.py index e0b73da2a2..c3d42fd537 100644 --- a/deepmd/tf/infer/deep_polar.py +++ b/deepmd/tf/infer/deep_polar.py @@ -1,165 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - TYPE_CHECKING, - List, - Optional, +from deepmd.infer.deep_polar import ( + DeepGlobalPolar, + DeepPolar, ) -import numpy as np - -from deepmd.tf.infer.deep_tensor import ( - DeepTensor, -) - -if TYPE_CHECKING: - from pathlib import ( - Path, - ) - - -class DeepPolar(DeepTensor): - """Constructor. - - Parameters - ---------- - model_file : Path - The name of the frozen model file. - load_prefix: str - The prefix in the load computational graph - default_tf_graph : bool - If uses the default tf graph, otherwise build a new tf graph for evaluation - input_map : dict, optional - The input map for tf.import_graph_def. Only work with default tf graph - neighbor_list : ase.neighborlist.NeighborList, optional - The neighbor list object. If None, then build the native neighbor list. - - Warnings - -------- - For developers: `DeepTensor` initializer must be called at the end after - `self.tensors` are modified because it uses the data in `self.tensors` dict. - Do not chanage the order! - """ - - def __init__( - self, - model_file: "Path", - load_prefix: str = "load", - default_tf_graph: bool = False, - input_map: Optional[dict] = None, - neighbor_list=None, - ) -> None: - # use this in favor of dict update to move attribute from class to - # instance namespace - self.tensors = dict( - { - # output tensor - "t_tensor": "o_polar:0", - }, - **self.tensors, - ) - - DeepTensor.__init__( - self, - model_file, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - neighbor_list=neighbor_list, - ) - - def get_dim_fparam(self) -> int: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") - - def get_dim_aparam(self) -> int: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") - - -class DeepGlobalPolar(DeepTensor): - """Constructor. - - Parameters - ---------- - model_file : str - The name of the frozen model file. - load_prefix: str - The prefix in the load computational graph - default_tf_graph : bool - If uses the default tf graph, otherwise build a new tf graph for evaluation - neighbor_list : ase.neighborlist.NeighborList, optional - The neighbor list object. If None, then build the native neighbor list. - """ - - def __init__( - self, - model_file: str, - load_prefix: str = "load", - default_tf_graph: bool = False, - neighbor_list=None, - ) -> None: - self.tensors.update( - { - "t_sel_type": "model_attr/sel_type:0", - # output tensor - "t_tensor": "o_global_polar:0", - } - ) - - DeepTensor.__init__( - self, - model_file, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - neighbor_list=None, - ) - - def eval( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - ) -> np.ndarray: - """Evaluate the model. - - Parameters - ---------- - coords - The coordinates of atoms. - The array should be of size nframes x natoms x 3 - cells - The cell of the region. - If None then non-PBC is assumed, otherwise using PBC. - The array should be of size nframes x 9 - atom_types - The atom types - The list should contain natoms ints - atomic - Not used in this model - fparam - Not used in this model - aparam - Not used in this model - efield - Not used in this model - - Returns - ------- - tensor - The returned tensor - If atomic == False then of size nframes x variable_dof - else of size nframes x natoms x variable_dof - """ - return DeepTensor.eval(self, coords, cells, atom_types, atomic=False) - - def get_dim_fparam(self) -> int: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") - - def get_dim_aparam(self) -> int: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") +__all__ = [ + "DeepPolar", + "DeepGlobalPolar", +] diff --git a/deepmd/tf/infer/deep_pot.py b/deepmd/tf/infer/deep_pot.py index 0663d2daee..587a13996a 100644 --- a/deepmd/tf/infer/deep_pot.py +++ b/deepmd/tf/infer/deep_pot.py @@ -1,692 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import logging -from typing import ( - TYPE_CHECKING, - Callable, - List, - Optional, - Tuple, - Union, +from deepmd.infer import ( + DeepPot, ) -import numpy as np - -from deepmd.infer.deep_pot import DeepPot as DeepPotBase -from deepmd.tf.common import ( - make_default_mesh, -) -from deepmd.tf.infer.data_modifier import ( - DipoleChargeModifier, -) -from deepmd.tf.infer.deep_eval import ( - DeepEval, -) -from deepmd.tf.utils.batch_size import ( - AutoBatchSize, -) -from deepmd.tf.utils.sess import ( - run_sess, -) - -if TYPE_CHECKING: - from pathlib import ( - Path, - ) - -log = logging.getLogger(__name__) - - -class DeepPot(DeepEval, DeepPotBase): - """Constructor. - - Parameters - ---------- - model_file : Path - The name of the frozen model file. - load_prefix: str - The prefix in the load computational graph - default_tf_graph : bool - If uses the default tf graph, otherwise build a new tf graph for evaluation - auto_batch_size : bool or int or AutomaticBatchSize, default: True - If True, automatic batch size will be used. If int, it will be used - as the initial batch size. - input_map : dict, optional - The input map for tf.import_graph_def. Only work with default tf graph - neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional - The ASE neighbor list class to produce the neighbor list. If None, the - neighbor list will be built natively in the model. - - Examples - -------- - >>> from deepmd.tf.infer import DeepPot - >>> import numpy as np - >>> dp = DeepPot('graph.pb') - >>> coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1]) - >>> cell = np.diag(10 * np.ones(3)).reshape([1, -1]) - >>> atype = [1,0,1] - >>> e, f, v = dp.eval(coord, cell, atype) - - where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively. - - Warnings - -------- - For developers: `DeepTensor` initializer must be called at the end after - `self.tensors` are modified because it uses the data in `self.tensors` dict. - Do not chanage the order! - """ - - def __init__( - self, - model_file: "Path", - load_prefix: str = "load", - default_tf_graph: bool = False, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, - input_map: Optional[dict] = None, - neighbor_list=None, - ) -> None: - # add these tensors on top of what is defined by DeepTensor Class - # use this in favor of dict update to move attribute from class to - # instance namespace - self.tensors = { - # descrpt attrs - "t_ntypes": "descrpt_attr/ntypes:0", - "t_rcut": "descrpt_attr/rcut:0", - # fitting attrs - "t_dfparam": "fitting_attr/dfparam:0", - "t_daparam": "fitting_attr/daparam:0", - # model attrs - "t_tmap": "model_attr/tmap:0", - # inputs - "t_coord": "t_coord:0", - "t_type": "t_type:0", - "t_natoms": "t_natoms:0", - "t_box": "t_box:0", - "t_mesh": "t_mesh:0", - # add output tensors - "t_energy": "o_energy:0", - "t_force": "o_force:0", - "t_virial": "o_virial:0", - "t_ae": "o_atom_energy:0", - "t_av": "o_atom_virial:0", - "t_descriptor": "o_descriptor:0", - } - DeepEval.__init__( - self, - model_file, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - auto_batch_size=auto_batch_size, - input_map=input_map, - neighbor_list=neighbor_list, - ) - - # load optional tensors - operations = [op.name for op in self.graph.get_operations()] - # check if the graph has these operations: - # if yes add them - - if ("%s/t_efield" % load_prefix) in operations: - self.tensors.update({"t_efield": "t_efield:0"}) - self.has_efield = True - else: - log.debug("Could not get tensor 't_efield:0'") - self.t_efield = None - self.has_efield = False - - if ("%s/t_fparam" % load_prefix) in operations: - self.tensors.update({"t_fparam": "t_fparam:0"}) - self.has_fparam = True - else: - log.debug("Could not get tensor 't_fparam:0'") - self.t_fparam = None - self.has_fparam = False - - if ("%s/t_aparam" % load_prefix) in operations: - self.tensors.update({"t_aparam": "t_aparam:0"}) - self.has_aparam = True - else: - log.debug("Could not get tensor 't_aparam:0'") - self.t_aparam = None - self.has_aparam = False - - if ("%s/spin_attr/ntypes_spin" % load_prefix) in operations: - self.tensors.update({"t_ntypes_spin": "spin_attr/ntypes_spin:0"}) - self.has_spin = True - else: - self.ntypes_spin = 0 - self.has_spin = False - - # now load tensors to object attributes - for attr_name, tensor_name in self.tensors.items(): - try: - self._get_tensor(tensor_name, attr_name) - except KeyError: - if attr_name != "t_descriptor": - raise - - self._run_default_sess() - self.tmap = self.tmap.decode("UTF-8").split() - - # setup modifier - try: - t_modifier_type = self._get_tensor("modifier_attr/type:0") - self.modifier_type = run_sess(self.sess, t_modifier_type).decode("UTF-8") - except (ValueError, KeyError): - self.modifier_type = None - - try: - t_jdata = self._get_tensor("train_attr/training_script:0") - jdata = run_sess(self.sess, t_jdata).decode("UTF-8") - import json - - jdata = json.loads(jdata) - self.descriptor_type = jdata["model"]["descriptor"]["type"] - except (ValueError, KeyError): - self.descriptor_type = None - - if self.modifier_type == "dipole_charge": - t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0") - t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0") - t_sys_charge_map = self._get_tensor("modifier_attr/sys_charge_map:0") - t_ewald_h = self._get_tensor("modifier_attr/ewald_h:0") - t_ewald_beta = self._get_tensor("modifier_attr/ewald_beta:0") - [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = run_sess( - self.sess, - [ - t_mdl_name, - t_mdl_charge_map, - t_sys_charge_map, - t_ewald_h, - t_ewald_beta, - ], - ) - mdl_name = mdl_name.decode("UTF-8") - mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode("UTF-8").split()] - sys_charge_map = [int(ii) for ii in sys_charge_map.decode("UTF-8").split()] - self.dm = DipoleChargeModifier( - mdl_name, - mdl_charge_map, - sys_charge_map, - ewald_h=ewald_h, - ewald_beta=ewald_beta, - ) - - def _run_default_sess(self): - if self.has_spin is True: - [ - self.ntypes, - self.ntypes_spin, - self.rcut, - self.dfparam, - self.daparam, - self.tmap, - ] = run_sess( - self.sess, - [ - self.t_ntypes, - self.t_ntypes_spin, - self.t_rcut, - self.t_dfparam, - self.t_daparam, - self.t_tmap, - ], - ) - else: - [self.ntypes, self.rcut, self.dfparam, self.daparam, self.tmap] = run_sess( - self.sess, - [ - self.t_ntypes, - self.t_rcut, - self.t_dfparam, - self.t_daparam, - self.t_tmap, - ], - ) - - def get_ntypes(self) -> int: - """Get the number of atom types of this model.""" - return self.ntypes - - def get_ntypes_spin(self): - """Get the number of spin atom types of this model.""" - return self.ntypes_spin - - def get_rcut(self) -> float: - """Get the cut-off radius of this model.""" - return self.rcut - - def get_type_map(self) -> List[str]: - """Get the type map (element name of the atom types) of this model.""" - return self.tmap - - def get_sel_type(self) -> List[int]: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") - - def get_descriptor_type(self) -> List[int]: - """Get the descriptor type of this model.""" - return self.descriptor_type - - def get_dim_fparam(self) -> int: - """Get the number (dimension) of frame parameters of this DP.""" - return self.dfparam - - def get_dim_aparam(self) -> int: - """Get the number (dimension) of atomic parameters of this DP.""" - return self.daparam - - def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: - """Wrapper method with auto batch size. - - Parameters - ---------- - inner_func : Callable - the method to be wrapped - numb_test : int - number of tests - natoms : int - number of atoms - - Returns - ------- - Callable - the wrapper - """ - if self.auto_batch_size is not None: - - def eval_func(*args, **kwargs): - return self.auto_batch_size.execute_all( - inner_func, numb_test, natoms, *args, **kwargs - ) - - else: - eval_func = inner_func - return eval_func - - def _get_natoms_and_nframes( - self, - coords: np.ndarray, - atom_types: Union[List[int], np.ndarray], - mixed_type: bool = False, - ) -> Tuple[int, int]: - if mixed_type: - natoms = len(atom_types[0]) - else: - natoms = len(atom_types) - if natoms == 0: - assert coords.size == 0 - else: - coords = np.reshape(np.array(coords), [-1, natoms * 3]) - nframes = coords.shape[0] - return natoms, nframes - - def eval( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - mixed_type: bool = False, - ) -> Tuple[np.ndarray, ...]: - """Evaluate the energy, force and virial by using this DP. - - Parameters - ---------- - coords - The coordinates of atoms. - The array should be of size nframes x natoms x 3 - cells - The cell of the region. - If None then non-PBC is assumed, otherwise using PBC. - The array should be of size nframes x 9 - atom_types - The atom types - The list should contain natoms ints - atomic - Calculate the atomic energy and virial - fparam - The frame parameter. - The array can be of size : - - nframes x dim_fparam. - - dim_fparam. Then all frames are assumed to be provided with the same fparam. - aparam - The atomic parameter - The array can be of size : - - nframes x natoms x dim_aparam. - - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - - dim_aparam. Then all frames and atoms are provided with the same aparam. - efield - The external field on atoms. - The array should be of size nframes x natoms x 3 - mixed_type - Whether to perform the mixed_type mode. - If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), - in which frames in a system may have different natoms_vec(s), with the same nloc. - - Returns - ------- - energy - The system energy. - force - The force on each atom - virial - The virial - atom_energy - The atomic energy. Only returned when atomic == True - atom_virial - The atomic virial. Only returned when atomic == True - """ - # reshape coords before getting shape - natoms, numb_test = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - output = self._eval_func(self._eval_inner, numb_test, natoms)( - coords, - cells, - atom_types, - fparam=fparam, - aparam=aparam, - atomic=atomic, - efield=efield, - mixed_type=mixed_type, - ) - - if self.modifier_type is not None: - if atomic: - raise RuntimeError("modifier does not support atomic modification") - me, mf, mv = self.dm.eval(coords, cells, atom_types) - output = list(output) # tuple to list - e, f, v = output[:3] - output[0] += me.reshape(e.shape) - output[1] += mf.reshape(f.shape) - output[2] += mv.reshape(v.shape) - output = tuple(output) - return output - - def _prepare_feed_dict( - self, - coords, - cells, - atom_types, - fparam=None, - aparam=None, - efield=None, - mixed_type=False, - ): - # standarize the shape of inputs - natoms, nframes = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - if mixed_type: - atom_types = np.array(atom_types, dtype=int).reshape([-1, natoms]) - else: - atom_types = np.array(atom_types, dtype=int).reshape([-1]) - coords = np.reshape(np.array(coords), [nframes, natoms * 3]) - if cells is None: - pbc = False - # make cells to work around the requirement of pbc - cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9]) - else: - pbc = True - cells = np.array(cells).reshape([nframes, 9]) - - if self.has_fparam: - assert fparam is not None - fparam = np.array(fparam) - if self.has_aparam: - assert aparam is not None - aparam = np.array(aparam) - if self.has_efield: - assert ( - efield is not None - ), "you are using a model with external field, parameter efield should be provided" - efield = np.array(efield) - - # reshape the inputs - if self.has_fparam: - fdim = self.get_dim_fparam() - if fparam.size == nframes * fdim: - fparam = np.reshape(fparam, [nframes, fdim]) - elif fparam.size == fdim: - fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) - else: - raise RuntimeError( - "got wrong size of frame param, should be either %d x %d or %d" - % (nframes, fdim, fdim) - ) - if self.has_aparam: - fdim = self.get_dim_aparam() - if aparam.size == nframes * natoms * fdim: - aparam = np.reshape(aparam, [nframes, natoms * fdim]) - elif aparam.size == natoms * fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) - elif aparam.size == fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) - else: - raise RuntimeError( - "got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d" - % (nframes, natoms, fdim, natoms, fdim, fdim) - ) - - # sort inputs - coords, atom_types, imap = self.sort_input( - coords, atom_types, mixed_type=mixed_type - ) - if self.has_efield: - efield = np.reshape(efield, [nframes, natoms, 3]) - efield = efield[:, imap, :] - efield = np.reshape(efield, [nframes, natoms * 3]) - if self.has_aparam: - aparam = np.reshape(aparam, [nframes, natoms, fdim]) - aparam = aparam[:, imap, :] - aparam = np.reshape(aparam, [nframes, natoms * fdim]) - - # make natoms_vec and default_mesh - if self.neighbor_list is None: - natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) - assert natoms_vec[0] == natoms - mesh = make_default_mesh(pbc, mixed_type) - ghost_map = None - else: - if nframes > 1: - raise NotImplementedError( - "neighbor_list does not support multiple frames" - ) - ( - natoms_vec, - coords, - atom_types, - mesh, - imap, - ghost_map, - ) = self.build_neighbor_list( - coords, - cells if cells is not None else None, - atom_types, - imap, - self.neighbor_list, - ) - - # evaluate - feed_dict_test = {} - feed_dict_test[self.t_natoms] = natoms_vec - if mixed_type: - feed_dict_test[self.t_type] = atom_types.reshape([-1]) - else: - feed_dict_test[self.t_type] = np.tile(atom_types, [nframes, 1]).reshape( - [-1] - ) - feed_dict_test[self.t_coord] = np.reshape(coords, [-1]) - - if len(self.t_box.shape) == 1: - feed_dict_test[self.t_box] = np.reshape(cells, [-1]) - elif len(self.t_box.shape) == 2: - feed_dict_test[self.t_box] = cells - else: - raise RuntimeError - if self.has_efield: - feed_dict_test[self.t_efield] = np.reshape(efield, [-1]) - feed_dict_test[self.t_mesh] = mesh - if self.has_fparam: - feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1]) - if self.has_aparam: - feed_dict_test[self.t_aparam] = np.reshape(aparam, [-1]) - return feed_dict_test, imap, natoms_vec, ghost_map - - def _eval_inner( - self, - coords, - cells, - atom_types, - fparam=None, - aparam=None, - atomic=False, - efield=None, - mixed_type=False, - ): - natoms, nframes = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - feed_dict_test, imap, natoms_vec, ghost_map = self._prepare_feed_dict( - coords, cells, atom_types, fparam, aparam, efield, mixed_type=mixed_type - ) - - nloc = natoms_vec[0] - nall = natoms_vec[1] - - t_out = [self.t_energy, self.t_force, self.t_virial] - if atomic: - t_out += [self.t_ae, self.t_av] - - v_out = run_sess(self.sess, t_out, feed_dict=feed_dict_test) - energy = v_out[0] - force = v_out[1] - virial = v_out[2] - if atomic: - ae = v_out[3] - av = v_out[4] - - if self.has_spin: - ntypes_real = self.ntypes - self.ntypes_spin - natoms_real = sum( - [ - np.count_nonzero(np.array(atom_types) == ii) - for ii in range(ntypes_real) - ] - ) - else: - natoms_real = natoms - if ghost_map is not None: - # add the value of ghost atoms to real atoms - force = np.reshape(force, [nframes, -1, 3]) - np.add.at(force[0], ghost_map, force[0, nloc:]) - if atomic: - av = np.reshape(av, [nframes, -1, 9]) - np.add.at(av[0], ghost_map, av[0, nloc:]) - - # reverse map of the outputs - force = self.reverse_map(np.reshape(force, [nframes, -1, 3]), imap) - if atomic: - ae = self.reverse_map(np.reshape(ae, [nframes, -1, 1]), imap[:natoms_real]) - av = self.reverse_map(np.reshape(av, [nframes, -1, 9]), imap) - - energy = np.reshape(energy, [nframes, 1]) - force = np.reshape(force, [nframes, nall, 3]) - if nloc < nall: - force = force[:, :nloc, :] - virial = np.reshape(virial, [nframes, 9]) - if atomic: - ae = np.reshape(ae, [nframes, natoms_real, 1]) - av = np.reshape(av, [nframes, nall, 9]) - if nloc < nall: - av = av[:, :nloc, :] - return energy, force, virial, ae, av - else: - return energy, force, virial - - def eval_descriptor( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - mixed_type: bool = False, - ) -> np.array: - """Evaluate descriptors by using this DP. - - Parameters - ---------- - coords - The coordinates of atoms. - The array should be of size nframes x natoms x 3 - cells - The cell of the region. - If None then non-PBC is assumed, otherwise using PBC. - The array should be of size nframes x 9 - atom_types - The atom types - The list should contain natoms ints - fparam - The frame parameter. - The array can be of size : - - nframes x dim_fparam. - - dim_fparam. Then all frames are assumed to be provided with the same fparam. - aparam - The atomic parameter - The array can be of size : - - nframes x natoms x dim_aparam. - - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - - dim_aparam. Then all frames and atoms are provided with the same aparam. - efield - The external field on atoms. - The array should be of size nframes x natoms x 3 - mixed_type - Whether to perform the mixed_type mode. - If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), - in which frames in a system may have different natoms_vec(s), with the same nloc. - - Returns - ------- - descriptor - Descriptors. - """ - natoms, numb_test = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)( - coords, - cells, - atom_types, - fparam=fparam, - aparam=aparam, - efield=efield, - mixed_type=mixed_type, - ) - return descriptor - - def _eval_descriptor_inner( - self, - coords: np.ndarray, - cells: np.ndarray, - atom_types: List[int], - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, - mixed_type: bool = False, - ) -> np.array: - natoms, nframes = self._get_natoms_and_nframes( - coords, atom_types, mixed_type=mixed_type - ) - feed_dict_test, imap, natoms_vec, ghost_map = self._prepare_feed_dict( - coords, cells, atom_types, fparam, aparam, efield, mixed_type=mixed_type - ) - (descriptor,) = run_sess( - self.sess, [self.t_descriptor], feed_dict=feed_dict_test - ) - imap = imap[:natoms] - return self.reverse_map(np.reshape(descriptor, [nframes, natoms, -1]), imap) +__all__ = ["DeepPot"] diff --git a/deepmd/tf/infer/deep_tensor.py b/deepmd/tf/infer/deep_tensor.py index 9b064114b4..9e8acf8241 100644 --- a/deepmd/tf/infer/deep_tensor.py +++ b/deepmd/tf/infer/deep_tensor.py @@ -13,9 +13,7 @@ from deepmd.tf.common import ( make_default_mesh, ) -from deepmd.tf.infer.deep_eval import ( - DeepEval, -) +from deepmd.tf.infer.deep_eval import DeepEvalOld as DeepEval from deepmd.tf.utils.sess import ( run_sess, ) diff --git a/deepmd/tf/infer/deep_wfc.py b/deepmd/tf/infer/deep_wfc.py index aa0dff6f38..f7674bdde7 100644 --- a/deepmd/tf/infer/deep_wfc.py +++ b/deepmd/tf/infer/deep_wfc.py @@ -1,68 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - TYPE_CHECKING, - Optional, +from deepmd.infer.deep_wfc import ( + DeepWFC, ) -from deepmd.tf.infer.deep_tensor import ( - DeepTensor, -) - -if TYPE_CHECKING: - from pathlib import ( - Path, - ) - - -class DeepWFC(DeepTensor): - """Constructor. - - Parameters - ---------- - model_file : Path - The name of the frozen model file. - load_prefix: str - The prefix in the load computational graph - default_tf_graph : bool - If uses the default tf graph, otherwise build a new tf graph for evaluation - input_map : dict, optional - The input map for tf.import_graph_def. Only work with default tf graph - - Warnings - -------- - For developers: `DeepTensor` initializer must be called at the end after - `self.tensors` are modified because it uses the data in `self.tensors` dict. - Do not chanage the order! - """ - - def __init__( - self, - model_file: "Path", - load_prefix: str = "load", - default_tf_graph: bool = False, - input_map: Optional[dict] = None, - ) -> None: - # use this in favor of dict update to move attribute from class to - # instance namespace - self.tensors = dict( - { - # output tensor - "t_tensor": "o_wfc:0", - }, - **self.tensors, - ) - DeepTensor.__init__( - self, - model_file, - load_prefix=load_prefix, - default_tf_graph=default_tf_graph, - input_map=input_map, - ) - - def get_dim_fparam(self) -> int: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") - - def get_dim_aparam(self) -> int: - """Unsupported in this model.""" - raise NotImplementedError("This model type does not support this attribute") +__all__ = [ + "DeepWFC", +] From b428fc29461075247cf296626160961f99957a79 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Feb 2024 08:24:38 -0500 Subject: [PATCH 02/20] fix cycle import Signed-off-by: Jinzhe Zeng --- deepmd/tf/infer/deep_eval.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 1a917d4f64..93417a10b3 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -45,9 +45,6 @@ default_tf_session_config, tf, ) -from deepmd.tf.infer.data_modifier import ( - DipoleChargeModifier, -) from deepmd.tf.utils.batch_size import ( AutoBatchSize, ) @@ -136,6 +133,10 @@ def __init__( # looks ugly... if self.modifier_type == "dipole_charge": + from deepmd.tf.infer.data_modifier import ( + DipoleChargeModifier, + ) + t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0") t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0") t_sys_charge_map = self._get_tensor("modifier_attr/sys_charge_map:0") From eb0b7d2565145dec7ed465483cc46f14ed973463 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Feb 2024 15:57:56 -0500 Subject: [PATCH 03/20] fix errors Signed-off-by: Jinzhe Zeng --- deepmd/infer/backend.py | 1 + deepmd/infer/deep_eval.py | 6 +++--- deepmd/infer/deep_pot.py | 12 +++++++++++- deepmd/pt/infer/deep_eval.py | 2 +- deepmd/tf/infer/deep_eval.py | 12 +++++------- deepmd/tf/model/frozen.py | 10 +++++++++- 6 files changed, 30 insertions(+), 13 deletions(-) diff --git a/deepmd/infer/backend.py b/deepmd/infer/backend.py index 809e19466b..26eef22eb4 100644 --- a/deepmd/infer/backend.py +++ b/deepmd/infer/backend.py @@ -21,6 +21,7 @@ def detect_backend(filename: str) -> DPBackend: filename : str The model file name """ + filename = str(filename).lower() if filename.endswith(".pb"): return DPBackend.TensorFlow elif filename.endswith(".pth") or filename.endswith(".pt"): diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 92b449eb36..266b3b2d4d 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -239,7 +239,7 @@ def get_has_efield(self): return False @abstractmethod - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model.""" @@ -458,10 +458,10 @@ def _get_sel_natoms(self, atype) -> int: return np.sum(np.isin(atype, self.get_sel_type()).astype(int)) @property - def has_efield(self): + def has_efield(self) -> bool: """Check if the model has efield.""" return self.deep_eval.get_has_efield() - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model.""" return self.deep_eval.get_ntypes_spin() diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 3f207d03f8..d77fd1172a 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -148,7 +148,17 @@ def eval( virial = results["energy_derv_c_redu"].reshape(nframes, 9) if atomic: - atomic_energy = results["energy"].reshape(nframes, natoms, 1) + if self.get_ntypes_spin() > 0: + ntypes_real = self.get_ntypes() - self.get_ntypes_spin() + natoms_real = sum( + [ + np.count_nonzero(np.array(atom_types[0]) == ii) + for ii in range(ntypes_real) + ] + ) + else: + natoms_real = natoms + atomic_energy = results["energy"].reshape(nframes, natoms_real, 1) atomic_virial = results["energy_derv_c"].reshape(nframes, natoms, 9) return ( energy, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 31f61f68ca..8eb7e06868 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -112,7 +112,7 @@ def get_sel_type(self) -> List[int]: def get_numb_dos(self) -> int: """Get the number of DOS.""" - raise 0 + return 0 def get_has_efield(self): """Check if the model has efield.""" diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 93417a10b3..8cb9229905 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -208,7 +208,7 @@ def _init_tensors(self): "efield": "t_efield:0", "fparam": "t_fparam:0", "aparam": "t_aparam:0", - "ntypes_spin": "descrpt_attr/ntypes_spin:0", + "ntypes_spin": "spin_attr/ntypes_spin:0", # descriptor "descriptor": "o_descriptor:0", } @@ -435,7 +435,7 @@ def sort_input( """ natoms = atom_type.shape[1] if sel_atoms is not None: - selection = [False] * natoms + selection = np.array([False] * natoms, dtype=bool) for ii in sel_atoms: selection += atom_type[0] == ii sel_atom_type = atom_type[:, selection] @@ -628,7 +628,7 @@ def get_ntypes(self) -> int: """Get the number of atom types of this model.""" return self.ntypes - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model.""" return self.ntypes_spin @@ -770,7 +770,6 @@ def eval( output_dict["energy_redu"] += me.reshape(e.shape) output_dict["energy_deri_r"] += mf.reshape(f.shape) output_dict["energy_deri_c_redu"] += mv.reshape(v.shape) - output = tuple(output) return output_dict def _prepare_feed_dict( @@ -787,7 +786,7 @@ def _prepare_feed_dict( coords, atom_types, ) - atom_types = np.array(atom_types, dtype=int).reshape([-1, natoms]) + atom_types = np.array(atom_types, dtype=int).reshape([nframes, natoms]) coords = np.reshape(np.array(coords), [nframes, natoms * 3]) if cells is None: pbc = False @@ -942,7 +941,7 @@ def _eval_inner( ntypes_real = self.ntypes - self.ntypes_spin natoms_real = sum( [ - np.count_nonzero(np.array(atom_types) == ii) + np.count_nonzero(np.array(atom_types[0]) == ii) for ii in range(ntypes_real) ] ) @@ -977,7 +976,6 @@ def _eval_inner( odef_shape = self._get_output_shape( odef.name, nframes, natoms_real, odef.shape ) - # tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] v_out[ii] = self.reverse_map( np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real] ) diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 8732fec8f4..4151b1b0e4 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -7,6 +7,9 @@ Union, ) +from deepmd.infer.deep_pot import ( + DeepPot, +) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, MODEL_VERSION, @@ -40,7 +43,12 @@ def __init__(self, model_file: str, **kwargs): super().__init__(**kwargs) self.model_file = model_file self.model = DeepPotential(model_file) - self.model_type = self.model.model_type + if isinstance(self.model, DeepPot): + self.model_type = "ener" + else: + raise NotImplementedError( + "This model type has not been implemented. " "Contribution is welcome!" + ) def build( self, From 02b3a6d93ef4b835c1509e2fe3bd0d70278027a0 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Feb 2024 16:13:22 -0500 Subject: [PATCH 04/20] fix PT tests Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_eval.py | 4 ++-- deepmd/pt/infer/deep_eval.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 266b3b2d4d..c8cc1186e7 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -322,10 +322,10 @@ def get_dim_aparam(self) -> int: def _get_natoms_and_nframes( self, coords: np.ndarray, - atom_types: Union[List[int], np.ndarray], + atom_types: np.ndarray, mixed_type: bool = False, ) -> Tuple[int, int]: - if mixed_type: + if mixed_type or atom_types.ndim > 1: natoms = len(atom_types[0]) else: natoms = len(atom_types) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 8eb7e06868..41bf8a4bce 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -126,7 +126,8 @@ def get_ntypes_spin(self): "energy": "atom_energy", "energy_redu": "energy", "energy_derv_r": "force", - "energy_derv_c": "atom_virial", + # not same as TF... + "energy_derv_c": "atomic_virial", "energy_derv_c_redu": "virial", } From 4fa87547a6635810c7327b93fedd3eda795f07f4 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Feb 2024 16:44:48 -0500 Subject: [PATCH 05/20] fix FrozenModel Signed-off-by: Jinzhe Zeng --- deepmd/tf/model/frozen.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 4151b1b0e4..f06ae954d1 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -130,14 +130,26 @@ def build( ) if self.model_type == "ener": return { - "energy": tf.identity(self.model.t_energy, name="o_energy" + suffix), - "force": tf.identity(self.model.t_force, name="o_force" + suffix), - "virial": tf.identity(self.model.t_virial, name="o_virial" + suffix), + # must visit the backend class + "energy": tf.identity( + self.model.deep_eval.output_tensors["energy_redu"], + name="o_energy" + suffix, + ), + "force": tf.identity( + self.model.deep_eval.output_tensors["energy_derv_r"], + name="o_force" + suffix, + ), + "virial": tf.identity( + self.model.deep_eval.output_tensors["energy_derv_c_redu"], + name="o_virial" + suffix, + ), "atom_ener": tf.identity( - self.model.t_ae, name="o_atom_energy" + suffix + self.model.deep_eval.output_tensors["energy"], + name="o_atom_energy" + suffix, ), "atom_virial": tf.identity( - self.model.t_av, name="o_atom_virial" + suffix + self.model.deep_eval.output_tensors["energy_derv_c"], + name="o_atom_virial" + suffix, ), "coord": coord_, "atype": atype_, From 6026ab73de83b3eb5b1cb24d7e986cbfd9cf69fb Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Feb 2024 17:27:05 -0500 Subject: [PATCH 06/20] add docs Signed-off-by: Jinzhe Zeng --- deepmd/__init__.py | 21 ++++++++++++++++++--- deepmd/infer/deep_dipole.py | 18 ++++++++++++++++++ deepmd/infer/deep_dos.py | 6 +++++- deepmd/infer/deep_polar.py | 18 ++++++++++++++++++ deepmd/infer/deep_pot.py | 4 ++++ deepmd/infer/deep_tensor.py | 18 ++++++++++++++++++ deepmd/infer/deep_wfc.py | 18 ++++++++++++++++++ deepmd/pt/infer/deep_eval.py | 20 ++++++++++++++++++++ deepmd/tf/infer/__init__.py | 31 +++++++++++++++++++++++++++++-- deepmd/tf/infer/deep_eval.py | 12 +++++++++--- 10 files changed, 157 insertions(+), 9 deletions(-) diff --git a/deepmd/__init__.py b/deepmd/__init__.py index 4df06e2b9e..5664c3edc6 100644 --- a/deepmd/__init__.py +++ b/deepmd/__init__.py @@ -16,11 +16,26 @@ def DeepPotential(*args, **kwargs): - from deepmd.infer.deep_eval import ( - DeepEval, + """Factory function that forwards to DeepEval (for compatbility + and performance). + + Parameters + ---------- + *args + positional arguments + **kwargs + keyword arguments + + Returns + ------- + DeepEval + potentials + """ + from deepmd.infer import ( + DeepPotential, ) - return DeepEval(*args, **kwargs) + return DeepPotential(*args, **kwargs) __all__ = [ diff --git a/deepmd/infer/deep_dipole.py b/deepmd/infer/deep_dipole.py index fe51202a43..b443b54417 100644 --- a/deepmd/infer/deep_dipole.py +++ b/deepmd/infer/deep_dipole.py @@ -5,6 +5,24 @@ class DeepDipole(DeepTensor): + """Deep dipole model. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + @property def output_tensor_name(self) -> str: return "dipole" diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index 7ece365d95..cdc62bd530 100644 --- a/deepmd/infer/deep_dos.py +++ b/deepmd/infer/deep_dos.py @@ -20,18 +20,22 @@ class DeepDOS(DeepEval): - """Potential energy model. + """Deep density of states model. Parameters ---------- model_file : Path The name of the frozen model file. + *args : list + Positional arguments. auto_batch_size : bool or int or AutoBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. """ @property diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 96ad70dcc0..2373252a8b 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -11,6 +11,24 @@ class DeepPolar(DeepTensor): + """Deep polar model. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + @property def output_tensor_name(self) -> str: return "polar" diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index d77fd1172a..e25e2142a0 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -26,12 +26,16 @@ class DeepPot(DeepEval): ---------- model_file : Path The name of the frozen model file. + *args : list + Positional arguments. auto_batch_size : bool or int or AutoBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. Examples -------- diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 64c2fc1130..1d0a97763c 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -20,6 +20,24 @@ class DeepTensor(DeepEval): + """Deep Tensor Model. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + def eval( self, coords: np.ndarray, diff --git a/deepmd/infer/deep_wfc.py b/deepmd/infer/deep_wfc.py index d1c7a92ad2..deed938e04 100644 --- a/deepmd/infer/deep_wfc.py +++ b/deepmd/infer/deep_wfc.py @@ -5,6 +5,24 @@ class DeepWFC(DeepTensor): + """Deep WFC model. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + @property def output_tensor_name(self) -> str: return "wfc" diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 41bf8a4bce..b304c0883c 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -45,6 +45,26 @@ class DeepEval(DeepEvalBase): + """PyTorch backend implementaion of DeepEval. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + output_def : ModelOutputDef + The output definition of the model. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutomaticBatchSize, default: False + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + def __init__( self, model_file: str, diff --git a/deepmd/tf/infer/__init__.py b/deepmd/tf/infer/__init__.py index f0cd86a885..e0168198ff 100644 --- a/deepmd/tf/infer/__init__.py +++ b/deepmd/tf/infer/__init__.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Submodule containing all the implemented potentials.""" -from deepmd import ( - DeepPotential, +from typing import ( + TYPE_CHECKING, ) from .data_modifier import ( @@ -34,6 +34,11 @@ calc_model_devi, ) +if TYPE_CHECKING: + from deepmd.infer.deep_eval import ( + DeepEval, + ) + __all__ = [ "DeepPotential", "DeepDipole", @@ -47,3 +52,25 @@ "EwaldRecp", "calc_model_devi", ] + + +def DeepPotential(*args, **kwargs) -> "DeepEval": + """Factory function that forwards to DeepEval (for compatbility). + + Parameters + ---------- + *args + positional arguments + **kwargs + keyword arguments + + Returns + ------- + DeepEval + potentials + """ + from deepmd.infer.deep_eval import ( + DeepEval, + ) + + return DeepEval(*args, **kwargs) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 8cb9229905..0113dedd7c 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -59,12 +59,16 @@ class DeepEval(DeepEvalBase): - """Common methods for DeepPot, DeepWFC, DeepPolar, ... + """TensorFlow backend implementation for DeepEval. Parameters ---------- model_file : Path The name of the frozen model file. + output_def : ModelOutputDef + The output definition of the model. + *args : list + Positional arguments. load_prefix: str The prefix in the load computational graph default_tf_graph : bool @@ -77,19 +81,21 @@ class DeepEval(DeepEvalBase): neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. """ - load_prefix: str # set by subclass - def __init__( self, model_file: "Path", output_def: ModelOutputDef, + *args: list, load_prefix: str = "load", default_tf_graph: bool = False, auto_batch_size: Union[bool, int, AutoBatchSize] = False, input_map: Optional[dict] = None, neighbor_list=None, + **kwargs: dict, ): self.graph = self._load_graph( model_file, From c262ef6da7b5ba9c2a002567b2c5bc87583204d9 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 2 Feb 2024 01:11:59 -0500 Subject: [PATCH 07/20] rename .model_format to .dpmodel Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_dos.py | 2 +- deepmd/infer/deep_eval.py | 2 +- deepmd/infer/deep_pot.py | 2 +- deepmd/infer/deep_tensor.py | 8 ++++---- deepmd/pt/infer/deep_eval.py | 6 +++--- deepmd/tf/infer/deep_eval.py | 6 +++--- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index cdc62bd530..69741fc263 100644 --- a/deepmd/infer/deep_dos.py +++ b/deepmd/infer/deep_dos.py @@ -8,7 +8,7 @@ import numpy as np -from deepmd.model_format.output_def import ( +from deepmd.dpmodel.output_def import ( FittingOutputDef, ModelOutputDef, OutputVariableDef, diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index c8cc1186e7..483460de10 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -15,7 +15,7 @@ import numpy as np -from deepmd.model_format.output_def import ( +from deepmd.dpmodel.output_def import ( FittingOutputDef, ModelOutputDef, ) diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index e25e2142a0..c3c30b5abf 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -8,7 +8,7 @@ import numpy as np -from deepmd.model_format.output_def import ( +from deepmd.dpmodel.output_def import ( FittingOutputDef, ModelOutputDef, OutputVariableDef, diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 1d0a97763c..0f6ac8161a 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -9,14 +9,14 @@ import numpy as np -from deepmd.infer.deep_eval import ( - DeepEval, -) -from deepmd.model_format.output_def import ( +from deepmd.dpmodel.output_def import ( FittingOutputDef, ModelOutputDef, OutputVariableDef, ) +from deepmd.infer.deep_eval import ( + DeepEval, +) class DeepTensor(DeepEval): diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index b304c0883c..14e53bd124 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -14,15 +14,15 @@ import numpy as np import torch +from deepmd.dpmodel.output_def import ( + ModelOutputDef, +) from deepmd.infer.deep_eval import ( DeepEvalBase, ) from deepmd.infer.deep_pot import ( DeepPot, ) -from deepmd.model_format.output_def import ( - ModelOutputDef, -) from deepmd.pt.model.model import ( get_model, ) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 0113dedd7c..e179e16e46 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -18,6 +18,9 @@ from deepmd.common import ( make_default_mesh, ) +from deepmd.dpmodel.output_def import ( + ModelOutputDef, +) from deepmd.infer.deep_dipole import ( DeepDipole, ) @@ -37,9 +40,6 @@ from deepmd.infer.deep_wfc import ( DeepWFC, ) -from deepmd.model_format.output_def import ( - ModelOutputDef, -) from deepmd.tf.env import ( MODEL_VERSION, default_tf_session_config, From 4fe9511ea5b0e3d66987746d0cff1e420b9644bf Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 2 Feb 2024 16:06:58 -0500 Subject: [PATCH 08/20] rename DeepEvalBase to DeepEvalBackend as the meaning of base is unclear Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_eval.py | 8 ++++---- deepmd/pt/infer/deep_eval.py | 4 ++-- deepmd/tf/infer/deep_eval.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 483460de10..8c7d6d5295 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -32,7 +32,7 @@ import ase.neighborlist -class DeepEvalBase(ABC): +class DeepEvalBackend(ABC): """Low-level Deep Evaluator interface. Backends should inherbit implement this interface. High-level interface @@ -67,7 +67,7 @@ def __init__( pass def __new__(cls, model_file: str, *args, **kwargs): - if cls is DeepEvalBase: + if cls is DeepEvalBackend: backend = detect_backend(model_file) if backend == DPBackend.TensorFlow: from deepmd.tf.infer.deep_eval import DeepEval as DeepEvalTF @@ -268,7 +268,7 @@ class DeepEval(ABC): def __new__(cls, model_file: str, *args, **kwargs): if cls is DeepEval: - deep_eval = DeepEvalBase( + deep_eval = DeepEvalBackend( model_file, ModelOutputDef(FittingOutputDef([])), *args, @@ -285,7 +285,7 @@ def __init__( neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Dict[str, Any], ) -> None: - self.deep_eval = DeepEvalBase( + self.deep_eval = DeepEvalBackend( model_file, self.output_def, *args, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 14e53bd124..bbadc0a3cb 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -18,7 +18,7 @@ ModelOutputDef, ) from deepmd.infer.deep_eval import ( - DeepEvalBase, + DeepEvalBackend, ) from deepmd.infer.deep_pot import ( DeepPot, @@ -44,7 +44,7 @@ import ase.neighborlist -class DeepEval(DeepEvalBase): +class DeepEval(DeepEvalBackend): """PyTorch backend implementaion of DeepEval. Parameters diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index e179e16e46..fbec943161 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -28,7 +28,7 @@ DeepDOS, ) from deepmd.infer.deep_eval import ( - DeepEvalBase, + DeepEvalBackend, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -58,7 +58,7 @@ ) -class DeepEval(DeepEvalBase): +class DeepEval(DeepEvalBackend): """TensorFlow backend implementation for DeepEval. Parameters From 5b16abf651e58d2adccee0ce3acd2fcc877317ed Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 2 Feb 2024 16:09:28 -0500 Subject: [PATCH 09/20] ckean OutputVariableDef Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_pot.py | 5 ----- deepmd/infer/deep_tensor.py | 5 ----- 2 files changed, 10 deletions(-) diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index c3c30b5abf..6cbc6c058c 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -63,11 +63,6 @@ def output_def(self) -> ModelOutputDef: differentiable=True, atomic=True, ), - OutputVariableDef( - # ugly... - "energy_derv_c_redu", - shape=[1, 3, 3], - ), ] ) ) diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 0f6ac8161a..b9346c4633 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -230,11 +230,6 @@ def output_def(self) -> ModelOutputDef: differentiable=True, atomic=True, ), - OutputVariableDef( - # ugly... - f"{self.output_tensor_name}_derv_c_redu", - shape=[-1, 3, 3], - ), ] ) ) From 51b5e7b75a7cdf7d4c4220a01eea16aa1702843e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 3 Feb 2024 01:42:32 -0500 Subject: [PATCH 10/20] make atom_types to be int32 Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 8c7d6d5295..be7ac20e39 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -441,7 +441,7 @@ def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type) coords = np.array(coords) if cells is not None: cells = np.array(cells) - atom_types = np.array(atom_types) + atom_types = np.array(atom_types, dtype=np.int32) if fparam is not None: fparam = np.array(fparam) if aparam is not None: From 3516792eb0b8eee20cf08d9906aeac9f8ea76637 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 3 Feb 2024 01:58:10 -0500 Subject: [PATCH 11/20] add docs Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_eval.py | 13 +++++++++++-- deepmd/pt/infer/deep_eval.py | 10 ++++++++++ deepmd/tf/infer/deep_eval.py | 9 ++++++--- 3 files changed, 27 insertions(+), 5 deletions(-) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index be7ac20e39..4af2031e14 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -217,8 +217,17 @@ def eval_typeebd(self) -> np.ndarray: """ raise NotImplementedError - def _check_distinguished_types(self, atom_types: np.ndarray) -> bool: - """Check if atom types of each frame.""" + def _check_mixed_types(self, atom_types: np.ndarray) -> bool: + """Check if atom types of all frames are the same. + + Traditional descriptors like se_e2_a requires all the frames to + have the same atom types. + + Parameters + ---------- + atom_types : np.ndarray + The atom types of all frames, in shape nframes * natoms. + """ return np.all(np.equal(atom_types, atom_types[0])) @property diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index bbadc0a3cb..c9b8f1b874 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -210,6 +210,9 @@ def eval( out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, atomic ) + # when atomic is True, all output_def are requested + # when atomic is False, only energy, force, and virial are requested + # the condition here should be the same as that in _eval_method return dict( zip( [ @@ -306,6 +309,8 @@ def _eval_model( results = [] for ii, odef in enumerate(self.output_def.var_defs.values()): + # when atomic is True, all output_def are requested + # when atomic is False, only energy, force, and virial are requested if not atomic and "_redu" not in odef.name and "_derv_r" not in odef.name: continue pt_name = self._OUTDEF_DP2PT[odef.name] @@ -318,15 +323,20 @@ def _eval_model( def _get_output_shape(self, name, nframes, natoms, shape): if "_redu" in name: if "_derv_c" in name: + # virial return [nframes, *shape[:-2], 9] else: + # energy return [nframes, *shape, 1] else: if "_derv_c" in name: + # atom_virial return [nframes, *shape[:-2], natoms, 9] elif "_derv_r" in name: + # force return [nframes, *shape[:-1], natoms, 3] else: + # atom_energy, atom_tensor # Something wrong here? # return [nframes, *shape, natoms, 1] return [nframes, natoms, *shape, 1] diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index fbec943161..9b2befdfa1 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -857,9 +857,7 @@ def _prepare_feed_dict( if self.neighbor_list is None: natoms_vec = self.make_natoms_vec(atom_types) assert natoms_vec[0] == natoms - mesh = make_default_mesh( - pbc, not self._check_distinguished_types(atom_types) - ) + mesh = make_default_mesh(pbc, not self._check_mixed_types(atom_types)) ghost_map = None else: if nframes > 1: @@ -994,15 +992,20 @@ def _eval_inner( def _get_output_shape(self, name, nframes, natoms, shape): if "_redu" in name: if "_derv_c" in name: + # virial return [nframes, *shape[:-2], 9] else: + # energy return [nframes, *shape, 1] else: if "_derv_c" in name: + # atom_virial return [nframes, *shape[:-2], natoms, 9] elif "_derv_r" in name: + # force return [nframes, *shape[:-1], natoms, 3] else: + # atom_energy, atom_tensor # Something wrong here? # return [nframes, *shape, natoms, 1] return [nframes, natoms, *shape, 1] From 3bc7cb5e11fb913b03cb59385ccfb9d98c9cdad5 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 4 Feb 2024 00:31:47 -0500 Subject: [PATCH 12/20] pass list[OutputVariableDef] instead of bool Signed-off-by: Jinzhe Zeng --- deepmd/pt/infer/deep_eval.py | 63 ++++++++++++++++++++++-------------- 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index c9b8f1b874..111a2406c1 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -16,6 +16,7 @@ from deepmd.dpmodel.output_def import ( ModelOutputDef, + OutputVariableDef, ) from deepmd.infer.deep_eval import ( DeepEvalBackend, @@ -43,6 +44,8 @@ if TYPE_CHECKING: import ase.neighborlist + from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper + class DeepEval(DeepEvalBackend): """PyTorch backend implementaion of DeepEval. @@ -122,7 +125,7 @@ def get_dim_aparam(self) -> int: return 0 @property - def model_type(self) -> "DeepEval": + def model_type(self) -> "DeepEvalWrapper": """The the evaluator of the model type.""" return DeepPot @@ -155,7 +158,7 @@ def eval( self, coords: np.ndarray, cells: np.ndarray, - atom_types: List[int], + atom_types: np.ndarray, atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -207,23 +210,43 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + request_defs = self._get_request_defs(atomic) out = self._eval_func(self._eval_model, numb_test, natoms)( - coords, cells, atom_types, atomic + coords, cells, atom_types, request_defs ) - # when atomic is True, all output_def are requested - # when atomic is False, only energy, force, and virial are requested - # the condition here should be the same as that in _eval_method return dict( zip( - [ - x.name - for x in self.output_def.var_defs.values() - if atomic or "_redu" in x.name or "_derv_r" in x.name - ], + [x.name for x in request_defs], out, ) ) + def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: + """Get the requested output definitions. + + When atomic is True, all output_def are requested. + When atomic is False, only energy (tensor), force, and virial + are requested. + + Parameters + ---------- + atomic : bool + Whether to request the atomic output. + + Returns + ------- + list[OutputVariableDef] + The requested output definitions. + """ + if atomic: + return list(self.output_def.var_defs.values()) + else: + return [ + x + for x in self.output_def.var_defs.values() + if x.name.endswith("_redu") or x.name.endswith("_derv_r") + ] + def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: """Wrapper method with auto batch size. @@ -255,7 +278,7 @@ def eval_func(*args, **kwargs): def _get_natoms_and_nframes( self, coords: np.ndarray, - atom_types: Union[List[int], np.ndarray], + atom_types: np.ndarray, mixed_type: bool = False, ) -> Tuple[int, int]: if mixed_type: @@ -274,14 +297,9 @@ def _eval_model( coords: np.ndarray, cells: Optional[np.ndarray], atom_types: np.ndarray, - atomic: bool = False, + request_defs: List[OutputVariableDef], ): model = self.dp.to(DEVICE) - energy_out = None - atomic_energy_out = None - force_out = None - virial_out = None - atomic_virial_out = None nframes = coords.shape[0] if len(atom_types.shape) == 1: @@ -301,18 +319,15 @@ def _eval_model( else: box_input = None + do_atomic_virial = any(x.name.endswith("_derv_c") for x in request_defs) batch_output = model( - coord_input, type_input, box=box_input, do_atomic_virial=atomic + coord_input, type_input, box=box_input, do_atomic_virial=do_atomic_virial ) if isinstance(batch_output, tuple): batch_output = batch_output[0] results = [] - for ii, odef in enumerate(self.output_def.var_defs.values()): - # when atomic is True, all output_def are requested - # when atomic is False, only energy, force, and virial are requested - if not atomic and "_redu" not in odef.name and "_derv_r" not in odef.name: - continue + for odef in request_defs: pt_name = self._OUTDEF_DP2PT[odef.name] if pt_name in batch_output: shape = self._get_output_shape(odef.name, nframes, natoms, odef.shape) From acfd367d84b392f82aa8c7cd8e364d9a65452c09 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 4 Feb 2024 00:40:48 -0500 Subject: [PATCH 13/20] fix py38 compatibility Signed-off-by: Jinzhe Zeng --- deepmd/pt/infer/deep_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 111a2406c1..a422f12cd2 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -221,7 +221,7 @@ def eval( ) ) - def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: + def _get_request_defs(self, atomic: bool) -> List[OutputVariableDef]: """Get the requested output definitions. When atomic is True, all output_def are requested. From be4e42faf443b6a31c738208817fa70ca36cdedd Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 4 Feb 2024 20:49:51 -0500 Subject: [PATCH 14/20] use consistent output name for different backends Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_eval.py | 21 +++++++++++++++++++++ deepmd/pt/infer/deep_eval.py | 12 +----------- deepmd/tf/infer/deep_eval.py | 23 +---------------------- 3 files changed, 23 insertions(+), 33 deletions(-) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 4af2031e14..663c05eb5f 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -6,6 +6,7 @@ from typing import ( TYPE_CHECKING, Any, + ClassVar, Dict, List, Optional, @@ -54,6 +55,26 @@ class DeepEvalBackend(ABC): Keyword arguments. """ + _OUTDEF_DP2BACKEND: ClassVar[dict] = { + "energy": "atom_energy", + "energy_redu": "energy", + "energy_derv_r": "force", + "energy_derv_c": "atom_virial", + "energy_derv_c_redu": "virial", + "polar": "polar", + "polar_redu": "global_polar", + "polar_derv_r": "force", + "polar_derv_c": "atom_virial", + "polar_derv_c_redu": "virial", + "dipole": "dipole", + "dipole_redu": "global_dipole", + "dipole_derv_r": "force", + "dipole_derv_c": "atom_virial", + "dipole_derv_c_redu": "virial", + "dos": "atom_dos", + "dos_redu": "dos", + } + @abstractmethod def __init__( self, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 70673907a9..e3ef411b1d 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -3,7 +3,6 @@ TYPE_CHECKING, Any, Callable, - ClassVar, Dict, List, Optional, @@ -145,15 +144,6 @@ def get_ntypes_spin(self): """Get the number of spin atom types of this model.""" return 0 - _OUTDEF_DP2PT: ClassVar[dict] = { - "energy": "atom_energy", - "energy_redu": "energy", - "energy_derv_r": "force", - # not same as TF... - "energy_derv_c": "atomic_virial", - "energy_derv_c_redu": "virial", - } - def eval( self, coords: np.ndarray, @@ -328,7 +318,7 @@ def _eval_model( results = [] for odef in request_defs: - pt_name = self._OUTDEF_DP2PT[odef.name] + pt_name = self._OUTDEF_DP2BACKEND[odef.name] if pt_name in batch_output: shape = self._get_output_shape(odef.name, nframes, natoms, odef.shape) out = batch_output[pt_name].reshape(shape).detach().cpu().numpy() diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 9b2befdfa1..d5304df1ed 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -5,7 +5,6 @@ from typing import ( TYPE_CHECKING, Callable, - ClassVar, Dict, List, Optional, @@ -169,26 +168,6 @@ def __init__( ewald_beta=ewald_beta, ) - _OUTDEF_DP2TF: ClassVar[dict] = { - "energy": "atom_energy", - "energy_redu": "energy", - "energy_derv_r": "force", - "energy_derv_c": "atom_virial", - "energy_derv_c_redu": "virial", - "polar": "polar", - "polar_redu": "global_polar", - "polar_derv_r": "force", - "polar_derv_c": "atom_virial", - "polar_derv_c_redu": "virial", - "dipole": "dipole", - "dipole_redu": "global_dipole", - "dipole_derv_r": "force", - "dipole_derv_c": "atom_virial", - "dipole_derv_c_redu": "virial", - "dos": "atom_dos", - "dos_redu": "dos", - } - def _init_tensors(self): tensor_names = { # descrpt attrs @@ -221,7 +200,7 @@ def _init_tensors(self): # output tensors output_tensor_names = {} for vv in self.output_def.var_defs: - output_tensor_names[vv] = f"o_{self._OUTDEF_DP2TF[vv]}:0" + output_tensor_names[vv] = f"o_{self._OUTDEF_DP2BACKEND[vv]}:0" self.tensors = {} for tensor_key, tensor_name in tensor_names.items(): From 199147cfe2b718b2d2411fb1fe29bd8ca384f045 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 4 Feb 2024 21:19:12 -0500 Subject: [PATCH 15/20] remove the commented out code Signed-off-by: Jinzhe Zeng --- deepmd/tf/infer/deep_eval.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index d5304df1ed..f7eba1a485 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -456,8 +456,6 @@ def reverse_map(vec: np.ndarray, imap: List[int]) -> np.ndarray: Reverse mapped vector. """ ret = np.zeros(vec.shape) - # for idx,ii in enumerate(imap) : - # ret[:,ii,:] = vec[:,idx,:] ret[:, imap, :] = vec return ret From 8686fa19c356a5044926b2979c2bb5f20753995c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 5 Feb 2024 23:28:58 -0500 Subject: [PATCH 16/20] use OutputVariableCategory to check odef Signed-off-by: Jinzhe Zeng --- deepmd/pt/infer/deep_eval.py | 52 ++++++++++++--------- deepmd/tf/infer/deep_eval.py | 91 +++++++++++++++++++----------------- 2 files changed, 77 insertions(+), 66 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index e3ef411b1d..3241ea9596 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -15,6 +15,7 @@ from deepmd.dpmodel.output_def import ( ModelOutputDef, + OutputVariableCategory, OutputVariableDef, ) from deepmd.infer.deep_eval import ( @@ -234,7 +235,12 @@ def _get_request_defs(self, atomic: bool) -> List[OutputVariableDef]: return [ x for x in self.output_def.var_defs.values() - if x.name.endswith("_redu") or x.name.endswith("_derv_r") + if x.category + in ( + OutputVariableCategory.REDU, + OutputVariableCategory.DERV_R, + OutputVariableCategory.DERV_C_REDU, + ) ] def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: @@ -309,7 +315,9 @@ def _eval_model( else: box_input = None - do_atomic_virial = any(x.name.endswith("_derv_c") for x in request_defs) + do_atomic_virial = any( + x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs + ) batch_output = model( coord_input, type_input, box=box_input, do_atomic_virial=do_atomic_virial ) @@ -320,31 +328,31 @@ def _eval_model( for odef in request_defs: pt_name = self._OUTDEF_DP2BACKEND[odef.name] if pt_name in batch_output: - shape = self._get_output_shape(odef.name, nframes, natoms, odef.shape) + shape = self._get_output_shape(odef, nframes, natoms) out = batch_output[pt_name].reshape(shape).detach().cpu().numpy() results.append(out) return tuple(results) - def _get_output_shape(self, name, nframes, natoms, shape): - if "_redu" in name: - if "_derv_c" in name: - # virial - return [nframes, *shape[:-2], 9] - else: - # energy - return [nframes, *shape, 1] + def _get_output_shape(self, odef, nframes, natoms): + if odef.category == OutputVariableCategory.DERV_C_REDU: + # virial + return [nframes, *odef.shape[:-2], 9] + elif odef.category == OutputVariableCategory.REDU: + # energy + return [nframes, *odef.shape, 1] + elif odef.category == OutputVariableCategory.DERV_C: + # atom_virial + return [nframes, *odef.shape[:-2], natoms, 9] + elif odef.category == OutputVariableCategory.DERV_R: + # force + return [nframes, *odef.shape[:-1], natoms, 3] + elif odef.category == OutputVariableCategory.OUT: + # atom_energy, atom_tensor + # Something wrong here? + # return [nframes, *shape, natoms, 1] + return [nframes, natoms, *odef.shape, 1] else: - if "_derv_c" in name: - # atom_virial - return [nframes, *shape[:-2], natoms, 9] - elif "_derv_r" in name: - # force - return [nframes, *shape[:-1], natoms, 3] - else: - # atom_energy, atom_tensor - # Something wrong here? - # return [nframes, *shape, natoms, 1] - return [nframes, natoms, *shape, 1] + raise RuntimeError("unknown category") # For tests only diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index f7eba1a485..c5a8835cc7 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -19,6 +19,7 @@ ) from deepmd.dpmodel.output_def import ( ModelOutputDef, + OutputVariableCategory, ) from deepmd.infer.deep_dipole import ( DeepDipole, @@ -932,60 +933,62 @@ def _eval_inner( # add the value of ghost atoms to real atoms for ii, odef in enumerate(self.output_def.var_defs.values()): # when the shape is nall - if "_redu" not in odef.name and "_derv" in odef.name: - odef_shape = self._get_output_shape( - odef.name, nframes, nall, odef.shape - ) + if odef.category in ( + OutputVariableCategory.DERV_R, + OutputVariableCategory.DERV_C, + ): + odef_shape = self._get_output_shape(odef, nframes, nall) tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] v_out[ii] = np.reshape(v_out[ii], tmp_shape) for jj in range(v_out[ii].shape[0]): np.add.at(v_out[ii][jj], ghost_map, v_out[ii][jj, nloc:]) for ii, odef in enumerate(self.output_def.var_defs.values()): - if "_redu" not in odef.name: - if "_derv" in odef.name: - odef_shape = self._get_output_shape( - odef.name, nframes, nall, odef.shape - ) - tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] - # reverse map of the outputs - v_out[ii] = self.reverse_map(np.reshape(v_out[ii], tmp_shape), imap) - v_out[ii] = np.reshape(v_out[ii], odef_shape) - if nloc < nall: - v_out[ii] = v_out[ii][:, :, :nloc] - else: - odef_shape = self._get_output_shape( - odef.name, nframes, natoms_real, odef.shape - ) - v_out[ii] = self.reverse_map( - np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real] - ) - v_out[ii] = np.reshape(v_out[ii], odef_shape) - else: - odef_shape = self._get_output_shape(odef.name, nframes, 0, odef.shape) + odef_shape = self._get_output_shape(odef, nframes, nall) + if odef.category in ( + OutputVariableCategory.DERV_R, + OutputVariableCategory.DERV_C, + ): + tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] + # reverse map of the outputs + v_out[ii] = self.reverse_map(np.reshape(v_out[ii], tmp_shape), imap) v_out[ii] = np.reshape(v_out[ii], odef_shape) + if nloc < nall: + v_out[ii] = v_out[ii][:, :, :nloc] + elif odef.category == OutputVariableCategory.OUT: + v_out[ii] = self.reverse_map( + np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real] + ) + v_out[ii] = np.reshape(v_out[ii], odef_shape) + elif odef.category in ( + OutputVariableCategory.REDU, + OutputVariableCategory.DERV_C_REDU, + ): + v_out[ii] = np.reshape(v_out[ii], odef_shape) + else: + raise RuntimeError("unknown category") return tuple(v_out) - def _get_output_shape(self, name, nframes, natoms, shape): - if "_redu" in name: - if "_derv_c" in name: - # virial - return [nframes, *shape[:-2], 9] - else: - # energy - return [nframes, *shape, 1] + def _get_output_shape(self, odef, nframes, natoms): + if odef.category == OutputVariableCategory.DERV_C_REDU: + # virial + return [nframes, *odef.shape[:-2], 9] + elif odef.category == OutputVariableCategory.REDU: + # energy + return [nframes, *odef.shape, 1] + elif odef.category == OutputVariableCategory.DERV_C: + # atom_virial + return [nframes, *odef.shape[:-2], natoms, 9] + elif odef.category == OutputVariableCategory.DERV_R: + # force + return [nframes, *odef.shape[:-1], natoms, 3] + elif odef.category == OutputVariableCategory.OUT: + # atom_energy, atom_tensor + # Something wrong here? + # return [nframes, *shape, natoms, 1] + return [nframes, natoms, *odef.shape, 1] else: - if "_derv_c" in name: - # atom_virial - return [nframes, *shape[:-2], natoms, 9] - elif "_derv_r" in name: - # force - return [nframes, *shape[:-1], natoms, 3] - else: - # atom_energy, atom_tensor - # Something wrong here? - # return [nframes, *shape, natoms, 1] - return [nframes, natoms, *shape, 1] + raise RuntimeError("unknown category") def eval_descriptor( self, From c0a4251913ab1055750461649fb205f742af4676 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 5 Feb 2024 23:59:45 -0500 Subject: [PATCH 17/20] fix typo Signed-off-by: Jinzhe Zeng --- deepmd/tf/infer/deep_eval.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index c5a8835cc7..4bcb70349b 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -944,11 +944,11 @@ def _eval_inner( np.add.at(v_out[ii][jj], ghost_map, v_out[ii][jj, nloc:]) for ii, odef in enumerate(self.output_def.var_defs.values()): - odef_shape = self._get_output_shape(odef, nframes, nall) if odef.category in ( OutputVariableCategory.DERV_R, OutputVariableCategory.DERV_C, ): + odef_shape = self._get_output_shape(odef, nframes, nall) tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]] # reverse map of the outputs v_out[ii] = self.reverse_map(np.reshape(v_out[ii], tmp_shape), imap) @@ -956,6 +956,7 @@ def _eval_inner( if nloc < nall: v_out[ii] = v_out[ii][:, :, :nloc] elif odef.category == OutputVariableCategory.OUT: + odef_shape = self._get_output_shape(odef, nframes, natoms_real) v_out[ii] = self.reverse_map( np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real] ) @@ -964,6 +965,7 @@ def _eval_inner( OutputVariableCategory.REDU, OutputVariableCategory.DERV_C_REDU, ): + odef_shape = self._get_output_shape(odef, nframes, 0) v_out[ii] = np.reshape(v_out[ii], odef_shape) else: raise RuntimeError("unknown category") From d7bf11331a5a65e552822cef13392bc7ece83a67 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Feb 2024 03:40:30 -0500 Subject: [PATCH 18/20] improve docs and type hints Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_dos.py | 6 ++++-- deepmd/infer/deep_eval.py | 18 ++++++++++++++---- deepmd/infer/deep_polar.py | 6 ++++-- deepmd/infer/deep_pot.py | 6 ++++-- deepmd/infer/deep_tensor.py | 6 ++++-- deepmd/pt/infer/deep_eval.py | 7 ++++++- deepmd/tf/infer/deep_eval.py | 17 ++++++++++++----- 7 files changed, 48 insertions(+), 18 deletions(-) diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index 69741fc263..d95d2a119f 100644 --- a/deepmd/infer/deep_dos.py +++ b/deepmd/infer/deep_dos.py @@ -2,8 +2,10 @@ from typing import ( Any, Dict, + List, Optional, Tuple, + Union, ) import numpy as np @@ -58,7 +60,7 @@ def eval( self, coords: np.ndarray, cells: Optional[np.ndarray], - atom_types: np.ndarray, + atom_types: Union[List[int], np.ndarray], atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -75,7 +77,7 @@ def eval( cells : np.ndarray The cell vectors of the system, in shape (nframes, 9). If the system is not periodic, set it to None. - atom_types : List[int] + atom_types : List[int] or np.ndarray The types of the atoms. If mixed_type is False, the shape is (natoms,); otherwise, the shape is (nframes, natoms). atomic : bool, optional diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 663c05eb5f..b1d17afc8b 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -107,7 +107,7 @@ def eval( self, coords: np.ndarray, cells: np.ndarray, - atom_types: List[int], + atom_types: np.ndarray, atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -174,7 +174,7 @@ def eval_descriptor( self, coords: np.ndarray, cells: np.ndarray, - atom_types: List[int], + atom_types: np.ndarray, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, efield: Optional[np.ndarray] = None, @@ -258,7 +258,12 @@ def model_type(self) -> "DeepEval": @abstractmethod def get_sel_type(self) -> List[int]: - """Get the selected atom types of this model.""" + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ def get_numb_dos(self) -> int: """Get the number of DOS.""" @@ -481,7 +486,12 @@ def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type) return coords, cells, atom_types, fparam, aparam, nframes, natoms def get_sel_type(self) -> List[int]: - """Get the selected atom types of this model.""" + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ return self.deep_eval.get_sel_type() def _get_sel_natoms(self, atype) -> int: diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 2373252a8b..f857619871 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + List, Optional, + Union, ) import numpy as np @@ -43,7 +45,7 @@ def eval( self, coords: np.ndarray, cells: Optional[np.ndarray], - atom_types: np.ndarray, + atom_types: Union[List[int], np.ndarray], atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -61,7 +63,7 @@ def eval( The cell of the region. If None then non-PBC is assumed, otherwise using PBC. The array should be of size nframes x 9 - atom_types + atom_types : list[int] or np.ndarray The atom types The list should contain natoms ints atomic diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 6cbc6c058c..23fe58834e 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -2,8 +2,10 @@ from typing import ( Any, Dict, + List, Optional, Tuple, + Union, ) import numpy as np @@ -71,7 +73,7 @@ def eval( self, coords: np.ndarray, cells: Optional[np.ndarray], - atom_types: np.ndarray, + atom_types: Union[List[int], np.ndarray], atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -88,7 +90,7 @@ def eval( cells : np.ndarray The cell vectors of the system, in shape (nframes, 9). If the system is not periodic, set it to None. - atom_types : List[int] + atom_types : List[int] or np.ndarray The types of the atoms. If mixed_type is False, the shape is (natoms,); otherwise, the shape is (nframes, natoms). atomic : bool, optional diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index b9346c4633..edef0df46b 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -3,8 +3,10 @@ abstractmethod, ) from typing import ( + List, Optional, Tuple, + Union, ) import numpy as np @@ -42,7 +44,7 @@ def eval( self, coords: np.ndarray, cells: Optional[np.ndarray], - atom_types: np.ndarray, + atom_types: Union[List[int], np.ndarray], atomic: bool = True, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -60,7 +62,7 @@ def eval( The cell of the region. If None then non-PBC is assumed, otherwise using PBC. The array should be of size nframes x 9 - atom_types + atom_types : list[int] or np.ndarray The atom types The list should contain natoms ints atomic diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 3241ea9596..eb3ae99ec0 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -130,7 +130,12 @@ def model_type(self) -> "DeepEvalWrapper": return DeepPot def get_sel_type(self) -> List[int]: - """Get the selected atom types of this model.""" + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ return [] def get_numb_dos(self) -> int: diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 4bcb70349b..9141e27650 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -57,6 +57,8 @@ Path, ) + from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper + class DeepEval(DeepEvalBackend): """TensorFlow backend implementation for DeepEval. @@ -259,7 +261,7 @@ def _init_attr(self): @property @lru_cache(maxsize=None) - def model_type(self) -> str: + def model_type(self) -> "DeepEvalWrapper": """Get type of model. :type:str @@ -625,7 +627,12 @@ def get_type_map(self) -> List[str]: return self.tmap def get_sel_type(self) -> Optional[np.ndarray]: - """Get selection type.""" + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ return np.array(self.sel_type).ravel() def get_dim_fparam(self) -> int: @@ -681,7 +688,7 @@ def eval( self, coords: np.ndarray, cells: np.ndarray, - atom_types: List[int], + atom_types: np.ndarray, atomic: bool = False, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, @@ -996,7 +1003,7 @@ def eval_descriptor( self, coords: np.ndarray, cells: np.ndarray, - atom_types: List[int], + atom_types: np.ndarray, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, efield: Optional[np.ndarray] = None, @@ -1053,7 +1060,7 @@ def _eval_descriptor_inner( self, coords: np.ndarray, cells: np.ndarray, - atom_types: List[int], + atom_types: np.ndarray, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, efield: Optional[np.ndarray] = None, From 295f4d76f0bd8eea32a834266aa82d3c416a9b01 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Feb 2024 03:56:11 -0500 Subject: [PATCH 19/20] change the OutputVariableDef argument from differentiable to r_differentiable and c_differentiable Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_pot.py | 3 ++- deepmd/infer/deep_tensor.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 23fe58834e..463a07115c 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -62,7 +62,8 @@ def output_def(self) -> ModelOutputDef: "energy", shape=[1], reduciable=True, - differentiable=True, + r_differentiable=True, + c_differentiable=True, atomic=True, ), ] diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index edef0df46b..a6cefa63c1 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -229,7 +229,8 @@ def output_def(self) -> ModelOutputDef: self.output_tensor_name, shape=[-1], reduciable=True, - differentiable=True, + r_differentiable=True, + c_differentiable=True, atomic=True, ), ] From 6b104e3ad2640185217ed83a6f2e280c076712fd Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Feb 2024 04:17:13 -0500 Subject: [PATCH 20/20] fix the virial shape since DERV_C changed from (3,3) to (9,) Signed-off-by: Jinzhe Zeng --- deepmd/pt/infer/deep_eval.py | 4 ++-- deepmd/tf/infer/deep_eval.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index eb3ae99ec0..b42bee1dbe 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -341,13 +341,13 @@ def _eval_model( def _get_output_shape(self, odef, nframes, natoms): if odef.category == OutputVariableCategory.DERV_C_REDU: # virial - return [nframes, *odef.shape[:-2], 9] + return [nframes, *odef.shape[:-1], 9] elif odef.category == OutputVariableCategory.REDU: # energy return [nframes, *odef.shape, 1] elif odef.category == OutputVariableCategory.DERV_C: # atom_virial - return [nframes, *odef.shape[:-2], natoms, 9] + return [nframes, *odef.shape[:-1], natoms, 9] elif odef.category == OutputVariableCategory.DERV_R: # force return [nframes, *odef.shape[:-1], natoms, 3] diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 9141e27650..c87a888c37 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -981,13 +981,13 @@ def _eval_inner( def _get_output_shape(self, odef, nframes, natoms): if odef.category == OutputVariableCategory.DERV_C_REDU: # virial - return [nframes, *odef.shape[:-2], 9] + return [nframes, *odef.shape[:-1], 9] elif odef.category == OutputVariableCategory.REDU: # energy return [nframes, *odef.shape, 1] elif odef.category == OutputVariableCategory.DERV_C: # atom_virial - return [nframes, *odef.shape[:-2], natoms, 9] + return [nframes, *odef.shape[:-1], natoms, 9] elif odef.category == OutputVariableCategory.DERV_R: # force return [nframes, *odef.shape[:-1], natoms, 3]