deepmd.tf.descriptor
Submodules
deepmd.tf.descriptor.descriptordeepmd.tf.descriptor.hybriddeepmd.tf.descriptor.loc_framedeepmd.tf.descriptor.sedeepmd.tf.descriptor.se_adeepmd.tf.descriptor.se_a_ebddeepmd.tf.descriptor.se_a_ebd_v2deepmd.tf.descriptor.se_a_efdeepmd.tf.descriptor.se_a_maskdeepmd.tf.descriptor.se_attendeepmd.tf.descriptor.se_atten_v2deepmd.tf.descriptor.se_rdeepmd.tf.descriptor.se_t
Package Contents
Classes
The abstract class for descriptors. All specific descriptors should | |
Concate a list of descriptors to form a new descriptor. | |
Defines a local frame at each atom, and the compute the descriptor as local | |
DeepPot-SE constructed from all information (both angular and radial) of | |
DeepPot-SE descriptor with type embedding approach. | |
A compressible se_a_ebd model. | |
Smooth edition descriptor with Ef. | |
Helper class for implementing DescrptSeAEf. | |
DeepPot-SE constructed from all information (both angular and radial) of | |
Smooth version descriptor with attention. | |
Smooth version 2.0 descriptor with attention. | |
DeepPot-SE constructed from radial information of atomic configurations. | |
DeepPot-SE constructed from all information (both angular and radial) of atomic |
- class deepmd.tf.descriptor.Descriptor[source]
Bases:
deepmd.tf.utils.PluginVariant,make_plugin_registry('descriptor')The abstract class for descriptors. All specific descriptors should be based on this class.
The descriptor \(\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.
Examples
>>> descript = Descriptor(type="se_e2_a", rcut=6.0, rcut_smth=0.5, sel=[50]) >>> type(descript) <class 'deepmd.tf.descriptor.se_a.DescrptSeA'>
- abstract get_rcut() float[source]
Returns the cut-off radius.
- Returns:
floatthe cut-off radius
Notes
This method must be implemented, as it’s called by other classes.
- abstract get_ntypes() int[source]
Returns the number of atom types.
- Returns:
intthe number of atom types
Notes
This method must be implemented, as it’s called by other classes.
- abstract get_dim_out() int[source]
Returns the output dimension of this descriptor.
- Returns:
intthe output dimension of this descriptor
Notes
This method must be implemented, as it’s called by other classes.
- abstract get_dim_rot_mat_1() int[source]
Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.
- Returns:
intthe first dimension of the rotation matrix
- abstract get_nlist() Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
Returns neighbor information.
- abstract compute_input_stats(data_coord: List[numpy.ndarray], data_box: List[numpy.ndarray], data_atype: List[numpy.ndarray], natoms_vec: List[numpy.ndarray], mesh: List[numpy.ndarray], input_dict: Dict[str, List[numpy.ndarray]], **kwargs) None[source]
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
deepmd.tf.model.model_stat.make_stat_input()- data_box
list[np.ndarray] The box. Can be generated by
deepmd.tf.model.model_stat.make_stat_input()- data_atype
list[np.ndarray] The atom types. Can be generated by
deepmd.tf.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
deepmd.tf.model.model_stat.make_stat_input()- mesh
list[np.ndarray] The mesh for neighbor searching. Can be generated by
deepmd.tf.model.model_stat.make_stat_input()- input_dict
dict[str,list[np.ndarray]] Dictionary for additional input
- **kwargs
Additional keyword arguments which may contain mixed_type and real_natoms_vec.
- data_coord
Notes
This method must be implemented, as it’s called by other classes.
- abstract build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: Dict[str, Any], reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
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
- reusebool,
optional The weights in the networks should be reused when get the variable.
- suffix
str,optional Name suffix to identify this descriptor
- coord_
- Returns:
- descriptor:
tf.Tensor The output descriptor
- descriptor:
Notes
This method must be implemented, as it’s called by other classes.
- abstract enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]
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
- graph
tf.Graph The graph of the model
- graph_def
tf.GraphDef The graph definition of the model
- 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
- suffix
str,optional The suffix of the scope
- min_nbor_dist
Notes
This method is called by others when the descriptor supported compression.
- abstract enable_mixed_precision(mixed_prec: dict | None = None) None[source]
Reveive the mixed precision setting.
- Parameters:
- mixed_prec
The mixed precision setting used in the embedding net
Notes
This method is called by others when the descriptor supported compression.
- abstract prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- Returns:
- abstract init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') None[source]
Init the embedding net variables with the given dict.
- Parameters:
Notes
This method is called by others when the descriptor supported initialization from the given variables.
- abstract pass_tensors_from_frz_model(*tensors: deepmd.tf.env.tf.Tensor) None[source]
Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
- Parameters:
- *tensors
tf.Tensor passed tensors
- *tensors
Notes
The number of parameters in the method must be equal to the numbers of returns in
get_tensor_names().
- build_type_exclude_mask(exclude_types: Set[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: deepmd.tf.env.tf.Tensor, shape0: deepmd.tf.env.tf.Tensor) deepmd.tf.env.tf.Tensor[source]
Build the type exclude mask for the descriptor.
- Parameters:
- exclude_types
List[Tuple[int,int]] The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction between type 0 and type 1 is excluded.
- ntypes
int The number of types.
- sel
List[int] The list of the number of selected neighbors for each type.
- ndescrpt
int The number of descriptors for each atom.
- atype
tf.Tensor The type of atoms, with the size of shape0.
- shape0
tf.Tensor The shape of the first dimension of the inputs, which is equal to nsamples * natoms.
- exclude_types
- Returns:
tf.TensorThe type exclude mask, with the shape of (shape0, ndescrpt), and the precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the interaction between two types is not excluded, and 0 otherwise.
Notes
To exclude the interaction between two types, the derivative of energy with respect to distances (or angles) between two atoms should be zero[Rdfe82bee38f2-1]_, i.e.
\[\forall i \in \text{type 1}, j \in \text{type 2}, \frac{\partial{E}}{\partial{r_{ij}}} = 0\]When embedding networks between every two types are built, we can just remove that network. But when type_one_side is enabled, a network may be built for multiple pairs of types. In this case, we need to build a mask to exclude the interaction between two types.
The mask assumes the descriptors are sorted by neighbro type with the fixed number of given sel and each neighbor has the same number of descriptors (for example 4).
References
[1]Jinzhe Zeng, Timothy J. Giese, ̧Sölen Ekesan, Darrin M. York, Development of Range-Corrected Deep Learning Potentials for Fast, Accurate Quantum Mechanical/molecular Mechanical Simulations of Chemical Reactions in Solution, J. Chem. Theory Comput., 2021, 17 (11), 6993-7009.
- abstract classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]
Update the selection and perform neighbor statistics.
- classmethod deserialize(data: dict, suffix: str = '') Descriptor[source]
Deserialize the model.
There is no suffix in a native DP model, but it is important for the TF backend.
- Parameters:
- Returns:
DescriptorThe deserialized descriptor
- class deepmd.tf.descriptor.DescrptHybrid(list: List[deepmd.tf.descriptor.descriptor.Descriptor | Dict[str, Any]], multi_task: bool = False, ntypes: int | None = None, spin: deepmd.tf.utils.spin.Spin | None = None, **kwargs)[source]
Bases:
deepmd.tf.descriptor.descriptor.DescriptorConcate a list of descriptors to form a new descriptor.
- Parameters:
- get_nlist() Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- get_nlist_i(ii: int) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
Get the neighbor information of the ii-th descriptor.
- Parameters:
- ii
int The index of the descriptor
- ii
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: list | None = None, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- 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.
- real_natoms_vec
If mixed_type is True, it takes in the real natoms_vec for each frame.
- **kwargs
Additional keyword arguments.
- merge_input_stats(stat_dict)[source]
Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
- Parameters:
- stat_dict
The dict of statisitcs computed from compute_input_stats, including:
- sumr
The sum of radial statisitcs.
- suma
The sum of relative coord statisitcs.
- sumn
The sum of neighbor numbers.
- sumr2
The sum of square of radial statisitcs.
- suma2
The sum of square of relative coord statisitcs.
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialThe total virial
atom_virialThe atomic virial
- enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]
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
- graph
tf.Graph The graph of the model
- graph_def
tf.GraphDef The graph_def of the model
- 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
- suffix
str,optional The suffix of the scope
- min_nbor_dist
- enable_mixed_precision(mixed_prec: dict | None = None) None[source]
Reveive the mixed precision setting.
- Parameters:
- mixed_prec
The mixed precision setting used in the embedding net
- init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') None[source]
Init the embedding net variables with the given dict.
- pass_tensors_from_frz_model(*tensors: deepmd.tf.env.tf.Tensor) None[source]
Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
- Parameters:
- *tensors
tf.Tensor passed tensors
- *tensors
- classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]
Update the selection and perform neighbor statistics.
- serialize(suffix: str = '') dict[source]
Serialize the model.
There is no suffix in a native DP model, but it is important for the TF backend.
- classmethod deserialize(data: dict, suffix: str = '') DescrptHybrid[source]
Deserialize the model.
There is no suffix in a native DP model, but it is important for the TF backend.
- Parameters:
- Returns:
DescriptorThe deserialized descriptor
- class deepmd.tf.descriptor.DescrptLocFrame(rcut: float, sel_a: List[int], sel_r: List[int], axis_rule: List[int], **kwargs)[source]
Bases:
deepmd.tf.descriptor.descriptor.DescriptorDefines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.
- Parameters:
- rcut
The cut-off radius
- sel_a
list[int] The length of the list should be the same as the number of atom types in the system. sel_a[i] gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor.
- sel_r
list[int] The length of the list should be the same as the number of atom types in the system. sel_r[i] gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.
- axis_rule: list[int]
The length should be 6 times of the number of types. - axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance. - axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom. - axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance. - axis_rule[i*6+3]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance. - axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom. - axis_rule[i*6+5]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance.
- get_nlist() Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialThe total virial
atom_virialThe atomic virial
- class deepmd.tf.descriptor.DescrptSeA(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, multi_task: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, stripped_type_embedding: bool = False, env_protection: float = 0.0, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se.DescrptSeDeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.
The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]
\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows
\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:
\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):
\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at
deepmd.tf.utils.network.embedding_net().- Parameters:
- rcut
The cut-off radius \(r_c\)
- rcut_smth
From where the environment matrix should be smoothed \(r_s\)
- sel
list[int] sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list[int] Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_one_side
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- multi_task
If the model has multi fitting nets to train.
- env_protection: float
Protection parameter to prevent division by zero errors during environment matrix calculations.
References
[1]Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
- get_dim_rot_mat_1() int[source]
Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.
- get_nlist() Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
Returns neighbor information.
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- merge_input_stats(stat_dict)[source]
Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
- Parameters:
- stat_dict
The dict of statisitcs computed from compute_input_stats, including:
- sumr
The sum of radial statisitcs.
- suma
The sum of relative coord statisitcs.
- sumn
The sum of neighbor numbers.
- sumr2
The sum of square of radial statisitcs.
- suma2
The sum of square of relative coord statisitcs.
- enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
- Parameters:
- min_nbor_dist
The nearest distance between atoms
- graph
tf.Graph The graph of the model
- graph_def
tf.GraphDef The graph_def of the model
- table_extrapolate
The scale of model extrapolation
- table_stride_1
The uniform stride of the first table
- table_stride_2
The uniform stride of the second table
- check_frequency
The overflow check frequency
- suffix
str,optional The suffix of the scope
- enable_mixed_precision(mixed_prec: dict | None = None) None[source]
Reveive the mixed precision setting.
- Parameters:
- mixed_prec
The mixed precision setting used in the embedding net
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialThe total virial
atom_virialThe atomic virial
- _concat_type_embedding(xyz_scatter, nframes, natoms, type_embedding)[source]
Concatenate type_embedding of neighbors and xyz_scatter. If not self.type_one_side, concatenate type_embedding of center atoms as well.
- Parameters:
- xyz_scatter:
shape is [nframes*natoms[0]*self.nnei, 1]
- nframes:
shape is []
- natoms:
shape is [1+1+self.ntypes]
- type_embedding:
shape is [self.ntypes, Y] where Y=jdata[‘type_embedding’][‘neuron’][-1]
- Returns:
- embedding:
environment of each atom represented by embedding.
- _filter_lower(type_i, type_input, start_index, incrs_index, inputs, nframes, natoms, type_embedding=None, is_exclude=False, activation_fn=None, bavg=0.0, stddev=1.0, trainable=True, suffix='')[source]
Input env matrix, returns R.G.
- _filter(inputs, type_input, natoms, type_embedding=None, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, trainable=True)[source]
- init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') None[source]
Init the embedding net variables with the given dict.
- class deepmd.tf.descriptor.DescrptSeAEbd(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, type_nchanl: int = 2, type_nlayer: int = 1, numb_aparam: int = 0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', exclude_types: List[List[int]] = [], **kwargs)[source]
Bases:
deepmd.tf.descriptor.se_a.DescrptSeADeepPot-SE descriptor with type embedding approach.
- Parameters:
- rcut
The cut-off radius
- rcut_smth
From where the environment matrix should be smoothed
- sel
list[int] sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list[int] Number of neurons in each hidden layers of the embedding net
- axis_neuron
Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_one_side
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- type_nchanl
Number of channels for type representation
- type_nlayer
Number of hidden layers for the type embedding net (skip connected).
- numb_aparam
Number of atomic parameters. If >0 it will be embedded with atom types.
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are {0}
- precision
The precision of the embedding net parameters. Supported options are {1}
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- _embedding_net(inputs, natoms, filter_neuron, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True)[source]
inputs: nf x na x (nei x 4) outputs: nf x na x nei x output_size.
- _type_embedding_net_two_sides(mat_g, atype, natoms, name='', reuse=None, seed=None, trainable=True)[source]
- _type_embedding_net_one_side(mat_g, atype, natoms, name='', reuse=None, seed=None, trainable=True)[source]
- class deepmd.tf.descriptor.DescrptSeAEbdV2(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, multi_task: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se_a.DescrptSeAA compressible se_a_ebd model.
This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True.
- class deepmd.tf.descriptor.DescrptSeAEf(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed=False, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se.DescrptSeSmooth edition descriptor with Ef.
- Parameters:
- rcut
The cut-off radius
- rcut_smth
From where the environment matrix should be smoothed
- sel
list[int] sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list[int] Number of neurons in each hidden layers of the embedding net
- axis_neuron
Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_one_side
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- get_dim_rot_mat_1() int[source]
Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.
- get_nlist() Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
Returns neighbor information.
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs. Should have ‘efield’.
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialThe total virial
atom_virialThe atomic virial
- class deepmd.tf.descriptor.DescrptSeAEfLower(op, rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False)[source]
Bases:
deepmd.tf.descriptor.se_a.DescrptSeAHelper class for implementing DescrptSeAEf.
- compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict, **kwargs)[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- build(coord_, atype_, natoms, box_, mesh, input_dict, suffix='', reuse=None)[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- class deepmd.tf.descriptor.DescrptSeAMask(sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, exclude_types: List[List[int]] = [], seed: int | None = None, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, stripped_type_embedding: bool = False, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se_a.DescrptSeADeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.
The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]
\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows
\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:
\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):
\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at
deepmd.tf.utils.network.embedding_net(). Specially for descriptor se_a_mask is a concise implementation of se_a. The difference is that se_a_mask only considered a non-pbc system. And accept a mask matrix to indicate the atom i in frame j is a real atom or not. (1 means real atom, 0 means ghost atom) Thus se_a_mask can accept a variable number of atoms in a frame.- Parameters:
- sel
list[int] sel[i] specifies the maxmum number of type i atoms in the neighbor list.
- neuron
list[int] Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_one_side
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- activation_function
The activation function in the embedding net. Supported options are {0}
- precision
The precision of the embedding net parameters. Supported options are {1}
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- sel
References
[1]Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: Dict[str, Any], reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialNone for se_a_mask op
atom_virialNone for se_a_mask op
- class deepmd.tf.descriptor.DescrptSeAtten(rcut: float, rcut_smth: float, sel: List[int] | int, ntypes: int, neuron: List[int] = [25, 50, 100], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, set_davg_zero: bool = True, exclude_types: List[List[int]] = [], activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, multi_task: bool = False, stripped_type_embedding: bool = False, smooth_type_embedding: bool = False, scaling_factor=1.0, normalize=True, temperature=None, trainable_ln: bool = True, ln_eps: float | None = 0.001, concat_output_tebd: bool = True, env_protection: float = 0.0, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se_a.DescrptSeASmooth version descriptor with attention.
- Parameters:
- rcut: float
The cut-off radius \(r_c\)
- rcut_smth: float
From where the environment matrix should be smoothed \(r_s\)
- sel: list[int], int
list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius
- neuron: list[int]
Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron: int
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt: bool
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable: bool
If the weights of embedding net are trainable.
- seed: int, Optional
Random seed for initializing the network parameters.
- type_one_side: bool
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- set_davg_zero: bool
Set the shift of embedding net input to zero.
- activation_function: str
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision: str
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- uniform_seed: bool
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- attn: int
The length of hidden vector during scale-dot attention computation.
- attn_layer: int
The number of layers in attention mechanism.
- attn_dotr: bool
Whether to dot the relative coordinates on the attention weights as a gated scheme.
- attn_mask: bool
Whether to mask the diagonal in the attention weights.
- ln_eps: float, Optional
The epsilon value for layer normalization.
- multi_task: bool
If the model has multi fitting nets to train.
- stripped_type_embedding: bool
Whether to strip the type embedding into a separated embedding network. Default value will be True in se_atten_v2 descriptor.
- smooth_type_embedding: bool
Whether to use smooth process in attention weights calculation. And when using stripped type embedding, whether to dot smooth factor on the network output of type embedding to keep the network smooth, instead of setting set_davg_zero to be True. Default value will be True in se_atten_v2 descriptor.
- Raises:
ValueErrorif ntypes is 0.
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: list | None = None, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. If mixed_type is True, this para is blank. See real_natoms_vec.
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- 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.
- real_natoms_vec
If mixed_type is True, it takes in the real natoms_vec for each frame.
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
- Parameters:
- min_nbor_dist
The nearest distance between atoms
- graph
tf.Graph The graph of the model
- graph_def
tf.GraphDef The graph_def of the model
- table_extrapolate
The scale of model extrapolation
- table_stride_1
The uniform stride of the first table
- table_stride_2
The uniform stride of the second table
- check_frequency
The overflow check frequency
- suffix
str,optional The suffix of the scope
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- _compute_dstats_sys_smth(data_coord, data_box, data_atype, natoms_vec, mesh, mixed_type=False, real_natoms_vec=None)[source]
- _lookup_type_embedding(xyz_scatter, natype, type_embedding)[source]
Concatenate type_embedding of neighbors and xyz_scatter. If not self.type_one_side, concatenate type_embedding of center atoms as well.
- Parameters:
- xyz_scatter:
shape is [nframes*natoms[0]*self.nnei, 1]
- natype:
neighbor atom type
- type_embedding:
shape is [self.ntypes, Y] where Y=jdata[‘type_embedding’][‘neuron’][-1]
- Returns:
- embedding:
environment of each atom represented by embedding.
- _scaled_dot_attn(Q, K, V, temperature, input_r, dotr=False, do_mask=False, layer=0, save_weights=True)[source]
- _attention_layers(input_xyz, layer_num, shape_i, outputs_size, input_r, dotr=False, do_mask=False, trainable=True, suffix='')[source]
- _filter_lower(type_i, type_input, start_index, incrs_index, inputs, type_embedding=None, atype=None, is_exclude=False, activation_fn=None, bavg=0.0, stddev=1.0, trainable=True, suffix='', name='filter_', reuse=None)[source]
Input env matrix, returns R.G.
- _filter(inputs, type_input, natoms, type_embedding=None, atype=None, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, suffix='', name='linear', reuse=None, trainable=True)[source]
- init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') None[source]
Init the embedding net variables with the given dict.
- build_type_exclude_mask_mixed(exclude_types: Set[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: deepmd.tf.env.tf.Tensor, shape0: deepmd.tf.env.tf.Tensor, nei_type_vec: deepmd.tf.env.tf.Tensor) deepmd.tf.env.tf.Tensor[source]
Build the type exclude mask for the attention descriptor.
- Parameters:
- exclude_types
List[Tuple[int,int]] The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction between type 0 and type 1 is excluded.
- ntypes
int The number of types.
- sel
List[int] The list of the number of selected neighbors for each type.
- ndescrpt
int The number of descriptors for each atom.
- atype
tf.Tensor The type of atoms, with the size of shape0.
- shape0
tf.Tensor The shape of the first dimension of the inputs, which is equal to nsamples * natoms.
- nei_type_vec
tf.Tensor The type of neighbors, with the size of (shape0, nnei).
- exclude_types
- Returns:
tf.TensorThe type exclude mask, with the shape of (shape0, ndescrpt), and the precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the interaction between two types is not excluded, and 0 otherwise.
Notes
This method has the similiar way to build the type exclude mask as
deepmd.tf.descriptor.descriptor.Descriptor.build_type_exclude_mask(). The mathmatical expression has been explained in that method. The difference is that the attention descriptor has provided the type of the neighbors (idx_j) that is not in order, so we use it from an extra input.
- classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]
Update the selection and perform neighbor statistics.
- serialize_attention_layers(nlayer: int, nnei: int, embed_dim: int, hidden_dim: int, dotr: bool, do_mask: bool, trainable_ln: bool, ln_eps: float, variables: dict, bias: bool = True, suffix: str = '') dict[source]
- classmethod deserialize_attention_layers(data: dict, suffix: str = '') dict[source]
Deserialize attention layers.
- class deepmd.tf.descriptor.DescrptSeAttenV2(rcut: float, rcut_smth: float, sel: int, ntypes: int, neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, set_davg_zero: bool = False, exclude_types: List[List[int]] = [], activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, multi_task: bool = False, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se_atten.DescrptSeAttenSmooth version 2.0 descriptor with attention.
- Parameters:
- rcut
The cut-off radius \(r_c\)
- rcut_smth
From where the environment matrix should be smoothed \(r_s\)
- sel
int sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list[int] Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_one_side
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- attn
The length of hidden vector during scale-dot attention computation.
- attn_layer
The number of layers in attention mechanism.
- attn_dotr
Whether to dot the relative coordinates on the attention weights as a gated scheme.
- attn_mask
Whether to mask the diagonal in the attention weights.
- multi_task
If the model has multi fitting nets to train.
- class deepmd.tf.descriptor.DescrptSeR(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, multi_task: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, env_protection: float = 0.0, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se.DescrptSeDeepPot-SE constructed from radial information of atomic configurations.
The embedding takes the distance between atoms as input.
- Parameters:
- rcut
The cut-off radius
- rcut_smth
From where the environment matrix should be smoothed
- sel
list[int] sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list[int] Number of neurons in each hidden layers of the embedding net
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_one_side
Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
- exclude_types
List[List[int]] The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
- activation_function
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- get_nlist()[source]
Returns neighbor information.
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict, **kwargs)[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- merge_input_stats(stat_dict)[source]
Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
- Parameters:
- stat_dict
The dict of statisitcs computed from compute_input_stats, including:
- sumr
The sum of radial statisitcs.
- sumn
The sum of neighbor numbers.
- sumr2
The sum of square of radial statisitcs.
- enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
- Parameters:
- min_nbor_dist
The nearest distance between atoms
- graph
tf.Graph The graph of the model
- graph_def
tf.GraphDef The graph_def of the model
- table_extrapolate
The scale of model extrapolation
- table_stride_1
The uniform stride of the first table
- table_stride_2
The uniform stride of the second table
- check_frequency
The overflow check frequency
- suffix
str,optional The suffix of the scope
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialThe total virial
atom_virialThe atomic virial
- _filter_r(inputs, type_input, natoms, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, trainable=True)[source]
- class deepmd.tf.descriptor.DescrptSeT(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, multi_task: bool = False, **kwargs)[source]
Bases:
deepmd.tf.descriptor.se.DescrptSeDeepPot-SE constructed from all information (both angular and radial) of atomic configurations.
The embedding takes angles between two neighboring atoms as input.
- Parameters:
- rcut
The cut-off radius
- rcut_smth
From where the environment matrix should be smoothed
- sel
list[int] sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list[int] Number of neurons in each hidden layers of the embedding net
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- get_nlist() Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]][source]
Returns neighbor information.
- Returns:
nlistNeighbor list
rijThe relative distance between the neighbor and the center atom.
sel_aThe number of neighbors with full information
sel_rThe number of neighbors with only radial information
- compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]
Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
- Parameters:
- data_coord
The coordinates. Can be generated by deepmd.tf.model.make_stat_input
- data_box
The box. Can be generated by deepmd.tf.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.tf.model.make_stat_input
- natoms_vec
The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- merge_input_stats(stat_dict)[source]
Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
- Parameters:
- stat_dict
The dict of statisitcs computed from compute_input_stats, including:
- sumr
The sum of radial statisitcs.
- suma
The sum of relative coord statisitcs.
- sumn
The sum of neighbor numbers.
- sumr2
The sum of square of radial statisitcs.
- suma2
The sum of square of relative coord statisitcs.
- enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
- Parameters:
- min_nbor_dist
The nearest distance between atoms
- graph
tf.Graph The graph of the model
- graph_def
tf.GraphDef The graph_def of the model
- table_extrapolate
The scale of model extrapolation
- table_stride_1
The uniform stride of the first table
- table_stride_2
The uniform stride of the second table
- check_frequency
The overflow check frequency
- suffix
str,optional The suffix of the scope
- build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') deepmd.tf.env.tf.Tensor[source]
Build the computational graph for the descriptor.
- Parameters:
- coord_
The coordinate of atoms
- atype_
The type of atoms
- 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
- box_
tf.Tensor The box of the system
- mesh
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
Dictionary for additional inputs
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
descriptorThe output descriptor
- prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]
Compute force and virial.
- Parameters:
- atom_ener
The atomic energy
- 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
- Returns:
forceThe force on atoms
virialThe total virial
atom_virialThe atomic virial