deepmd.jax.descriptor#
Submodules#
Classes#
Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model. | |
The unit operation of a native model. | |
Concate a list of descriptors to form a new descriptor. | |
Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model. | |
DeepPot-SE constructed from all information (both angular and radial) of | |
DeepPot-SE_R constructed from only the radial information of atomic configurations. | |
DeepPot-SE constructed from all information (both angular and radial) of atomic | |
Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input. |
Package Contents#
- class deepmd.jax.descriptor.DescrptDPA1(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, scaling_factor=1.0, normalize: bool = True, temperature: float | None = None, trainable_ln: bool = True, ln_eps: float | None = 1e-05, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, spin: Any | None = None, stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None)[source]#
Bases:
deepmd.dpmodel.descriptor.dpa1.DescrptDPA1
Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.
This descriptor, \(\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}\), is given by
\[\mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<,\]where \(\hat{\mathcal{G}}^i\) represents the embedding matrix:math:mathcal{G}^i after additional self-attention mechanism and \(\mathcal{R}^i\) is defined by the full case in the se_e2_a descriptor. Note that we obtain \(\mathcal{G}^i\) using the type embedding method by default in this descriptor.
To perform the self-attention mechanism, the queries \(\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), keys \(\mathcal{K}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), and values \(\mathcal{V}^{i,l} \in \mathbb{R}^{N_c\times d_v}\) are first obtained:
\[\left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]\[\left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]\[\left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]where \(Q_{l}\), \(K_{l}\), \(V_{l}\) represent three trainable linear transformations that output the queries and keys of dimension \(d_k\) and values of dimension \(d_v\), and \(l\) is the index of the attention layer. The input embedding matrix to the attention layers, denoted by \(\mathcal{G}^{i,0}\), is chosen as the two-body embedding matrix.
Then the scaled dot-product attention method is adopted:
\[A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l},\]where \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) \in \mathbb{R}^{N_c\times N_c}\) is attention weights. In the original attention method, one typically has \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}\right)=\mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)\), with \(\sqrt{d_{k}}\) being the normalization temperature. This is slightly modified to incorporate the angular information:
\[\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) = \mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right) \odot \hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T},\]- where \(\hat{\mathcal{R}}^{i} \in \mathbb{R}^{N_c\times 3}\) denotes normalized relative coordinates,
\(\hat{\mathcal{R}}^{i}_{j} = \frac{\boldsymbol{r}_{ij}}{\lVert \boldsymbol{r}_{ij} \lVert}\) and \(\odot\) means element-wise multiplication.
- Then layer normalization is added in a residual way to finally obtain the self-attention local embedding matrix
\(\hat{\mathcal{G}}^{i} = \mathcal{G}^{i,L_a}\) after \(L_a\) attention layers:[^1]
\[\mathcal{G}^{i,l} = \mathcal{G}^{i,l-1} + \mathrm{LayerNorm}(A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})).\]- Parameters:
- rcut: float
The cut-off radius \(r_c\)
- rcut_smth: float
From where the environment matrix should be smoothed \(r_s\)
- sel
list
[int
],int
list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius
- ntypes
int
Number of element types
- neuron
list
[int
] Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron: int
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- tebd_dim: int
Dimension of the type embedding
- tebd_input_mode: str
The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the radial embedding network output.
- resnet_dt: bool
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable: bool
If the weights of this descriptors are trainable.
- trainable_ln: bool
Whether to use trainable shift and scale weights in layer normalization.
- ln_eps: float, Optional
The epsilon value for layer normalization.
- type_one_side: bool
If ‘False’, type embeddings of both neighbor and central atoms are considered. If ‘True’, only type embeddings of neighbor atoms are considered. Default is ‘False’.
- attn: int
Hidden dimension of the attention vectors
- attn_layer: int
Number of attention layers
- attn_dotr: bool
If dot the angular gate to the attention weights
- attn_mask: bool
(Only support False to keep consistent with other backend references.) (Not used in this version. True option is not implemented.) If mask the diagonal of attention weights
- 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.
- env_protection: float
Protection parameter to prevent division by zero errors during environment matrix calculations.
- set_davg_zero: bool
Set the shift of embedding net input to zero.
- activation_function: str
The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
- precision: str
The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
- scaling_factor: float
The scaling factor of normalization in calculations of attention weights. If temperature is None, the scaling of attention weights is (N_dim * scaling_factor)**0.5
- normalize: bool
Whether to normalize the hidden vectors in attention weights calculation.
- temperature: float
If not None, the scaling of attention weights is temperature itself.
- smooth_type_embedding: bool
Whether to use smooth process in attention weights calculation.
- concat_output_tebd: bool
Whether to concat type embedding at the output of the descriptor.
- stripped_type_embedding: bool, Optional
(Deprecated, kept only for compatibility.) Whether to strip the type embedding into a separate embedding network. Setting this parameter to True is equivalent to setting tebd_input_mode to ‘strip’. Setting it to False is equivalent to setting tebd_input_mode to ‘concat’. The default value is None, which means the tebd_input_mode setting will be used instead.
- use_econf_tebd: bool, Optional
Whether to use electronic configuration type embedding.
- use_tebd_biasbool,
Optional
Whether to use bias in the type embedding layer.
- type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
- spin
(Only support None to keep consistent with other backend references.) (Not used in this version. Not-none option is not implemented.) The old implementation of deepspin.
References
[1]Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. 2022. DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. arXiv preprint arXiv:2208.08236.
- class deepmd.jax.descriptor.DescrptDPA2(ntypes: int, repinit: RepinitArgs | dict, repformer: RepformerArgs | dict, concat_output_tebd: bool = True, precision: str = 'float64', smooth: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None)[source]#
Bases:
deepmd.dpmodel.descriptor.dpa2.DescrptDPA2
The unit operation of a native model.
- class deepmd.jax.descriptor.DescrptHybrid(list: DescrptHybrid.__init__.list[deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor | dict[str, Any]], type_map: DescrptHybrid.__init__.list[str] | None = None, ntypes: int | None = None)[source]#
Bases:
deepmd.dpmodel.descriptor.hybrid.DescrptHybrid
Concate a list of descriptors to form a new descriptor.
- Parameters:
- class deepmd.jax.descriptor.DescrptSeAttenV2(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, scaling_factor=1.0, normalize: bool = True, temperature: float | None = None, trainable_ln: bool = True, ln_eps: float | None = 1e-05, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, spin: Any | None = None, stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None)[source]#
Bases:
deepmd.jax.descriptor.dpa1.DescrptDPA1
,deepmd.dpmodel.descriptor.se_atten_v2.DescrptSeAttenV2
Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.
This descriptor, \(\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}\), is given by
\[\mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<,\]where \(\hat{\mathcal{G}}^i\) represents the embedding matrix:math:mathcal{G}^i after additional self-attention mechanism and \(\mathcal{R}^i\) is defined by the full case in the se_e2_a descriptor. Note that we obtain \(\mathcal{G}^i\) using the type embedding method by default in this descriptor.
To perform the self-attention mechanism, the queries \(\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), keys \(\mathcal{K}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), and values \(\mathcal{V}^{i,l} \in \mathbb{R}^{N_c\times d_v}\) are first obtained:
\[\left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]\[\left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]\[\left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]where \(Q_{l}\), \(K_{l}\), \(V_{l}\) represent three trainable linear transformations that output the queries and keys of dimension \(d_k\) and values of dimension \(d_v\), and \(l\) is the index of the attention layer. The input embedding matrix to the attention layers, denoted by \(\mathcal{G}^{i,0}\), is chosen as the two-body embedding matrix.
Then the scaled dot-product attention method is adopted:
\[A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l},\]where \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) \in \mathbb{R}^{N_c\times N_c}\) is attention weights. In the original attention method, one typically has \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}\right)=\mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)\), with \(\sqrt{d_{k}}\) being the normalization temperature. This is slightly modified to incorporate the angular information:
\[\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) = \mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right) \odot \hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T},\]- where \(\hat{\mathcal{R}}^{i} \in \mathbb{R}^{N_c\times 3}\) denotes normalized relative coordinates,
\(\hat{\mathcal{R}}^{i}_{j} = \frac{\boldsymbol{r}_{ij}}{\lVert \boldsymbol{r}_{ij} \lVert}\) and \(\odot\) means element-wise multiplication.
- Then layer normalization is added in a residual way to finally obtain the self-attention local embedding matrix
\(\hat{\mathcal{G}}^{i} = \mathcal{G}^{i,L_a}\) after \(L_a\) attention layers:[^1]
\[\mathcal{G}^{i,l} = \mathcal{G}^{i,l-1} + \mathrm{LayerNorm}(A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})).\]- Parameters:
- rcut: float
The cut-off radius \(r_c\)
- rcut_smth: float
From where the environment matrix should be smoothed \(r_s\)
- sel
list
[int
],int
list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius
- ntypes
int
Number of element types
- neuron
list
[int
] Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron: int
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- tebd_dim: int
Dimension of the type embedding
- tebd_input_mode: str
The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the radial embedding network output.
- resnet_dt: bool
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable: bool
If the weights of this descriptors are trainable.
- trainable_ln: bool
Whether to use trainable shift and scale weights in layer normalization.
- ln_eps: float, Optional
The epsilon value for layer normalization.
- type_one_side: bool
If ‘False’, type embeddings of both neighbor and central atoms are considered. If ‘True’, only type embeddings of neighbor atoms are considered. Default is ‘False’.
- attn: int
Hidden dimension of the attention vectors
- attn_layer: int
Number of attention layers
- attn_dotr: bool
If dot the angular gate to the attention weights
- attn_mask: bool
(Only support False to keep consistent with other backend references.) (Not used in this version. True option is not implemented.) If mask the diagonal of attention weights
- 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.
- env_protection: float
Protection parameter to prevent division by zero errors during environment matrix calculations.
- set_davg_zero: bool
Set the shift of embedding net input to zero.
- activation_function: str
The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
- precision: str
The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
- scaling_factor: float
The scaling factor of normalization in calculations of attention weights. If temperature is None, the scaling of attention weights is (N_dim * scaling_factor)**0.5
- normalize: bool
Whether to normalize the hidden vectors in attention weights calculation.
- temperature: float
If not None, the scaling of attention weights is temperature itself.
- smooth_type_embedding: bool
Whether to use smooth process in attention weights calculation.
- concat_output_tebd: bool
Whether to concat type embedding at the output of the descriptor.
- stripped_type_embedding: bool, Optional
(Deprecated, kept only for compatibility.) Whether to strip the type embedding into a separate embedding network. Setting this parameter to True is equivalent to setting tebd_input_mode to ‘strip’. Setting it to False is equivalent to setting tebd_input_mode to ‘concat’. The default value is None, which means the tebd_input_mode setting will be used instead.
- use_econf_tebd: bool, Optional
Whether to use electronic configuration type embedding.
- use_tebd_biasbool,
Optional
Whether to use bias in the type embedding layer.
- type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
- spin
(Only support None to keep consistent with other backend references.) (Not used in this version. Not-none option is not implemented.) The old implementation of deepspin.
References
[1]Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. 2022. DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. arXiv preprint arXiv:2208.08236.
- class deepmd.jax.descriptor.DescrptSeA(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = True, exclude_types: list[list[int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, spin: Any | None = None, type_map: list[str] | None = None, ntypes: int | None = None, seed: int | list[int] | None = None)[source]#
Bases:
deepmd.dpmodel.descriptor.se_e2_a.DescrptSeAArrayAPI
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.tf.utils.network.embedding_net()
.- Parameters:
- rcut
The cut-off radius \(r_c\)
- rcut_smth
From where the environment matrix should be smoothed \(r_s\)
- sel
list
[int
] sel[i] specifies the maxmum number of type i atoms in the cut-off radius
- neuron
list
[int
] Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
- axis_neuron
Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- 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.
- env_protection: float
Protection parameter to prevent division by zero errors during environment matrix calculations.
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
- precision
The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
- spin
The deepspin object.
- type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
- ntypes
int
Number of element types. Not used in this descriptor, only to be compat with input.
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.
- class deepmd.jax.descriptor.DescrptSeR(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = True, exclude_types: list[list[int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, spin: Any | None = None, type_map: list[str] | None = None, ntypes: int | None = None, seed: int | list[int] | None = None)[source]#
Bases:
deepmd.dpmodel.descriptor.se_r.DescrptSeR
DeepPot-SE_R constructed from only the radial information of atomic configurations.
- 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}\)
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- trainable
If the weights of embedding net are trainable.
- 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 “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
- precision
The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
- spin
The deepspin object.
- type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
- ntypes
int
Number of element types. Not used in this descriptor, only to be compat with input.
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.
- class deepmd.jax.descriptor.DescrptSeT(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, set_davg_zero: bool = False, activation_function: str = 'tanh', env_protection: float = 0.0, exclude_types: list[tuple[int, int]] = [], precision: str = DEFAULT_PRECISION, trainable: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, ntypes: int | None = None)[source]#
Bases:
deepmd.dpmodel.descriptor.se_t.DescrptSeT
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
float
The cut-off radius
- rcut_smth
float
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_dtbool
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- set_davg_zerobool
Set the shift of embedding net input to zero.
- activation_function
str
The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
- env_protection
float
Protection parameter to prevent division by zero errors during environment matrix calculations.
- 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.
- precision
str
The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
- trainablebool
If the weights of embedding net are trainable.
- seed
int
,Optional
Random seed for initializing the network parameters.
- type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
- ntypes
int
Number of element types. Not used in this descriptor, only to be compat with input.
- rcut
- class deepmd.jax.descriptor.DescrptSeTTebd(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, set_davg_zero: bool = True, activation_function: str = 'tanh', env_protection: float = 0.0, exclude_types: list[tuple[int, int]] = [], precision: str = 'float64', trainable: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias=False, smooth: bool = True)[source]#
Bases:
deepmd.dpmodel.descriptor.se_t_tebd.DescrptSeTTebd
Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.
- Parameters:
- rcut
The cut-off radius
- rcut_smth
From where the environment matrix should be smoothed
- sel
Union
[list
[int
],int
] list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius
- ntypes
int
Number of element types
- neuron
list
[int
] Number of neurons in each hidden layers of the embedding net
- tebd_dim
int
Dimension of the type embedding
- tebd_input_mode
str
The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed angular information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the angular embedding network output.
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- set_davg_zero
Set the shift of embedding net input to zero.
- activation_function
The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
- env_protection: float
Protection parameter to prevent division by zero errors during environment matrix calculations.
- exclude_types
list
[tuple
[int
,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.
- precision
The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
- concat_output_tebd: bool
Whether to concat type embedding at the output of the descriptor.
- use_econf_tebd: bool, Optional
Whether to use electronic configuration type embedding.
- use_tebd_biasbool,
Optional
Whether to use bias in the type embedding layer.
- smooth: bool
Whether to use smooth process in calculation.