deepmd.pt.model.descriptor.se_atten

deepmd.pt.model.descriptor.se_atten#

Classes#

`DescrptBlockSeAtten`	The building block of descriptor.
`NeighborGatedAttention`
`NeighborGatedAttentionLayer`
`GatedAttentionLayer`

Functions#

tabulate_fusion_se_atten(→ list[torch.Tensor])

Module Contents#

deepmd.pt.model.descriptor.se_atten.tabulate_fusion_se_atten(argument0: torch.Tensor, argument1: torch.Tensor, argument2: torch.Tensor, argument3: torch.Tensor, argument4: torch.Tensor, argument5: int, argument6: bool) → list[torch.Tensor][source]#

class deepmd.pt.model.descriptor.se_atten.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[source]#

rcut_smth[source]#

neuron[source]#

filter_neuron[source]#

axis_neuron[source]#

tebd_dim[source]#

tebd_input_mode[source]#

set_davg_zero[source]#

attn_dim[source]#

attn_layer[source]#

attn_dotr[source]#

attn_mask[source]#

activation_function[source]#

precision[source]#

prec[source]#

resnet_dt[source]#

scaling_factor[source]#

normalize[source]#

temperature[source]#

smooth[source]#

type_one_side[source]#

env_protection[source]#

trainable_ln[source]#

seed[source]#

ln_eps[source]#

ntypes[source]#

sel[source]#

sec[source]#

split_sel[source]#

nnei[source]#

ndescrpt[source]#

dpa1_attention[source]#

tebd_dim_input[source]#

filter_layers_strip = None[source]#

filter_layers[source]#

stats = None[source]#

compress = False[source]#

is_sorted = False[source]#

compress_info[source]#

compress_data[source]#

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

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

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

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

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

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

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

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

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

property dim_emb[source]#: 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.se_atten.NeighborGatedAttention(layer_num: int, nnei: int, embed_dim: int, hidden_dim: int, dotr: bool = False, do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-05, smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)[source]#

Bases: torch.nn.Module

layer_num[source]#

nnei[source]#

embed_dim[source]#

hidden_dim[source]#

dotr[source]#

do_mask[source]#

scaling_factor[source]#

normalize[source]#

temperature[source]#

trainable_ln[source]#

ln_eps[source]#

smooth[source]#

precision[source]#

seed[source]#

network_type[source]#

attention_layers[source]#

forward(input_G, nei_mask, input_r: torch.Tensor | None = None, sw: torch.Tensor | None = None)[source]#

Compute the multi-layer gated self-attention.

Parameters:

input_G: inputs with shape: (nf x nloc) x nnei x embed_dim.
nei_mask: neighbor mask, with paddings being 0. shape: (nf x nloc) x nnei.
input_r: normalized radial. shape: (nf x nloc) x nnei x 3.
sw: The smooth switch function. shape: nf x nloc x nnei

__getitem__(key)[source]#

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

serialize() → dict[source]#

Serialize the networks to a dict.

Returns:

dict: The serialized networks.

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

Deserialize the networks from a dict.

Parameters:

datadict: The dict to deserialize from.

class deepmd.pt.model.descriptor.se_atten.NeighborGatedAttentionLayer(nnei: int, embed_dim: int, hidden_dim: int, dotr: bool = False, do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = None, smooth: bool = True, trainable_ln: bool = True, ln_eps: float = 1e-05, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)[source]#

Bases: torch.nn.Module

nnei[source]#

embed_dim[source]#

hidden_dim[source]#

dotr[source]#

do_mask[source]#

scaling_factor[source]#

normalize[source]#

temperature[source]#

precision[source]#

trainable_ln[source]#

ln_eps[source]#

seed[source]#

attention_layer[source]#

attn_layer_norm[source]#

forward(x, nei_mask, input_r: torch.Tensor | None = None, sw: torch.Tensor | None = None)[source]#

serialize() → dict[source]#

Serialize the networks to a dict.

Returns:

dict: The serialized networks.

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

Deserialize the networks from a dict.

Parameters:

datadict: The dict to deserialize from.

class deepmd.pt.model.descriptor.se_atten.GatedAttentionLayer(nnei: int, embed_dim: int, hidden_dim: int, num_heads: int = 1, dotr: bool = False, do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)[source]#

Bases: torch.nn.Module

nnei[source]#

embed_dim[source]#

hidden_dim[source]#

num_heads[source]#

head_dim[source]#

dotr[source]#

do_mask[source]#

bias[source]#

smooth[source]#

scaling_factor[source]#

temperature[source]#

precision[source]#

seed[source]#

scaling[source]#

normalize[source]#

in_proj[source]#

out_proj[source]#

forward(query, nei_mask, input_r: torch.Tensor | None = None, sw: torch.Tensor | None = None, attnw_shift: float = 20.0)[source]#

Compute the multi-head gated self-attention.

Parameters:

query: inputs with shape: (nf x nloc) x nnei x embed_dim.
nei_mask: neighbor mask, with paddings being 0. shape: (nf x nloc) x nnei.
input_r: normalized radial. shape: (nf x nloc) x nnei x 3.
sw: The smooth switch function. shape: (nf x nloc) x nnei
attnw_shiftfloat: The attention weight shift to preserve smoothness when doing padding before softmax.

serialize() → dict[source]#

Serialize the networks to a dict.

Returns:

dict: The serialized networks.

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

Deserialize the networks from a dict.

Parameters:

datadict: The dict to deserialize from.