deepmd.pt.model.descriptor.se_atten

deepmd.pt.model.descriptor.se_atten#

Attributes#

_DEGREE_GAIN_INIT_STD

Classes#

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

Functions#

`tabulate_fusion_se_atten`(→ list[torch.Tensor])
`_safe_direction`(→ tuple[torch.Tensor, torch.Tensor, ...)	Scale displacements by the protected distance denominator.
`_compute_angular_radial`(→ torch.Tensor)	Compute the zero-mean radial factor for non-scalar moments.
`_build_moment_basis`(→ torch.Tensor)	Build the Cartesian moment basis through angular degree `lmax`.
`_build_degree_weights`(→ torch.Tensor)	Expand non-negative per-degree Gram weights to packed moment rows.

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

deepmd.pt.model.descriptor.se_atten._DEGREE_GAIN_INIT_STD = 0.1[source]#

deepmd.pt.model.descriptor.se_atten._safe_direction(diff: torch.Tensor, protection: float) → tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]#

Scale displacements by the protected distance denominator.

Parameters:

diff: Neighbor displacement vectors with shape (..., 3).
protection: Distance protection added to the denominator.

Returns:

torch.Tensor: Protected displacement coordinates with shape (..., 3).
torch.Tensor: Unprotected distances with shape (..., 1).
torch.Tensor: Nonzero-distance mask with shape (..., 1).

deepmd.pt.model.descriptor.se_atten._compute_angular_radial(distance: torch.Tensor, direction_mask: torch.Tensor, switch: torch.Tensor, stddev: torch.Tensor, valid_mask: torch.Tensor, protection: float) → torch.Tensor[source]#

Compute the zero-mean radial factor for non-scalar moments.

Non-scalar SO(3) components cannot carry an additive statistical mean. The factor therefore applies the scalar environment standard deviation without subtracting its mean, which keeps every angular component zero at the cutoff.

deepmd.pt.model.descriptor.se_atten._build_moment_basis(rr: torch.Tensor, direction: torch.Tensor, radial: torch.Tensor, lmax: int) → torch.Tensor[source]#

Build the Cartesian moment basis through angular degree lmax.

The first four rows preserve the existing scalar and vector environment matrix exactly. Higher-degree rows are norm-normalized real spherical harmonics in the m=-l,...,l order:

Y_l(u) @ Y_l(v) = P_l(u @ v).

Parameters:

rr: Normalized environment matrix with shape (ncenter, nnei, 4).
direction: Protected displacement coordinates with shape (ncenter, nnei, 3).
radial: Zero-mean normalized radial amplitude with shape (ncenter, nnei, 1). Invalid and excluded neighbors are zero.
lmax: Maximum angular degree. Supported values are 1 through 4.

Returns:

torch.Tensor: Moment basis with shape (ncenter, nnei, (lmax + 1) ** 2).

deepmd.pt.model.descriptor.se_atten._build_degree_weights(raw_gain: torch.Tensor, lmax: int, reference: torch.Tensor) → torch.Tensor[source]#: Expand non-negative per-degree Gram weights to packed moment rows.

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: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = 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, trainable: bool = True, lmax: int = 1)[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 = [25, 50, 100][source]#

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

axis_neuron = 16[source]#

lmax = 1[source]#

tebd_dim = 8[source]#

tebd_input_mode = 'concat'[source]#

set_davg_zero = True[source]#

attn_dim = 128[source]#

attn_layer = 2[source]#

attn_dotr = True[source]#

attn_mask = False[source]#

activation_function = 'tanh'[source]#

precision = 'float64'[source]#

prec[source]#

resnet_dt = False[source]#

scaling_factor = 1.0[source]#

normalize = True[source]#

temperature = None[source]#

smooth = True[source]#

type_one_side = False[source]#

env_protection = 0.0[source]#

trainable_ln = True[source]#

seed = None[source]#

ln_eps = 1e-05[source]#

ntypes[source]#

sel[source]#

sec[source]#

split_sel[source]#

nnei[source]#

ndescrpt[source]#

dpa1_attention[source]#

tebd_dim_input = 8[source]#

filter_layers_strip = None[source]#

filter_layers[source]#

stats = None[source]#

tebd_compress = False[source]#

geo_compress = False[source]#

is_sorted = False[source]#

compress_info[source]#

compress_data[source]#

use_triton_infer[source]#

se_conv_eligible[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: str, value: Any) → None[source]#

__getitem__(key: str) → Any[source]#

mixed_types() → bool[source]#

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

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

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

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

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

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

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

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

Parameters:

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

list[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: 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, table_config: dict, lower: dict, upper: dict) → None[source]#

type_embedding_compression(type_embedding_net: deepmd.pt.model.network.network.TypeEmbedNet) → None[source]#

Enable type embedding compression for strip mode.

Precomputes embedding network outputs for all type combinations: - One-side: (ntypes+1) combinations (neighbor types only) - Two-side: (ntypes+1)² combinations (neighbor x center type pairs)

Parameters:

type_embedding_netTypeEmbedNet: The type embedding network that provides get_full_embedding() method

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) → tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | 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, trainable: bool = True)[source]#

Bases: torch.nn.Module

layer_num[source]#

nnei[source]#

embed_dim[source]#

hidden_dim[source]#

dotr = False[source]#

do_mask = False[source]#

scaling_factor = 1.0[source]#

normalize = True[source]#

temperature = None[source]#

trainable_ln = True[source]#

ln_eps = 1e-05[source]#

smooth = True[source]#

precision = 'float64'[source]#

seed = None[source]#

network_type[source]#

attention_layers[source]#

forward(input_G: torch.Tensor, nei_mask: torch.Tensor, input_r: torch.Tensor | None = None, sw: torch.Tensor | None = None) → torch.Tensor[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: int) → Any[source]#

__setitem__(key: int, value: Any) → 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, trainable: bool = True)[source]#

Bases: torch.nn.Module

nnei[source]#

embed_dim[source]#

hidden_dim[source]#

dotr = False[source]#

do_mask = False[source]#

scaling_factor = 1.0[source]#

normalize = True[source]#

temperature = None[source]#

precision = 'float64'[source]#

trainable_ln = True[source]#

ln_eps = 1e-05[source]#

seed = None[source]#

attention_layer[source]#

attn_layer_norm[source]#

forward(x: torch.Tensor, nei_mask: torch.Tensor, input_r: torch.Tensor | None = None, sw: torch.Tensor | None = None) → torch.Tensor[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, trainable: bool = True)[source]#

Bases: torch.nn.Module

nnei[source]#

embed_dim[source]#

hidden_dim[source]#

num_heads = 1[source]#

head_dim[source]#

dotr = False[source]#

do_mask = False[source]#

bias = True[source]#

smooth = True[source]#

scaling_factor = 1.0[source]#

temperature = None[source]#

precision = 'float64'[source]#

seed = None[source]#

scaling[source]#

normalize = True[source]#

in_proj[source]#

out_proj[source]#

forward(query: torch.Tensor, nei_mask: torch.Tensor, input_r: torch.Tensor | None = None, sw: torch.Tensor | None = None, attnw_shift: float = 20.0) → tuple[torch.Tensor, torch.Tensor][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.