deepmd.dpmodel.descriptor

deepmd.dpmodel.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.
`DescrptDPA3`	The DPA-3 descriptor.
`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.

Functions#

make_base_descriptor(t_tensor[, fwd_method_name])

Make the base class for the descriptor.

Package Contents#

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

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 “sigmoid”, “linear”, “gelu”, “none”, “silu”, “relu”, “gelu_tf”, “softplus”, “relu6”, “tanh”, “silut”.
precision: str: The precision of the embedding net parameters. Supported options are “float64”, “default”, “float16”, “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.

se_atten#

use_econf_tebd = False#

use_tebd_bias = False#

type_map = None#

type_embedding#

tebd_dim = 8#

concat_output_tebd = True#

trainable = True#

precision = 'float64'#

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

property dim_out#

property dim_emb#

compute_input_stats(merged: Callable[[], list[dict]] | list[dict], path: deepmd.utils.path.DPPath | 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: 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.

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.

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) → DescrptDPA1[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.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.NativeOP, deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor

The unit operation of a native model.

repinit_args#

repformer_args#

repinit#

use_three_body = False#

repformers#

rcsl_list#

rcut_list#

nsel_list#

use_econf_tebd = False#

use_tebd_bias = False#

type_map = None#

type_embedding#

concat_output_tebd = True#

precision = 'float64'#

smooth = True#

exclude_types = []#

env_protection = 0.0#

trainable = True#

add_tebd_to_repinit_out = False#

repinit_out_dim#

tebd_transform = None#

tebd_dim = 8#

rcut#

ntypes#

sel#

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 of this descriptor.

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

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

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.

property dim_out#

property dim_emb#: Returns the embedding dimension g2.

compute_input_stats(merged: Callable[[], list[dict]] | list[dict], path: deepmd.utils.path.DPPath | 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: list[numpy.ndarray], stddev: list[numpy.ndarray]) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[numpy.ndarray], list[numpy.ndarray]][source]#: Get mean and stddev for descriptor.

call(coord_ext: numpy.ndarray, atype_ext: numpy.ndarray, nlist: numpy.ndarray, 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, maps extended region index to local region.

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. shape: nf x nloc x nnei x ng
h2: The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw: The smooth switch function. shape: nf x nloc x nnei

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

classmethod deserialize(data: dict) → DescrptDPA2[source]#

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

BD: The deserialized descriptor

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.DescrptDPA3(ntypes: int, repflow: RepFlowArgs | dict, concat_output_tebd: bool = False, activation_function: str = 'silu', precision: str = 'float64', exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, type_map: list[str] | None = None)[source]#

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

The DPA-3 descriptor.

Parameters:

repflowUnion[RepFlowArgs, dict]: The arguments used to initialize the repflow block, see docstr in RepFlowArgs for details information.
concat_output_tebdbool, optional: Whether to concat type embedding at the output of the descriptor.
activation_functionstr, optional: The activation function in the embedding net.
precisionstr, optional: The precision of the embedding net parameters.
exclude_typeslist[list[int]], optional: 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_protectionfloat, optional: Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection.
trainablebool, optional: If the parameters are trainable.
seedint, optional: Random seed for parameter initialization.
use_econf_tebdbool, Optional: Whether to use electronic configuration type embedding.
use_tebd_biasbool, Optional: Whether to use bias in the type embedding layer.
use_loc_mappingbool, Optional: Whether to use local atom index mapping in training or non-parallel inference. When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.
type_maplist[str], Optional: A list of strings. Give the name to each type of atoms.

repflow_args#

activation_function = 'silu'#

repflows#

use_econf_tebd = False#

use_tebd_bias = False#

use_loc_mapping = True#

type_map = None#

tebd_dim = 128#

type_embedding#

concat_output_tebd = False#

precision = 'float64'#

exclude_types = []#

env_protection = 0.0#

trainable = True#

rcut#

rcut_smth#

sel#

ntypes#

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 of this descriptor.

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

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, shared_level, resume=False) → None[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.

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.

property dim_out#

property dim_emb#: Returns the embedding dimension g2.

compute_input_stats(merged: list[dict], path: deepmd.utils.path.DPPath | None = None)[source]#: Update mean and stddev for descriptor elements.

set_stat_mean_and_stddev(mean: list[numpy.ndarray], stddev: list[numpy.ndarray]) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[numpy.ndarray], list[numpy.ndarray]][source]#: Get mean and stddev for descriptor.

call(coord_ext: numpy.ndarray, atype_ext: numpy.ndarray, nlist: numpy.ndarray, 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, mapps extended region index to local region.

Returns:

node_ebd: The output descriptor. shape: nf x nloc x n_dim (or n_dim + tebd_dim)
rot_mat: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x e_dim x 3
edge_ebd: The edge embedding. shape: nf x nloc x nnei x e_dim
h2: The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw: The smooth switch function. shape: nf x nloc x nnei

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

classmethod deserialize(data: dict) → DescrptDPA3[source]#

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

BD: The deserialized descriptor

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.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.base_descriptor.BaseDescriptor, deepmd.dpmodel.common.NativeOP

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.

descrpt_list = []#

numb_descrpt = 0#

nlist_cut_idx = []#

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_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 output dimension.

mixed_types()[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. All descriptors should be the same.

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.

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.

compute_input_stats(merged: list[dict], path: deepmd.utils.path.DPPath | None = None) → None[source]#: Update mean and stddev for descriptor elements.

set_stat_mean_and_stddev(mean: list[numpy.ndarray | list[numpy.ndarray]], stddev: list[numpy.ndarray | list[numpy.ndarray]]) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[numpy.ndarray | list[numpy.ndarray]], list[numpy.ndarray | list[numpy.ndarray]]][source]#: Get mean and stddev for descriptor.

enable_compression(min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1) → None[source]#

Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_dist: The nearest distance between atoms
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

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, not required 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.
h2: The rotationally equivariant pair-partical representation.
sw: The smooth switch function.

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

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

classmethod deserialize(data: dict) → DescrptHybrid[source]#

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

BD: The deserialized descriptor

deepmd.dpmodel.descriptor.make_base_descriptor(t_tensor, fwd_method_name: str = 'forward')[source]#

Make the base class for the descriptor.

Parameters:

t_tensor: The type of the tensor. used in the type hint.
fwd_method_name: Name of the forward method. For dpmodels, it should be “call”. For torch models, it should be “forward”.

class deepmd.dpmodel.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, 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, 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