deepmd.dpmodel.descriptor.se_t_tebd

deepmd.dpmodel.descriptor.se_t_tebd#

Classes#

`DescrptSeTTebd`	Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.
`DescrptBlockSeTTebd`	The three-body descriptor block with type embedding.

Module Contents#

class deepmd.dpmodel.descriptor.se_t_tebd.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: bool = False, smooth: bool = True)[source]#

Bases: deepmd.dpmodel.NativeOP, deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor

Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.

The descriptor \(\mathcal{D}^i \in \mathbb{R}^{M}\) is given by

\[\mathcal{D}^i = \frac{1}{N_c^2} \sum_{j,k} \mathcal{N}(\cos\theta_{jik}, \mathcal{T}_j, \mathcal{T}_k),\]

where \(\theta_{jik}\) is the angle between neighbors \(j\) and \(k\) around the central atom \(i\), \(\mathcal{T}_j\) and \(\mathcal{T}_k\) are the type embeddings of atoms \(j\) and \(k\), and \(\mathcal{N}\) is the embedding network.

The cosine of the angle is computed from the normalized relative coordinates:

\[\cos\theta_{jik} = \frac{\boldsymbol{r}_{ij} \cdot \boldsymbol{r}_{ik}}{|\boldsymbol{r}_{ij}| |\boldsymbol{r}_{ik}|}.\]

The type embedding can be incorporated in two modes:

“concat”: Concatenate \([\cos\theta_{jik}, \mathcal{T}_j, \mathcal{T}_k]\) as input to the embedding network.
“strip”: Use separate embedding networks for \(\cos\theta_{jik}\) and \([\mathcal{T}_j, \mathcal{T}_k]\), then combine their outputs multiplicatively.

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 “silu”, “softplus”, “sigmoid”, “none”, “gelu_tf”, “relu6”, “tanh”, “linear”, “gelu”, “relu”, “silut”.
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 “float32”, “bfloat16”, “default”, “float64”, “float16”.
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.

_update_sel_cls[source]#

se_ttebd[source]#

use_econf_tebd = False[source]#

type_map = None[source]#

smooth = True[source]#

type_embedding[source]#

tebd_dim = 8[source]#

concat_output_tebd = True[source]#

trainable = True[source]#

precision = 'float64'[source]#

get_rcut() → float[source]#: Returns the cut-off radius.

get_rcut_smth() → float[source]#: Returns the radius where the neighbor information starts to smoothly decay to 0.

get_nsel() → int[source]#: Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int][source]#: Returns the number of selected atoms for each type.

get_ntypes() → int[source]#: Returns the number of element types.

get_type_map() → list[str][source]#: Get the name to each type of atoms.

get_dim_out() → int[source]#: Returns the output dimension.

get_dim_emb() → int[source]#: Returns the embedding dimension of g2.

mixed_types() → bool[source]#

If true, the descriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.

If false, the descriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.

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: DescrptSeTTebd, shared_level: int, resume: bool = 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.

property dim_out: int[source]#

property dim_emb: int[source]#

compute_input_stats(merged: collections.abc.Callable[[], list[dict]] | list[dict], path: deepmd.utils.path.DPPath | None = None) → None[source]#

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters:

mergedUnion[Callable[[], list[dict]], list[dict]]

list[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], list[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

set_stat_mean_and_stddev(mean: deepmd.dpmodel.array_api.Array, stddev: deepmd.dpmodel.array_api.Array) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#: Get mean and stddev for descriptor.

change_type_map(type_map: list[str], model_with_new_type_stat: DescrptSeTTebd | None = 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.

call(coord_ext: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][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) → DescrptSeTTebd[source]#: Deserialize from dict.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][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_t_tebd.DescrptBlockSeTTebd(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, tebd_input_mode: str = 'concat', set_davg_zero: bool = True, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, seed: int | list[int] | None = None, trainable: bool = True)[source]#

Bases: deepmd.dpmodel.NativeOP, deepmd.dpmodel.descriptor.descriptor.DescriptorBlock

The three-body descriptor block with type embedding.

This block computes an embedding using angles between two neighboring atoms and type embeddings. The descriptor is computed as:

\[\mathcal{D}^i = \frac{1}{N_c^2} \sum_{j,k} \mathcal{N}(\cos\theta_{jik}, \mathcal{T}_j, \mathcal{T}_k),\]

where \(\theta_{jik}\) is the angle between neighbors \(j\) and \(k\) around the central atom \(i\), \(\mathcal{T}_j\) and \(\mathcal{T}_k\) are the type embeddings of atoms \(j\) and \(k\).

The cosine of the angle is computed from the normalized relative coordinates:

\[\cos\theta_{jik} = \frac{\boldsymbol{r}_{ij} \cdot \boldsymbol{r}_{ik}}{|\boldsymbol{r}_{ij}| |\boldsymbol{r}_{ik}|}.\]

Parameters:

rcutfloat: The cut-off radius.
rcut_smthfloat: Where to start smoothing.
selUnion[list[int], int]: Maximally possible number of selected neighbors.
ntypesint: Number of element types.
neuronlist[int], optional: Number of neurons in each hidden layer of the embedding net.
tebd_dimint, optional: Dimension of the type embedding.
tebd_input_modestr, optional: The input mode of the type embedding. Supported modes are [“concat”, “strip”].
set_davg_zerobool, optional: Set the shift of embedding net input to zero.
activation_functionstr, optional: The activation function in the embedding net.
precisionstr, optional: The precision of the embedding net parameters.
resnet_dtbool, optional: Time-step dt in the resnet construction.
exclude_typeslist[tuple[int, int]], optional: The excluded pairs of types which have no interaction.
env_protectionfloat, optional: Protection parameter to prevent division by zero.
smoothbool, optional: Whether to use smoothness.
seedint, optional: Random seed for parameter initialization.
trainablebool, optional: If the parameters are trainable.

rcut[source]#

rcut_smth[source]#

neuron = [25, 50, 100][source]#

filter_neuron = [25, 50, 100][source]#

tebd_dim = 8[source]#

tebd_input_mode = 'concat'[source]#

set_davg_zero = True[source]#

activation_function = 'tanh'[source]#

precision = 'float64'[source]#

resnet_dt = False[source]#

env_protection = 0.0[source]#

seed = None[source]#

smooth = True[source]#

ntypes[source]#

sel[source]#

sec[source]#

split_sel[source]#

nnei[source]#

ndescrpt[source]#

tebd_dim_input = 16[source]#

embeddings[source]#

env_mat[source]#

mean[source]#

stddev[source]#

orig_sel[source]#

get_rcut() → float[source]#: Returns the cut-off radius.

get_rcut_smth() → float[source]#: Returns the radius where the neighbor information starts to smoothly decay to 0.

get_nsel() → int[source]#: Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int][source]#: Returns the number of selected atoms for each type.

get_ntypes() → int[source]#: Returns the number of element types.

get_dim_in() → int[source]#: Returns the input dimension.

get_dim_out() → int[source]#: Returns the output dimension.

get_dim_emb() → int[source]#: Returns the output dimension of embedding.

__setitem__(key: str, value: deepmd.dpmodel.array_api.Array) → None[source]#

__getitem__(key: str) → deepmd.dpmodel.array_api.Array[source]#

mixed_types() → bool[source]#

If true, the descriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.

If false, the descriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.

get_env_protection() → float[source]#: Returns the protection of building environment matrix.

property dim_out: int[source]#: Returns the output dimension of this descriptor.

property dim_in: int[source]#: Returns the atomic input dimension of this descriptor.

property dim_emb: int[source]#: Returns the output dimension of embedding.

compute_input_stats(merged: collections.abc.Callable[[], list[dict]] | list[dict], path: deepmd.utils.path.DPPath | None = None) → None[source]#

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters:

mergedUnion[Callable[[], list[dict]], list[dict]]

list[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: paddle.Tensor originating from the i-th data system.
Callable[[], list[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

get_stats() → dict[str, deepmd.utils.env_mat_stat.StatItem][source]#: Get the statistics of the descriptor.

reinit_exclude(exclude_types: list[tuple[int, int]] = []) → None[source]#

cal_g(ss: deepmd.dpmodel.array_api.Array, embedding_idx: int) → deepmd.dpmodel.array_api.Array[source]#

cal_g_strip(ss: deepmd.dpmodel.array_api.Array, embedding_idx: int) → deepmd.dpmodel.array_api.Array[source]#

call(nlist: deepmd.dpmodel.array_api.Array, coord_ext: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array, atype_embd_ext: deepmd.dpmodel.array_api.Array | None = None, mapping: deepmd.dpmodel.array_api.Array | None = None, type_embedding: deepmd.dpmodel.array_api.Array | None = None) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#: Forward pass in NumPy implementation.

has_message_passing() → bool[source]#: Returns whether the descriptor block has message passing.

need_sorted_nlist_for_lower() → bool[source]#: Returns whether the descriptor block needs sorted nlist when using forward_lower.

serialize() → dict[source]#: Serialize the descriptor to dict.

classmethod deserialize(data: dict) → DescrptSeTTebd[source]#: Deserialize from dict.