deepmd.jax.descriptor#

Submodules#

Classes#

DescrptDPA1

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

DescrptDPA2

The unit operation of a native model.

DescrptHybrid

Concate a list of descriptors to form a new descriptor.

DescrptSeAttenV2

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

DescrptSeA

DeepPot-SE constructed from all information (both angular and radial) of

DescrptSeR

DeepPot-SE_R constructed from only the radial information of atomic configurations.

DescrptSeT

DeepPot-SE constructed from all information (both angular and radial) of atomic

DescrptSeTTebd

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\)

sellist[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

ntypesint

Number of element types

neuronlist[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_typeslist[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.

__setattr__(name: str, value: Any) None[source]#
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.

__setattr__(name: str, value: Any) None[source]#
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:
listlistlist[Union[BaseDescriptor, dict[str, Any]]]

Build a descriptor from the concatenation of the list of descriptors. The descriptor can be either an object or a dictionary.

__setattr__(name: str, value: Any) None[source]#
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\)

sellist[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

ntypesint

Number of element types

neuronlist[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_typeslist[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\)

sellist[int]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[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_typeslist[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.

ntypesint

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.

__setattr__(name: str, value: Any) None[source]#
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\)

sellist[int]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[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_typeslist[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.

ntypesint

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.

__setattr__(name: str, value: Any) None[source]#
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:
rcutfloat

The cut-off radius

rcut_smthfloat

From where the environment matrix should be smoothed

sellist[int]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[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_functionstr

The activation function in the embedding net. Supported options are “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.

env_protectionfloat

Protection parameter to prevent division by zero errors during environment matrix calculations.

exclude_typeslist[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.

precisionstr

The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.

trainablebool

If the weights of embedding net are trainable.

seedint, 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.

ntypesint

Number of element types. Not used in this descriptor, only to be compat with input.

__setattr__(name: str, value: Any) None[source]#
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

selUnion[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

ntypesint

Number of element types

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

tebd_dimint

Dimension of the type embedding

tebd_input_modestr

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_typeslist[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.

__setattr__(name: str, value: Any) None[source]#