deepmd.pt.model.descriptor
Submodules
deepmd.pt.model.descriptor.base_descriptordeepmd.pt.model.descriptor.descriptordeepmd.pt.model.descriptor.dpa1deepmd.pt.model.descriptor.dpa2deepmd.pt.model.descriptor.env_matdeepmd.pt.model.descriptor.gaussian_lccdeepmd.pt.model.descriptor.hybriddeepmd.pt.model.descriptor.repformer_layerdeepmd.pt.model.descriptor.repformersdeepmd.pt.model.descriptor.se_adeepmd.pt.model.descriptor.se_attendeepmd.pt.model.descriptor.se_r
Package Contents
Classes
The building block of descriptor. | |
The building block of descriptor. | |
Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model. | |
Base class for all neural network modules. | |
Base class for all neural network modules. | |
The building block of descriptor. | |
Concate a list of descriptors to form a new descriptor. | |
The building block of descriptor. | |
The building block of descriptor. | |
Base descriptor provides the interfaces of descriptor. | |
Base descriptor provides the interfaces of descriptor. |
Functions
| |
| Generate smooth environment matrix from atom coordinates and other context. |
Attributes
- class deepmd.pt.model.descriptor.DescriptorBlock(*args, **kwargs)[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 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- abstract get_stats() Dict[str, deepmd.utils.env_mat_stat.StatItem][source]
Get the statistics of the descriptor.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- 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, type: str | None = None, old_impl: bool = False)[source]
Bases:
deepmd.pt.model.descriptor.descriptor.DescriptorBlockThe building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.
- 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.
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- get_stats() Dict[str, deepmd.utils.env_mat_stat.StatItem][source]
Get the statistics of the descriptor.
- forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: 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.
- Returns:
resultThe descriptor. shape: nf x nloc x (ng x axis_neuron)
g2The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
grThe rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
swThe smooth switch function. shape: nf x nloc x nnei
- 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 = False, spin=None, type: str | None = None, seed: int | None = None, old_impl: bool = False)[source]
Bases:
deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor,torch.nn.ModuleAttention-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\)
- sel
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
- ntypes
int Number of element types
- neuron
list[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 way to mix the type embeddings. Supported options are concat. (TODO need to support stripped_type_embedding option)
- 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_types
List[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 “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision: str
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- 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.
- 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.
- property dim_out
- property dim_emb
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- classmethod deserialize(data: dict) DescrptDPA1[source]
Deserialize the model.
- Parameters:
- data
dict The serialized data
- data
- Returns:
BDThe 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, not required by this descriptor.
- comm_dict
The data needed for communication for parallel inference.
- Returns:
descriptorThe descriptor. shape: nf x nloc x (ng x axis_neuron)
grThe rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
swThe smooth switch function. shape: nf x nloc x nnei
- class deepmd.pt.model.descriptor.DescrptDPA2(ntypes: int, repinit_rcut: float, repinit_rcut_smth: float, repinit_nsel: int, repformer_rcut: float, repformer_rcut_smth: float, repformer_nsel: int, tebd_dim: int = 8, concat_output_tebd: bool = True, repinit_neuron: List[int] = [25, 50, 100], repinit_axis_neuron: int = 16, repinit_set_davg_zero: bool = True, repinit_activation='tanh', repformer_nlayers: int = 3, repformer_g1_dim: int = 128, repformer_g2_dim: int = 16, repformer_axis_dim: int = 4, repformer_do_bn_mode: str = 'no', repformer_bn_momentum: float = 0.1, repformer_update_g1_has_conv: bool = True, repformer_update_g1_has_drrd: bool = True, repformer_update_g1_has_grrg: bool = True, repformer_update_g1_has_attn: bool = True, repformer_update_g2_has_g1g1: bool = True, repformer_update_g2_has_attn: bool = True, repformer_update_h2: bool = False, repformer_attn1_hidden: int = 64, repformer_attn1_nhead: int = 4, repformer_attn2_hidden: int = 16, repformer_attn2_nhead: int = 4, repformer_attn2_has_gate: bool = False, repformer_activation: str = 'tanh', repformer_update_style: str = 'res_avg', repformer_set_davg_zero: bool = True, repformer_add_type_ebd_to_seq: bool = False, env_protection: float = 0.0, trainable: bool = True, exclude_types: List[Tuple[int, int]] = [], type: str | None = None, rcut: float | None = None, rcut_smth: float | None = None, sel: int | None = None)[source]
Bases:
torch.nn.Module,deepmd.pt.model.descriptor.base_descriptor.BaseDescriptorBase class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:
import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x))
Submodules assigned in this way will be registered, and will have their parameters converted too when you call
to(), etc.Note
As per the example above, an
__init__()call to the parent class must be made before assignment on the child.- Variables:
training (bool) – Boolean represents whether this module is in training or evaluation mode.
- property dim_out
- property dim_emb
Returns the embedding dimension g2.
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- abstract classmethod deserialize() DescrptDPA2[source]
Deserialize from a dict.
- 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:
- 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.
- comm_dict
The data needed for communication for parallel inference.
- Returns:
descriptorThe descriptor. shape: nf x nloc x (ng x axis_neuron)
grThe rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
swThe smooth switch function. shape: nf x nloc x nnei
- 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:
Shapeis[nframes,natoms[1]*nnei*4].
- env_mat:
- class deepmd.pt.model.descriptor.DescrptGaussianLcc(rcut, rcut_smth, sel: int, ntypes: int, num_pair: int, embed_dim: int = 768, kernel_num: int = 128, pair_embed_dim: int = 64, num_block: int = 1, layer_num: int = 12, attn_head: int = 48, pair_hidden_dim: int = 16, ffn_embedding_dim: int = 768, dropout: float = 0.0, droppath_prob: float = 0.1, pair_dropout: float = 0.25, attention_dropout: float = 0.1, activation_dropout: float = 0.1, pre_ln: bool = True, do_tag_embedding: bool = False, tag_ener_pref: bool = False, atomic_sum_gbf: bool = False, pre_add_seq: bool = True, tri_update: bool = True, **kwargs)[source]
Bases:
torch.nn.Module,deepmd.pt.model.descriptor.base_descriptor.BaseDescriptorBase class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:
import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x))
Submodules assigned in this way will be registered, and will have their parameters converted too when you call
to(), etc.Note
As per the example above, an
__init__()call to the parent class must be made before assignment on the child.- Variables:
training (bool) – Boolean represents whether this module is in training or evaluation mode.
- property dim_out
Returns the output dimension of atomic representation.
- property dim_in
Returns the atomic input dimension of this descriptor.
- property dim_emb
Returns the output dimension of pair representation.
- compute_input_stats(merged: List[dict], path: deepmd.utils.path.DPPath | None = None)[source]
Update mean and stddev for descriptor elements.
- forward(extended_coord, nlist, atype, nlist_type, nlist_loc=None, atype_tebd=None, nlist_tebd=None, seq_input=None)[source]
Calculate the atomic and pair representations of this descriptor.
Args: - extended_coord: Copied atom coordinates with shape [nframes, nall, 3]. - nlist: Neighbor list with shape [nframes, nloc, nnei]. - atype: Atom type with shape [nframes, nloc]. - nlist_type: Atom type of neighbors with shape [nframes, nloc, nnei]. - nlist_loc: Local index of neighbor list with shape [nframes, nloc, nnei]. - atype_tebd: Atomic type embedding with shape [nframes, nloc, tebd_dim]. - nlist_tebd: Type embeddings of neighbor with shape [nframes, nloc, nnei, tebd_dim]. - seq_input: The sequential input from other descriptor with
shape [nframes, nloc, tebd_dim] or [nframes * nloc, 1 + nnei, tebd_dim]
- Returns:
- result:
descriptorwithshape[nframes,nloc,self.filter_neuron[-1] *self.axis_neuron].
- result:
- ret:
environmentmatrixwithshape[nframes,nloc,self.neei,out_size]
- ret:
- class deepmd.pt.model.descriptor.DescrptBlockHybrid(list, ntypes: int, tebd_dim: int = 8, tebd_input_mode: str = 'concat', hybrid_mode: str = 'concat', **kwargs)[source]
Bases:
deepmd.pt.model.descriptor.DescriptorBlockThe building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.
- property dim_out
Returns the output dimension of this descriptor.
- property dim_emb
Returns the output dimension of embedding.
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None)[source]
Calculate decoded embedding for each atom.
Args: - extended_coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - nlist: Tell atom types with shape [nframes, natoms[1]]. - atype: Tell atom count and element count. Its shape is [2+self.ntypes]. - nlist_type: Tell simulation box with shape [nframes, 9]. - atype_tebd: Tell simulation box with shape [nframes, 9]. - nlist_tebd: Tell simulation box with shape [nframes, 9].
- Returns:
- result:
descriptorwithshape[nframes,nloc,self.filter_neuron[-1] *self.axis_neuron].
- result:
- ret:
environmentmatrixwithshape[nframes,nloc,self.neei,out_size]
- ret:
- class deepmd.pt.model.descriptor.DescrptHybrid(list: List[deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor | Dict[str, Any]], **kwargs)[source]
Bases:
deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor,torch.nn.ModuleConcate a list of descriptors to form a new descriptor.
- Parameters:
- mixed_types()[source]
Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- compute_input_stats(merged: List[dict], path: deepmd.utils.path.DPPath | None = None)[source]
Update mean and stddev for descriptor elements.
- 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:
descriptorThe descriptor. shape: nf x nloc x (ng x axis_neuron)
grThe rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3. This descriptor returns None
g2The rotationally invariant pair-partical representation. this descriptor returns None
h2The rotationally equivariant pair-partical representation. this descriptor returns None
swThe smooth switch function. this descriptor returns None
- classmethod update_sel(global_jdata: dict, local_jdata: dict) dict[source]
Update the selection and perform neighbor statistics.
- classmethod deserialize(data: dict) DescrptHybrid[source]
Deserialize the model.
- Parameters:
- data
dict The serialized data
- data
- Returns:
BDThe 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_dim: int = 4, direct_dist: bool = False, do_bn_mode: str = 'no', bn_momentum: float = 0.1, 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', set_davg_zero: bool = True, smooth: bool = True, add_type_ebd_to_seq: bool = False, exclude_types: List[Tuple[int, int]] = [], env_protection: float = 0.0, type: str | None = None)[source]
Bases:
deepmd.pt.model.descriptor.descriptor.DescriptorBlockThe building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.
- 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.
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.
- forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: 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)[source]
Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.
- Parameters:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- get_stats() Dict[str, deepmd.utils.env_mat_stat.StatItem][source]
Get the statistics of the descriptor.
- 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, old_impl: bool = False, type_one_side: bool = True, trainable: bool = True, **kwargs)[source]
Bases:
deepmd.pt.model.descriptor.DescriptorBlockThe building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.
- property dim_out
Returns the output dimension of this descriptor.
- property dim_in
Returns the atomic input dimension of this descriptor.
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- get_stats() Dict[str, deepmd.utils.env_mat_stat.StatItem][source]
Get the statistics of the descriptor.
- forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: 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:
descriptormatrixwithshape[nframes,natoms[0]*self.filter_neuron[-1]*self.axis_neuron].
- torch.Tensor:
- 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, old_impl: bool = False, type_one_side: bool = True, **kwargs)[source]
Bases:
deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor,torch.nn.ModuleBase descriptor provides the interfaces of descriptor.
- property dim_out
Returns the output dimension of this descriptor.
- mixed_types()[source]
Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- 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:
descriptorThe descriptor. shape: nf x nloc x (ng x axis_neuron)
grThe rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2The rotationally invariant pair-partical representation. this descriptor returns None
h2The rotationally equivariant pair-partical representation. this descriptor returns None
swThe smooth switch function.
- classmethod deserialize(data: dict) DescrptSeA[source]
Deserialize the model.
- Parameters:
- data
dict The serialized data
- data
- Returns:
BDThe deserialized descriptor
- class deepmd.pt.model.descriptor.DescrptSeR(rcut, rcut_smth, sel, neuron=[25, 50, 100], 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, old_impl: bool = False, trainable: bool = True, **kwargs)[source]
Bases:
deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor,torch.nn.ModuleBase descriptor provides the interfaces of descriptor.
- mixed_types() bool[source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.
Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.
- 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:
- merged
Union[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.
- path
Optional[DPPath] The path to the stat file.
- merged
- get_stats() Dict[str, deepmd.utils.env_mat_stat.StatItem][source]
Get the statistics of the descriptor.
- forward(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: 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.
- Returns:
descriptorThe descriptor. shape: nf x nloc x (ng x axis_neuron)
grThe rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2The rotationally invariant pair-partical representation. this descriptor returns None
h2The rotationally equivariant pair-partical representation. this descriptor returns None
swThe smooth switch function.
- classmethod deserialize(data: dict) DescrptSeR[source]
Deserialize the model.
- Parameters:
- data
dict The serialized data
- data
- Returns:
BDThe deserialized descriptor