deepmd.pd.model.descriptor.se_atten

Classes

DescrptBlockSeAtten	The building block of descriptor.
NeighborGatedAttention
NeighborGatedAttentionLayer
GatedAttentionLayer

Module Contents

class deepmd.pd.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)

Bases: deepmd.pd.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 = [25, 50, 100]

filter_neuron = [25, 50, 100]

axis_neuron = 16

tebd_dim = 8

tebd_input_mode = 'concat'

set_davg_zero = True

attn_dim = 128

attn_layer = 2

attn_dotr = True

attn_mask = False

activation_function = 'tanh'

precision = 'float64'

prec

resnet_dt = False

scaling_factor = 1.0

normalize = True

temperature = None

smooth = True

type_one_side = False

env_protection = 0.0

trainable_ln = True

seed = None

ln_eps = 1e-05

ntypes

sel

sec

split_sel

nnei

ndescrpt

dpa1_attention

tebd_dim_input = 8

filter_layers_strip = None

filter_layers

stats = None

compress = False

is_sorted = False

compress_info

compress_data

get_rcut() → float

Returns the cut-off radius.

get_rcut_smth() → float

Returns the radius where the neighbor information starts to smoothly decay to 0.

get_buffer_rcut() → paddle.Tensor

Returns the cut-off radius as a buffer-style Tensor.

get_buffer_rcut_smth() → paddle.Tensor

Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.

get_nsel() → int

Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int]

Returns the number of selected atoms for each type.

get_ntypes() → int

Returns the number of element types.

get_dim_in() → int

Returns the input dimension.

get_dim_out() → int

Returns the output dimension.

get_dim_rot_mat_1() → int

Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

get_dim_emb() → int

Returns the output dimension of embedding.

__setitem__(key: str, value: paddle.Tensor) → None

__getitem__(key: str) → paddle.Tensor

mixed_types() → bool

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

Returns the protection of building environment matrix.

property dim_out: int

Returns the output dimension of this descriptor.

property dim_in: int

Returns the atomic input dimension of this descriptor.

property dim_emb: int

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

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: paddle.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]

Get the statistics of the descriptor.

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

enable_compression(table_data: dict[str, paddle.Tensor], table_config: list[float], lower: dict[str, float], upper: dict[str, float]) → None

forward(nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, extended_atype_embd: paddle.Tensor | None = None, mapping: paddle.Tensor | None = None, type_embedding: paddle.Tensor | None = None) → tuple[paddle.Tensor, paddle.Tensor | None, paddle.Tensor, paddle.Tensor, paddle.Tensor]

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

Returns whether the descriptor block has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the descriptor block needs sorted nlist when using forward_lower.

class deepmd.pd.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)

Bases: paddle.nn.Layer

layer_num

nnei

embed_dim

hidden_dim

dotr = False

do_mask = False

scaling_factor = 1.0

normalize = True

temperature = None

trainable_ln = True

ln_eps = 1e-05

smooth = True

precision = 'float64'

seed = None

network_type

attention_layers

forward(input_G: paddle.Tensor, nei_mask: paddle.Tensor, input_r: paddle.Tensor | None = None, sw: paddle.Tensor | None = None) → paddle.Tensor

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) → paddle.nn.Layer

__setitem__(key: int, value: paddle.nn.Layer | dict) → None

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → NeighborGatedAttention

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.

class deepmd.pd.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)

Bases: paddle.nn.Layer

nnei

embed_dim

hidden_dim

dotr = False

do_mask = False

scaling_factor = 1.0

normalize = True

temperature = None

precision = 'float64'

trainable_ln = True

ln_eps = 1e-05

seed = None

attention_layer

attn_layer_norm

forward(x: paddle.Tensor, nei_mask: paddle.Tensor, input_r: paddle.Tensor | None = None, sw: paddle.Tensor | None = None) → paddle.Tensor

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → NeighborGatedAttentionLayer

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.

class deepmd.pd.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)

Bases: paddle.nn.Layer

nnei

embed_dim

hidden_dim

num_heads = 1

head_dim

dotr = False

do_mask = False

bias = True

smooth = True

scaling_factor = 1.0

temperature = None

precision = 'float64'

seed = None

scaling

normalize = True

in_proj

out_proj

forward(query: paddle.Tensor, nei_mask: paddle.Tensor, input_r: paddle.Tensor | None = None, sw: paddle.Tensor | None = None, attnw_shift: float = 20.0) → tuple[paddle.Tensor, paddle.Tensor]

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

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → GatedAttentionLayer

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.