deepmd.dpmodel.descriptor.se_e2_a#

Classes#

DescrptSeA

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

DescrptSeAArrayAPI

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

Module Contents#

class deepmd.dpmodel.descriptor.se_e2_a.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.NativeOP, deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor

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 “gelu_tf”, “relu6”, “relu”, “linear”, “gelu”, “none”, “softplus”, “tanh”, “sigmoid”.

precision

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

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.

rcut[source]#
rcut_smth[source]#
sel[source]#
ntypes[source]#
neuron = [24, 48, 96][source]#
axis_neuron = 8[source]#
resnet_dt = False[source]#
trainable = True[source]#
type_one_side = True[source]#
env_protection = 0.0[source]#
set_davg_zero = False[source]#
activation_function = 'tanh'[source]#
precision = 'float64'[source]#
spin = None[source]#
type_map = None[source]#
embeddings[source]#
env_mat[source]#
nnei[source]#
davg[source]#
dstd[source]#
orig_sel[source]#
sel_cumsum[source]#
__setitem__(key, value) None[source]#
__getitem__(key)[source]#
property dim_out[source]#

Returns the output dimension of this descriptor.

get_dim_out()[source]#

Returns the output dimension of this descriptor.

get_dim_emb()[source]#

Returns the embedding (g2) dimension of this descriptor.

get_rcut()[source]#

Returns cutoff radius.

get_rcut_smth() float[source]#

Returns the radius where the neighbor information starts to smoothly decay to 0.

get_sel()[source]#

Returns cutoff radius.

mixed_types() bool[source]#

Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

has_message_passing() bool[source]#

Returns whether the descriptor has message passing.

need_sorted_nlist_for_lower() bool[source]#

Returns whether the descriptor needs sorted nlist when using forward_lower.

get_env_protection() float[source]#

Returns the protection of building environment matrix.

abstractmethod share_params(base_class, shared_level, resume=False) NoReturn[source]#

Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.

abstractmethod change_type_map(type_map: list[str], model_with_new_type_stat=None) None[source]#

Change the type related params to new ones, according to type_map and the original one in the model. If there are new types in type_map, statistics will be updated accordingly to model_with_new_type_stat for these new types.

get_ntypes() int[source]#

Returns the number of element types.

get_type_map() list[str][source]#

Get the name to each type of atoms.

abstractmethod compute_input_stats(merged: list[dict], path: deepmd.utils.path.DPPath | None = None) NoReturn[source]#

Update mean and stddev for descriptor elements.

set_stat_mean_and_stddev(mean: numpy.ndarray, stddev: numpy.ndarray) None[source]#

Update mean and stddev for descriptor.

get_stat_mean_and_stddev() tuple[numpy.ndarray, numpy.ndarray][source]#

Get mean and stddev for descriptor.

cal_g(ss, embedding_idx)[source]#
reinit_exclude(exclude_types: list[tuple[int, int]] = []) None[source]#
call(coord_ext, atype_ext, nlist, mapping: numpy.ndarray | None = None)[source]#

Compute the descriptor.

Parameters:
coord_ext

The extended coordinates of atoms. shape: nf x (nallx3)

atype_ext

The extended aotm types. shape: nf x nall

nlist

The neighbor list. shape: nf x nloc x nnei

mapping

The index mapping from extended to local region. not used by this descriptor.

Returns:
descriptor

The descriptor. shape: nf x nloc x (ng x axis_neuron)

gr

The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3

g2

The rotationally invariant pair-partical representation. this descriptor returns None

h2

The rotationally equivariant pair-partical representation. this descriptor returns None

sw

The smooth switch function.

serialize() dict[source]#

Serialize the descriptor to dict.

classmethod deserialize(data: dict) DescrptSeA[source]#

Deserialize from dict.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) tuple[dict, float | None][source]#

Update the selection and perform neighbor statistics.

Parameters:
train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:
dict

The updated local data

float

The minimum distance between two atoms

class deepmd.dpmodel.descriptor.se_e2_a.DescrptSeAArrayAPI(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: 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.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 “gelu_tf”, “relu6”, “relu”, “linear”, “gelu”, “none”, “softplus”, “tanh”, “sigmoid”.

precision

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

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.

call(coord_ext, atype_ext, nlist, mapping: numpy.ndarray | None = None)[source]#

Compute the descriptor.

Parameters:
coord_ext

The extended coordinates of atoms. shape: nf x (nallx3)

atype_ext

The extended aotm types. shape: nf x nall

nlist

The neighbor list. shape: nf x nloc x nnei

mapping

The index mapping from extended to local region. not used by this descriptor.

Returns:
descriptor

The descriptor. shape: nf x nloc x (ng x axis_neuron)

gr

The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3

g2

The rotationally invariant pair-partical representation. this descriptor returns None

h2

The rotationally equivariant pair-partical representation. this descriptor returns None

sw

The smooth switch function.