deepmd.descriptor package
- class deepmd.descriptor.Descriptor(*args, **kwargs)[source]
Bases:
PluginVariant
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., rcut_smth=0.5, sel=[50]) >>> type(descript) <class 'deepmd.descriptor.se_a.DescrptSeA'>
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(*tensors)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') 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.
- build_type_exclude_mask(exclude_types: Set[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor) 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.Tensor
The 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[R08579741114c-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 compute_input_stats(data_coord: List[ndarray], data_box: List[ndarray], data_atype: List[ndarray], natoms_vec: List[ndarray], mesh: List[ndarray], input_dict: Dict[str, List[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.model.model_stat.make_stat_input()
- data_box
list
[np.ndarray
] The box. Can be generated by
deepmd.model.model_stat.make_stat_input()
- data_atype
list
[np.ndarray
] The atom types. Can be generated by
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
deepmd.model.model_stat.make_stat_input()
- mesh
list
[np.ndarray
] The mesh for neighbor searching. Can be generated by
deepmd.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.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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.
- enable_mixed_precision(mixed_prec: Optional[dict] = 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 get_dim_out() int [source]
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.
- 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
int
the first dimension of the rotation matrix
- abstract get_ntypes() int [source]
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.
- abstract get_rcut() float [source]
Returns the cut-off radius.
- Returns
float
the cut-off radius
Notes
This method must be implemented, as it’s called by other classes.
- init_variables(graph: Graph, graph_def: 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.
- pass_tensors_from_frz_model(*tensors: 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()
.
- abstract prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor] [source]
Compute force and virial.
- Parameters
- Returns
- class deepmd.descriptor.DescrptHybrid(*args, **kwargs)[source]
Bases:
Descriptor
Concate a list of descriptors to form a new descriptor.
- Parameters
- list
list
Build a descriptor from the concatenation of the list of descriptors.
- list
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.
get_nlist_i
(ii)Get the neighbor information of the ii-th descriptor.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(*tensors)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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: Optional[list] = 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.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: Graph, graph_def: 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: Optional[dict] = None) None [source]
Reveive the mixed precision setting.
- Parameters
- mixed_prec
The mixed precision setting used in the embedding net
- get_nlist() Tuple[Tensor, 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
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- get_nlist_i(ii: int) Tuple[Tensor, 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
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- 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.
- pass_tensors_from_frz_model(*tensors: 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
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
- class deepmd.descriptor.DescrptLocFrame(*args, **kwargs)[source]
Bases:
Descriptor
Defines 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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
- Returns
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(*tensors)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- get_nlist() Tuple[Tensor, Tensor, List[int], List[int]] [source]
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
- class deepmd.descriptor.DescrptSeA(*args, **kwargs)[source]
Bases:
DescrptSe
DeepPot-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.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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- 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.
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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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: Optional[dict] = None) None [source]
Reveive the mixed precision setting.
- Parameters
- mixed_prec
The mixed precision setting used in the embedding net
- 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[Tensor, Tensor, List[int], List[int]] [source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- 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.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
- class deepmd.descriptor.DescrptSeAEbd(*args, **kwargs)[source]
Bases:
DescrptSeA
DeepPot-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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- class deepmd.descriptor.DescrptSeAEbdV2(*args, **kwargs)[source]
Bases:
DescrptSeA
A compressible se_a_ebd model.
This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- class deepmd.descriptor.DescrptSeAEf(*args, **kwargs)[source]
Bases:
DescrptSe
Smooth 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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- 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[Tensor, Tensor, List[int], List[int]] [source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
- class deepmd.descriptor.DescrptSeAEfLower(*args, **kwargs)[source]
Bases:
DescrptSeA
Helper class for implementing DescrptSeAEf.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- class deepmd.descriptor.DescrptSeAMask(*args, **kwargs)[source]
Bases:
DescrptSeA
DeepPot-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.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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cutoff radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
None for se_a_mask op
atom_virial
None for se_a_mask op
- class deepmd.descriptor.DescrptSeAtten(*args, **kwargs)[source]
Bases:
DescrptSeA
Smooth version 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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- 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.
- stripped_type_embedding
Whether to strip the type embedding into a separated embedding network. Default value will be True in se_atten_v2 descriptor.
- smooth_type_embedding
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
ValueError
if ntypes is 0.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
build_type_exclude_mask_mixed
(exclude_types, ...)Build the type exclude mask for the attention descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- build_type_exclude_mask_mixed(exclude_types: Set[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor, nei_type_vec: Tensor) 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.Tensor
The 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.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.
- 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: Optional[list] = 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.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: Graph, graph_def: 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
- class deepmd.descriptor.DescrptSeAttenV2(*args, **kwargs)[source]
Bases:
DescrptSeAtten
Smooth 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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- 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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
build_type_exclude_mask_mixed
(exclude_types, ...)Build the type exclude mask for the attention descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- class deepmd.descriptor.DescrptSeR(*args, **kwargs)[source]
Bases:
DescrptSe
DeepPot-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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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
- get_nlist()[source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- 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.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
- class deepmd.descriptor.DescrptSeT(*args, **kwargs)[source]
Bases:
DescrptSe
DeepPot-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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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
- get_nlist() Tuple[Tensor, Tensor, List[int], List[int]] [source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- 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.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
Submodules
deepmd.descriptor.descriptor module
- class deepmd.descriptor.descriptor.Descriptor(*args, **kwargs)[source]
Bases:
PluginVariant
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., rcut_smth=0.5, sel=[50]) >>> type(descript) <class 'deepmd.descriptor.se_a.DescrptSeA'>
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(*tensors)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') 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.
- build_type_exclude_mask(exclude_types: Set[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor) 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.Tensor
The 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[Rafc1ae60e195-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 compute_input_stats(data_coord: List[ndarray], data_box: List[ndarray], data_atype: List[ndarray], natoms_vec: List[ndarray], mesh: List[ndarray], input_dict: Dict[str, List[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.model.model_stat.make_stat_input()
- data_box
list
[np.ndarray
] The box. Can be generated by
deepmd.model.model_stat.make_stat_input()
- data_atype
list
[np.ndarray
] The atom types. Can be generated by
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
deepmd.model.model_stat.make_stat_input()
- mesh
list
[np.ndarray
] The mesh for neighbor searching. Can be generated by
deepmd.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.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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.
- enable_mixed_precision(mixed_prec: Optional[dict] = 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 get_dim_out() int [source]
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.
- 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
int
the first dimension of the rotation matrix
- abstract get_ntypes() int [source]
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.
- abstract get_rcut() float [source]
Returns the cut-off radius.
- Returns
float
the cut-off radius
Notes
This method must be implemented, as it’s called by other classes.
- init_variables(graph: Graph, graph_def: 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.
- pass_tensors_from_frz_model(*tensors: 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()
.
- abstract prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor] [source]
Compute force and virial.
- Parameters
- Returns
deepmd.descriptor.hybrid module
- class deepmd.descriptor.hybrid.DescrptHybrid(*args, **kwargs)[source]
Bases:
Descriptor
Concate a list of descriptors to form a new descriptor.
- Parameters
- list
list
Build a descriptor from the concatenation of the list of descriptors.
- list
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.
get_nlist_i
(ii)Get the neighbor information of the ii-th descriptor.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(*tensors)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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: Optional[list] = 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.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: Graph, graph_def: 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: Optional[dict] = None) None [source]
Reveive the mixed precision setting.
- Parameters
- mixed_prec
The mixed precision setting used in the embedding net
- get_nlist() Tuple[Tensor, 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
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- get_nlist_i(ii: int) Tuple[Tensor, 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
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- 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.
- pass_tensors_from_frz_model(*tensors: 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
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
deepmd.descriptor.loc_frame module
- class deepmd.descriptor.loc_frame.DescrptLocFrame(*args, **kwargs)[source]
Bases:
Descriptor
Defines 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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
- Returns
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(*tensors)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- get_nlist() Tuple[Tensor, Tensor, List[int], List[int]] [source]
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
deepmd.descriptor.se module
- class deepmd.descriptor.se.DescrptSe(*args, **kwargs)[source]
Bases:
Descriptor
A base class for smooth version of descriptors.
Notes
All of these descriptors have an environmental matrix and an embedding network (
deepmd.utils.network.embedding_net()
), so they can share some similiar methods without defining them twice.- Attributes
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- pass_tensors_from_frz_model(descrpt_reshape: Tensor, descrpt_deriv: Tensor, rij: Tensor, nlist: Tensor)[source]
Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
- Parameters
- descrpt_reshape
The passed descrpt_reshape tensor
- descrpt_deriv
The passed descrpt_deriv tensor
- rij
The passed rij tensor
- nlist
The passed nlist tensor
- property precision: DType
Precision of filter network.
deepmd.descriptor.se_a module
- class deepmd.descriptor.se_a.DescrptSeA(*args, **kwargs)[source]
Bases:
DescrptSe
DeepPot-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.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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- 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.
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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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: Optional[dict] = None) None [source]
Reveive the mixed precision setting.
- Parameters
- mixed_prec
The mixed precision setting used in the embedding net
- 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[Tensor, Tensor, List[int], List[int]] [source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') None [source]
Init the embedding net variables with the given dict.
- 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.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
deepmd.descriptor.se_a_ebd module
- class deepmd.descriptor.se_a_ebd.DescrptSeAEbd(*args, **kwargs)[source]
Bases:
DescrptSeA
DeepPot-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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
deepmd.descriptor.se_a_ebd_v2 module
- class deepmd.descriptor.se_a_ebd_v2.DescrptSeAEbdV2(*args, **kwargs)[source]
Bases:
DescrptSeA
A compressible se_a_ebd model.
This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
deepmd.descriptor.se_a_ef module
- class deepmd.descriptor.se_a_ef.DescrptSeAEf(*args, **kwargs)[source]
Bases:
DescrptSe
Smooth 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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- 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[Tensor, Tensor, List[int], List[int]] [source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
- class deepmd.descriptor.se_a_ef.DescrptSeAEfLower(*args, **kwargs)[source]
Bases:
DescrptSeA
Helper class for implementing DescrptSeAEf.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
deepmd.descriptor.se_a_mask module
- class deepmd.descriptor.se_a_mask.DescrptSeAMask(*args, **kwargs)[source]
Bases:
DescrptSeA
DeepPot-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.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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cutoff radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
None for se_a_mask op
atom_virial
None for se_a_mask op
deepmd.descriptor.se_atten module
- class deepmd.descriptor.se_atten.DescrptSeAtten(*args, **kwargs)[source]
Bases:
DescrptSeA
Smooth version 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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- 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.
- stripped_type_embedding
Whether to strip the type embedding into a separated embedding network. Default value will be True in se_atten_v2 descriptor.
- smooth_type_embedding
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
ValueError
if ntypes is 0.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
build_type_exclude_mask_mixed
(exclude_types, ...)Build the type exclude mask for the attention descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- build_type_exclude_mask_mixed(exclude_types: Set[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor, nei_type_vec: Tensor) 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.Tensor
The 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.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.
- 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: Optional[list] = 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.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: Graph, graph_def: 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
deepmd.descriptor.se_atten_v2 module
- class deepmd.descriptor.se_atten_v2.DescrptSeAttenV2(*args, **kwargs)[source]
Bases:
DescrptSeAtten
Smooth 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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- 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.
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
build_type_exclude_mask_mixed
(exclude_types, ...)Build the type exclude mask for the attention descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
get_dim_out
()Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
get_nlist
()Returns neighbor information.
get_ntypes
()Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_rot_mat
()Get rotational matrix.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
deepmd.descriptor.se_r module
- class deepmd.descriptor.se_r.DescrptSeR(*args, **kwargs)[source]
Bases:
DescrptSe
DeepPot-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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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
- get_nlist()[source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- 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.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial
deepmd.descriptor.se_t module
- class deepmd.descriptor.se_t.DescrptSeT(*args, **kwargs)[source]
Bases:
DescrptSe
DeepPot-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”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
- precision
The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- Attributes
explicit_ntypes
Explicit ntypes with type embedding.
precision
Precision of filter network.
Methods
build
(coord_, atype_, natoms, box_, mesh, ...)Build the computational graph for the descriptor.
build_type_exclude_mask
(exclude_types, ...)Build the type exclude mask for the descriptor.
compute_input_stats
(data_coord, data_box, ...)Compute the statisitcs (avg and std) of the training data.
enable_compression
(min_nbor_dist, graph, ...)Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
enable_mixed_precision
([mixed_prec])Reveive the mixed precision setting.
Returns the output dimension of this descriptor.
get_dim_rot_mat_1
()Returns the first dimension of the rotation matrix.
Returns neighbor information.
Returns the number of atom types.
get_rcut
()Returns the cut-off radius.
get_tensor_names
([suffix])Get names of tensors.
init_variables
(graph, graph_def[, suffix])Init the embedding net variables with the given dict.
merge_input_stats
(stat_dict)Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
pass_tensors_from_frz_model
(descrpt_reshape, ...)Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
prod_force_virial
(atom_ener, natoms)Compute force and virial.
register
(key)Register a descriptor plugin.
update_sel
(global_jdata, local_jdata)Update the selection and perform neighbor statistics.
get_class_by_input
- build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') 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
descriptor
The output descriptor
- 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.model.make_stat_input
- data_box
The box. Can be generated by deepmd.model.make_stat_input
- data_atype
The atom types. Can be generated by deepmd.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.model.make_stat_input
- mesh
The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
- input_dict
Dictionary for additional input
- **kwargs
Additional keyword arguments.
- enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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
- get_nlist() Tuple[Tensor, Tensor, List[int], List[int]] [source]
Returns neighbor information.
- Returns
nlist
Neighbor list
rij
The relative distance between the neighbor and the center atom.
sel_a
The number of neighbors with full information
sel_r
The number of neighbors with only radial information
- 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.
- prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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
force
The force on atoms
virial
The total virial
atom_virial
The atomic virial