From 44939c281fb84b6e3c2edaf4a4986d71521108d8 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 1 Sep 2021 02:34:39 -0400 Subject: [PATCH 1/3] add ABC for descriptors I'm going to add abstract base classes for different object, where a list of methods and attributes is defined to normalize classes and their external call by other classes. It's also useful to develop and extend new classes. The first one I did is the descriptor. --- deepmd/descriptor/descriptor.py | 275 ++++++++++++++++++++++++++++++++ deepmd/descriptor/hybrid.py | 3 +- deepmd/descriptor/loc_frame.py | 3 +- deepmd/descriptor/se_a.py | 3 +- deepmd/descriptor/se_a_ef.py | 3 +- deepmd/descriptor/se_ar.py | 3 +- deepmd/descriptor/se_r.py | 3 +- deepmd/descriptor/se_t.py | 3 +- 8 files changed, 289 insertions(+), 7 deletions(-) create mode 100644 deepmd/descriptor/descriptor.py diff --git a/deepmd/descriptor/descriptor.py b/deepmd/descriptor/descriptor.py new file mode 100644 index 0000000000..20d26b49b9 --- /dev/null +++ b/deepmd/descriptor/descriptor.py @@ -0,0 +1,275 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +if TYPE_CHECKING: + import numpy as np + from deepmd.env import tf + + +class Descriptor(ABC): + """The abstract class for descriptors. All specific descriptors should + be based on this class. + + The descriptor :math:`\mathcal{D}` describes the environment of an atom, + which should be a function of coordinates and types of its neighbour atoms. + + Notes + ----- + Only methods and attributes defined in this class are generally public, + that can be called by other classes. + """ + + @abstractmethod + def get_rcut(self) -> float: + """ + Returns the cut-off radius. + + Returns + ------- + float + the cut-off radius + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + @abstractmethod + def get_ntypes(self) -> int: + """ + Returns the number of atom types. + + Returns + ------- + int + the number of atom types + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + @abstractmethod + def get_dim_out(self) -> int: + """ + Returns the output dimension of this descriptor. + + Returns + ------- + int + the output dimension of this descriptor + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + def get_dim_rot_mat_1(self) -> int: + """ + Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3 + + Returns + ------- + int + the first dimension of the rotation matrix + """ + # TODO: I think this method should be implemented as it's called by dipole and + # polar fitting network. However, currently not all descriptors have this method. + raise NotImplementedError + + def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns neighbor information. + + Returns + ------- + nlist : tf.Tensor + Neighbor list + rij : tf.Tensor + The relative distance between the neighbor and the center atom. + sel_a : list[int] + The number of neighbors with full information + sel_r : list[int] + The number of neighbors with only radial information + """ + # TODO: I think this method should be implemented as it's called by energy model. + # However, se_ar and hybrid doesn't have this method. + raise NotImplementedError + + @abstractmethod + def compute_input_stats(self, + data_coord: List[np.ndarray], + data_box: List[np.ndarray], + data_atype: List[np.ndarray], + natoms_vec: List[np.ndarray], + mesh: List[np.ndarray], + input_dict: Dict[str, List[np.ndarray]] + ) -> None: + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord : list[np.ndarray] + The coordinates. Can be generated by :meth:`deepmd.model.make_stat_input` + data_box : list[np.ndarray] + The box. Can be generated by :meth:`deepmd.model.make_stat_input` + data_atype : list[np.ndarray] + The atom types. Can be generated by :meth:`deepmd.model.make_stat_input` + natoms_vec : list[np.ndarray] + The vector for the number of atoms of the system and different types of atoms. Can be generated by :meth:`deepmd.model.make_stat_input` + mesh : list[np.ndarray] + The mesh for neighbor searching. Can be generated by :meth:`deepmd.model.make_stat_input` + input_dict : dict[str, list[np.ndarray]] + Dictionary for additional input + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + @abstractmethod + def build(self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box_: tf.Tensor, + mesh: tf.Tensor, + input_dict: Dict[str, Any], + reuse: bool = None, + suffix: str = '', + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor. + + Parameters + ---------- + coord_ : tf.Tensor + The coordinate of atoms + atype_ : tf.Tensor + The type of atoms + natoms : tf.Tensor + 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 + box : tf.Tensor + The box of frames + mesh : tf.Tensor + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict : dict[str, Any] + Dictionary for additional inputs + reuse : bool, optional + The weights in the networks should be reused when get the variable. + suffix : str, optional + Name suffix to identify this descriptor + + Returns + ------- + descriptor: tf.Tensor + The output descriptor + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + def enable_compression(self, + min_nbor_dist: float, + model_file: str = 'frozon_model.pb', + table_extrapolate: float = 5., + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1 + ) -> None: + """ + Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data. + + Parameters + ---------- + min_nbor_dist : float + The nearest distance between atoms + model_file : str, default: 'frozon_model.pb' + The original frozen model, which will be compressed by the program + table_extrapolate : float, default: 5. + The scale of model extrapolation + table_stride_1 : float, default: 0.01 + The uniform stride of the first table + table_stride_2 : float, default: 0.1 + The uniform stride of the second table + check_frequency : int, default: -1 + The overflow check frequency + + Notes + ----- + This method is called by others when the descriptor supported compression. + """ + raise NotImplementedError( + "Descriptor %s doesn't support compression!" % self.__name__) + + @abstractmethod + def prod_force_virial(self, + atom_ener: tf.Tensor, + natoms: tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial. + + Parameters + ---------- + atom_ener : tf.Tensor + The atomic energy + natoms : tf.Tensor + 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 + + Returns + ------- + force : tf.Tensor + The force on atoms + virial : tf.Tensor + The total virial + atom_virial : tf.Tensor + The atomic virial + """ + + def get_feed_dict(self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box: tf.Tensor, + mesh: tf.Tensor + ) -> Dict[str, tf.Tensor]: + """ + Generate the feed_dict for current descriptor + + Parameters + ---------- + coord_ : tf.Tensor + The coordinate of atoms + atype_ : tf.Tensor + The type of atoms + natoms : tf.Tensor + 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 + box : tf.Tensor + The box. Can be generated by deepmd.model.make_stat_input + mesh : tf.Tensor + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + + Returns + ------- + feed_dict : dict[str, tf.Tensor] + The output feed_dict of current descriptor + """ + # TODO: currently only SeA has this method, but I think the method can be + # moved here as it doesn't contain anything related to a specific descriptor + raise NotImplementedError diff --git a/deepmd/descriptor/hybrid.py b/deepmd/descriptor/hybrid.py index 013ee9f753..d4e6d795b7 100644 --- a/deepmd/descriptor/hybrid.py +++ b/deepmd/descriptor/hybrid.py @@ -12,6 +12,7 @@ # from deepmd.descriptor import DescrptSeAEbd # from deepmd.descriptor import DescrptSeAEf # from deepmd.descriptor import DescrptSeR +from .descriptor import Descriptor from .se_a import DescrptSeA from .se_r import DescrptSeR from .se_ar import DescrptSeAR @@ -20,7 +21,7 @@ from .se_a_ef import DescrptSeAEf from .loc_frame import DescrptLocFrame -class DescrptHybrid (): +class DescrptHybrid (Descriptor): """Concate a list of descriptors to form a new descriptor. Parameters diff --git a/deepmd/descriptor/loc_frame.py b/deepmd/descriptor/loc_frame.py index be57403d60..a664384ba6 100644 --- a/deepmd/descriptor/loc_frame.py +++ b/deepmd/descriptor/loc_frame.py @@ -7,8 +7,9 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from deepmd.utils.sess import run_sess +from .descriptor import Descriptor -class DescrptLocFrame () : +class DescrptLocFrame (Descriptor) : """Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame. diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 3777368c0f..494c7e5fa6 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -14,8 +14,9 @@ from deepmd.utils.type_embed import embed_atom_type from deepmd.utils.sess import run_sess from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph +from .descriptor import Descriptor -class DescrptSeA (): +class DescrptSeA (Descriptor): r"""DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input. diff --git a/deepmd/descriptor/se_a_ef.py b/deepmd/descriptor/se_a_ef.py index 996de6ff2d..e2424aeae4 100644 --- a/deepmd/descriptor/se_a_ef.py +++ b/deepmd/descriptor/se_a_ef.py @@ -10,8 +10,9 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from .se_a import DescrptSeA +from .descriptor import Descriptor -class DescrptSeAEf (): +class DescrptSeAEf (Descriptor): """ Parameters diff --git a/deepmd/descriptor/se_ar.py b/deepmd/descriptor/se_ar.py index 8ded2c5849..518cf9315c 100644 --- a/deepmd/descriptor/se_ar.py +++ b/deepmd/descriptor/se_ar.py @@ -5,8 +5,9 @@ from .se_a import DescrptSeA from .se_r import DescrptSeR from deepmd.env import op_module +from .descriptor import Descriptor -class DescrptSeAR (): +class DescrptSeAR (Descriptor): def __init__ (self, jdata): args = ClassArg()\ .add('a', dict, must = True) \ diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index fafa03a350..b6de76be76 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -10,8 +10,9 @@ from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift from deepmd.utils.sess import run_sess +from .descriptor import Descriptor -class DescrptSeR (): +class DescrptSeR (Descriptor): """DeepPot-SE constructed from radial information of atomic configurations. The embedding takes the distance between atoms as input. diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index 94298c6eb4..29f9323b34 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -10,8 +10,9 @@ from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift from deepmd.utils.sess import run_sess +from .descriptor import Descriptor -class DescrptSeT (): +class DescrptSeT (Descriptor): """DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. From 742662045cae086d1ba3bfadc7616f7b21cd6d5d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 1 Sep 2021 02:42:14 -0400 Subject: [PATCH 2/3] TYPE_CHECKING doesn't work in python 3.6 --- deepmd/descriptor/descriptor.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/deepmd/descriptor/descriptor.py b/deepmd/descriptor/descriptor.py index 20d26b49b9..079c6d230d 100644 --- a/deepmd/descriptor/descriptor.py +++ b/deepmd/descriptor/descriptor.py @@ -1,13 +1,12 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Tuple +from typing import Any, Dict, List, Tuple -if TYPE_CHECKING: - import numpy as np - from deepmd.env import tf +import numpy as np +from deepmd.env import tf class Descriptor(ABC): - """The abstract class for descriptors. All specific descriptors should + r"""The abstract class for descriptors. All specific descriptors should be based on this class. The descriptor :math:`\mathcal{D}` describes the environment of an atom, From 993100bedff17acd4eeb25565c971994bfaa8504 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 1 Sep 2021 02:53:33 -0400 Subject: [PATCH 3/3] fix warnings --- deepmd/descriptor/descriptor.py | 34 ++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/deepmd/descriptor/descriptor.py b/deepmd/descriptor/descriptor.py index 079c6d230d..c7ca158e05 100644 --- a/deepmd/descriptor/descriptor.py +++ b/deepmd/descriptor/descriptor.py @@ -65,7 +65,8 @@ def get_dim_out(self) -> int: def get_dim_rot_mat_1(self) -> int: """ - Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3 + Returns the first dimension of the rotation matrix. The rotation is of shape + dim_1 x 3 Returns ------- @@ -73,7 +74,8 @@ def get_dim_rot_mat_1(self) -> int: the first dimension of the rotation matrix """ # TODO: I think this method should be implemented as it's called by dipole and - # polar fitting network. However, currently not all descriptors have this method. + # polar fitting network. However, currently not all descriptors have this + # method. raise NotImplementedError def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: @@ -91,8 +93,8 @@ def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: sel_r : list[int] The number of neighbors with only radial information """ - # TODO: I think this method should be implemented as it's called by energy model. - # However, se_ar and hybrid doesn't have this method. + # TODO: I think this method should be implemented as it's called by energy + # model. However, se_ar and hybrid doesn't have this method. raise NotImplementedError @abstractmethod @@ -105,20 +107,25 @@ def compute_input_stats(self, input_dict: Dict[str, List[np.ndarray]] ) -> None: """ - Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + Compute the statisitcs (avg and std) of the training data. The input will be + normalized by the statistics. Parameters ---------- data_coord : list[np.ndarray] - The coordinates. Can be generated by :meth:`deepmd.model.make_stat_input` + The coordinates. Can be generated by + :meth:`deepmd.model.model_stat.make_stat_input` data_box : list[np.ndarray] - The box. Can be generated by :meth:`deepmd.model.make_stat_input` + The box. Can be generated by + :meth:`deepmd.model.model_stat.make_stat_input` data_atype : list[np.ndarray] - The atom types. Can be generated by :meth:`deepmd.model.make_stat_input` + The atom types. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input` natoms_vec : list[np.ndarray] - The vector for the number of atoms of the system and different types of atoms. Can be generated by :meth:`deepmd.model.make_stat_input` + The vector for the number of atoms of the system and different types of + atoms. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input` mesh : list[np.ndarray] - The mesh for neighbor searching. Can be generated by :meth:`deepmd.model.make_stat_input` + The mesh for neighbor searching. Can be generated by + :meth:`deepmd.model.model_stat.make_stat_input` input_dict : dict[str, list[np.ndarray]] Dictionary for additional input @@ -156,8 +163,8 @@ def build(self, The box of frames mesh : tf.Tensor For historical reasons, only the length of the Tensor matters. - if size of mesh == 6, pbc is assumed. - if size of mesh == 0, no-pbc is assumed. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. input_dict : dict[str, Any] Dictionary for additional inputs reuse : bool, optional @@ -184,7 +191,8 @@ def enable_compression(self, check_frequency: int = -1 ) -> None: """ - Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data. + Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the + training data. Parameters ----------