deepmd.pt.model.descriptor

deepmd.pt.model.descriptor#

Submodules#

Attributes#

BaseDescriptor

Classes#

`DescriptorBlock`	The building block of descriptor.
`DescrptBlockSeAtten`	The building block of descriptor.
`DescrptDPA1`	Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.
`DescrptDPA2`	Base descriptor provides the interfaces of descriptor.
`DescrptHybrid`	Concate a list of descriptors to form a new descriptor.
`DescrptBlockRepformers`	The building block of descriptor.
`DescrptBlockSeA`	The building block of descriptor.
`DescrptSeA`	Base descriptor provides the interfaces of descriptor.
`DescrptSeAttenV2`	Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.
`DescrptSeR`	Base descriptor provides the interfaces of descriptor.
`DescrptSeT`	DeepPot-SE constructed from all information (both angular and radial) of atomic
`DescrptBlockSeTTebd`	The building block of descriptor.
`DescrptSeTTebd`	Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.

Functions#

`make_default_type_embedding`(ntypes)
`prod_env_mat`(extended_coord, nlist, atype, mean, ...)	Generate smooth environment matrix from atom coordinates and other context.

Package Contents#

deepmd.pt.model.descriptor.BaseDescriptor[source]#

class deepmd.pt.model.descriptor.DescriptorBlock[source]#

Bases: torch.nn.Module, abc.ABC, make_plugin_registry('DescriptorBlock')

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

local_cluster = False#

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

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

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

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

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

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

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

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

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

abstract compute_input_stats(merged: Callable[[], list[dict]] | list[dict], path: deepmd.utils.path.DPPath | None = None) → NoReturn[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.

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

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.

abstract forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, type_embedding: torch.Tensor | None = None)[source]#: Calculate DescriptorBlock.

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

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

deepmd.pt.model.descriptor.make_default_type_embedding(ntypes)[source]#

class deepmd.pt.model.descriptor.DescrptBlockSeAtten(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, tebd_input_mode: str = 'concat', set_davg_zero: bool = True, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, activation_function='tanh', precision: str = 'float64', resnet_dt: bool = False, scaling_factor=1.0, normalize=True, temperature=None, smooth: bool = True, type_one_side: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable_ln: bool = True, ln_eps: float | None = 1e-05, seed: int | list[int] | None = None, type: str | None = None)[source]#

Bases: deepmd.pt.model.descriptor.descriptor.DescriptorBlock

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

rcut#

rcut_smth#

neuron#

filter_neuron#

axis_neuron#

tebd_dim#

tebd_input_mode#

set_davg_zero#

attn_dim#

attn_layer#

attn_dotr#

attn_mask#

activation_function#

precision#

prec#

resnet_dt#

scaling_factor#

normalize#

temperature#

smooth#

type_one_side#

env_protection#

trainable_ln#

seed#

ln_eps#

ntypes#

sel#

sec#

split_sel#

nnei#

ndescrpt#

dpa1_attention#

tebd_dim_input#

filter_layers_strip = None#

filter_layers#

stats = None#

compress = False#

is_sorted = False#

compress_info#

compress_data#

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_rot_mat_1() → int[source]#: Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

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

__setitem__(key, value) → None[source]#

__getitem__(key)[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#: Returns the output dimension of this descriptor.

property dim_in#: Returns the atomic input dimension of this descriptor.

property dim_emb#: Returns the output dimension of embedding.

compute_input_stats(merged: 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.

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]#

enable_compression(table_data, table_config, lower, upper) → None[source]#

forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, type_embedding: torch.Tensor | None = None)[source]#

Compute the descriptor.

Parameters:

nlist: The neighbor list. shape: nf x nloc x nnei
extended_coord: The extended coordinates of atoms. shape: nf x (nallx3)
extended_atype: The extended aotm types. shape: nf x nall x nt
extended_atype_embd: The extended type embedding of atoms. shape: nf x nall
mapping: The index mapping, not required by this descriptor.
type_embedding: Full type embeddings. shape: (ntypes+1) x nt Required for stripped type embeddings.

Returns:

result: The descriptor. shape: nf x nloc x (ng x axis_neuron)
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
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
sw: The smooth switch function. shape: nf x nloc x nnei

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.

class deepmd.pt.model.descriptor.DescrptDPA1(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, tebd_input_mode: str = 'concat', set_davg_zero: bool = True, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, scaling_factor: int = 1.0, normalize=True, temperature=None, concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, ln_eps: float | None = 1e-05, smooth_type_embedding: bool = True, type_one_side: bool = False, stripped_type_embedding: bool | None = None, seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, spin=None, type: str | None = None)[source]#

Bases: deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.Module

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.
seed: int, Optional: Random seed for parameter initialization.
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.

tebd_input_mode#

se_atten#

use_econf_tebd#

use_tebd_bias#

type_map#

compress = False#

type_embedding#

prec#

tebd_dim#

concat_output_tebd#

trainable#

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.

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.

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: torch.Tensor, stddev: torch.Tensor) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[torch.Tensor, torch.Tensor][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.

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

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

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

BD: The deserialized 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

forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None)[source]#

Compute the descriptor.

Parameters:

extended_coord: The extended coordinates of atoms. shape: nf x (nallx3)
extended_atype: 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.
comm_dict: The data needed for communication for parallel inference.

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

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.pt.model.descriptor.DescrptDPA2(ntypes: int, repinit: deepmd.dpmodel.descriptor.dpa2.RepinitArgs | dict, repformer: deepmd.dpmodel.descriptor.dpa2.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.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.Module

Base descriptor provides the interfaces of descriptor.

repinit_args#

repformer_args#

tebd_input_mode#

repinit#

use_three_body#

repformers#

rcsl_list#

rcut_list#

nsel_list#

use_econf_tebd#

use_tebd_bias#

type_map#

type_embedding#

concat_output_tebd#

precision#

prec#

smooth#

exclude_types#

env_protection#

trainable#

add_tebd_to_repinit_out#

repinit_out_dim#

tebd_transform = None#

tebd_dim#

rcut#

rcut_smth#

ntypes#

sel#

compress = False#

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.

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: 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: list[torch.Tensor], stddev: list[torch.Tensor]) → None[source]#: Update mean and stddev for descriptor.

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

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

forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None)[source]#

Compute the descriptor.

Parameters:

extended_coord: The extended coordinates of atoms. shape: nf x (nallx3)
extended_atype: 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.
comm_dict: The data needed for communication for parallel inference.

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

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

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 statistics (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

deepmd.pt.model.descriptor.prod_env_mat(extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float, radial_only: bool = False, protection: float = 0.0)[source]#

Generate smooth environment matrix from atom coordinates and other context.

Args: - extended_coord: Copied atom coordinates with shape [nframes, nall*3]. - atype: Atom types with shape [nframes, nloc]. - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4 or 1]. - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4 or 1]. - rcut: Cut-off radius. - rcut_smth: Smooth hyper-parameter for pair force & energy. - radial_only: Whether to return a full description or a radial-only descriptor. - protection: Protection parameter to prevent division by zero errors during calculations.

Returns:

env_mat: Shape is [nframes, natoms[1]*nnei*4].

class deepmd.pt.model.descriptor.DescrptHybrid(list: DescrptHybrid.__init__.list[deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor | dict[str, Any]], **kwargs)[source]#

Bases: deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.Module

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.

nlist_cut_idx: list[torch.Tensor]#

descrpt_list#

numb_descrpt#

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.

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.

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[torch.Tensor | list[torch.Tensor]], stddev: list[torch.Tensor | list[torch.Tensor]]) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[torch.Tensor | list[torch.Tensor]], list[torch.Tensor | list[torch.Tensor]]][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

forward(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | 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.
comm_dict: The data needed for communication for parallel inference.

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. This descriptor returns None
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. this descriptor returns None

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

class deepmd.pt.model.descriptor.DescrptBlockRepformers(rcut, rcut_smth, sel: int, ntypes: int, nlayers: int = 3, g1_dim=128, g2_dim=16, axis_neuron: int = 4, direct_dist: bool = False, update_g1_has_conv: bool = True, update_g1_has_drrd: bool = True, update_g1_has_grrg: bool = True, update_g1_has_attn: bool = True, update_g2_has_g1g1: bool = True, update_g2_has_attn: bool = True, update_h2: bool = False, attn1_hidden: int = 64, attn1_nhead: int = 4, attn2_hidden: int = 16, attn2_nhead: int = 4, attn2_has_gate: bool = False, activation_function: str = 'tanh', update_style: str = 'res_avg', update_residual: float = 0.001, update_residual_init: str = 'norm', set_davg_zero: bool = True, smooth: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, precision: str = 'float64', trainable_ln: bool = True, ln_eps: float | None = 1e-05, seed: int | list[int] | None = None, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True)[source]#

Bases: deepmd.pt.model.descriptor.descriptor.DescriptorBlock

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

rcut#

rcut_smth#

ntypes#

nlayers#

nnei#

ndescrpt#

sel#

sec#

split_sel#

axis_neuron#

set_davg_zero#

g1_dim#

g2_dim#

update_g1_has_conv#

update_g1_has_drrd#

update_g1_has_grrg#

update_g1_has_attn#

update_g2_has_g1g1#

update_g2_has_attn#

update_h2#

attn1_hidden#

attn1_nhead#

attn2_has_gate#

attn2_hidden#

attn2_nhead#

activation_function#

update_style#

update_residual#

update_residual_init#

direct_dist#

act#

smooth#

use_sqrt_nnei#

g1_out_conv#

g1_out_mlp#

env_protection#

precision#

prec#

trainable_ln#

ln_eps#

epsilon = 0.0001#

seed#

g2_embd#

layers#

stats = None#

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_out() → int[source]#: Returns the output dimension.

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

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

__setitem__(key, value) → None[source]#

__getitem__(key)[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#: Returns the output dimension of this descriptor.

property dim_in#: Returns the atomic input dimension of this descriptor.

property dim_emb#: Returns the embedding dimension g2.

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

forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, type_embedding: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None)[source]#: Calculate DescriptorBlock.

compute_input_stats(merged: 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.

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

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.

class deepmd.pt.model.descriptor.DescrptBlockSeA(rcut, rcut_smth, sel, neuron=[25, 50, 100], axis_neuron=16, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, seed: int | list[int] | None = None, **kwargs)[source]#

Bases: deepmd.pt.model.descriptor.DescriptorBlock

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

ndescrpt: Final[int]#

__constants__: ClassVar[list] = ['ndescrpt']#

rcut#

rcut_smth#

neuron#

filter_neuron#

axis_neuron#

set_davg_zero#

activation_function#

precision#

prec#

resnet_dt#

env_protection#

ntypes#

type_one_side#

seed#

sel#

sec#

split_sel#

nnei#

filter_layers#

stats = None#

trainable#

compress = False#

compress_info#

compress_data#

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_out() → int[source]#: Returns the output dimension.

get_dim_rot_mat_1() → int[source]#: Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

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

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

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#: Returns the output dimension of this descriptor.

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

__setitem__(key, value) → None[source]#

__getitem__(key)[source]#

compute_input_stats(merged: 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.

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]#

enable_compression(table_data: dict[str, torch.Tensor], table_config: list[int | float], lower: dict[str, int], upper: dict[str, int]) → None[source]#

forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, type_embedding: torch.Tensor | None = None)[source]#

Calculate decoded embedding for each atom.

Args: - coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - atype: Tell atom types with shape [nframes, natoms[1]]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - box: Tell simulation box with shape [nframes, 9].

Returns:

torch.Tensor: descriptor matrix with shape [nframes, natoms[0]*self.filter_neuron[-1]*self.axis_neuron].

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.

class deepmd.pt.model.descriptor.DescrptSeA(rcut, rcut_smth, sel, neuron=[25, 50, 100], axis_neuron=16, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, seed: int | list[int] | None = None, ntypes: int | None = None, type_map: list[str] | None = None, spin=None)[source]#

Bases: deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.Module

Base descriptor provides the interfaces of descriptor.

type_map#

compress = False#

prec#

sea#

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

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.

property dim_out#: Returns the output dimension of this descriptor.

abstract 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: 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.

reinit_exclude(exclude_types: list[tuple[int, int]] = []) → None[source]#: Update the type exclusions.

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

forward(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | 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.
comm_dict: The data needed for communication for parallel inference.

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.

set_stat_mean_and_stddev(mean: torch.Tensor, stddev: torch.Tensor) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[torch.Tensor, torch.Tensor][source]#: Get mean and stddev for descriptor.

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

classmethod deserialize(data: dict) → DescrptSeA[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.pt.model.descriptor.DescrptSeAttenV2(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, set_davg_zero: bool = True, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, scaling_factor: int = 1.0, normalize=True, temperature=None, concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, ln_eps: float | None = 1e-05, type_one_side: bool = False, stripped_type_embedding: bool | None = None, seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, spin=None, type: str | None = None)[source]#

Bases: deepmd.pt.model.descriptor.dpa1.DescrptDPA1