deepmd.pd.model.descriptor.repformer_layer

Classes

Atten2Map
Atten2MultiHeadApply
Atten2EquiVarApply
LocalAtten
RepformerLayer

Functions

get_residual(→ paddle.Tensor)	Get residual tensor for one update vector.
_make_nei_g1(→ paddle.Tensor)	Make neighbor-wise atomic invariant rep.
_apply_nlist_mask(→ paddle.Tensor)	Apply nlist mask to neighbor-wise rep tensors.
_apply_switch(→ paddle.Tensor)	Apply switch function to neighbor-wise rep tensors.

Module Contents

deepmd.pd.model.descriptor.repformer_layer.get_residual(_dim: int, _scale: float, _mode: str = 'norm', trainable: bool = True, precision: str = 'float64', seed: int | list[int] | None = None) → paddle.Tensor

Get residual tensor for one update vector.

Parameters:

_dimint

The dimension of the update vector.

_scale

The initial scale of the residual tensor. See _mode for details.

_mode

The mode of residual initialization for the residual tensor. - “norm” (default): init residual using normal with _scale std. - “const”: init residual using element-wise constants of _scale.

trainable

Whether the residual tensor is trainable.

precision

The precision of the residual tensor.

seedint, optional

Random seed for parameter initialization.

deepmd.pd.model.descriptor.repformer_layer._make_nei_g1(g1_ext: paddle.Tensor, nlist: paddle.Tensor) → paddle.Tensor

Make neighbor-wise atomic invariant rep.

Parameters:

g1_ext

Extended atomic invariant rep, with shape nb x nall x ng1.

nlist

Neighbor list, with shape nb x nloc x nnei.

Returns:

gg1: paddle.Tensor

Neighbor-wise atomic invariant rep, with shape nb x nloc x nnei x ng1.

deepmd.pd.model.descriptor.repformer_layer._apply_nlist_mask(gg: paddle.Tensor, nlist_mask: paddle.Tensor) → paddle.Tensor

Apply nlist mask to neighbor-wise rep tensors.

Parameters:

gg

Neighbor-wise rep tensors, with shape nf x nloc x nnei x d.

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nf x nloc x nnei.

deepmd.pd.model.descriptor.repformer_layer._apply_switch(gg: paddle.Tensor, sw: paddle.Tensor) → paddle.Tensor

Apply switch function to neighbor-wise rep tensors.

Parameters:

gg

Neighbor-wise rep tensors, with shape nf x nloc x nnei x d.

sw

The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, and remains 0 beyond rcut, with shape nf x nloc x nnei.

class deepmd.pd.model.descriptor.repformer_layer.Atten2Map(input_dim: int, hidden_dim: int, head_num: int, has_gate: bool = False, smooth: bool = True, attnw_shift: float = 20.0, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: paddle.nn.Layer

input_dim

hidden_dim

head_num

mapqk

has_gate = False

smooth = True

attnw_shift = 20.0

precision = 'float64'

forward(g2: paddle.Tensor, h2: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor) → paddle.Tensor

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → Atten2Map

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.

class deepmd.pd.model.descriptor.repformer_layer.Atten2MultiHeadApply(input_dim: int, head_num: int, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: paddle.nn.Layer

input_dim

head_num

mapv

head_map

precision = 'float64'

forward(AA: paddle.Tensor, g2: paddle.Tensor) → paddle.Tensor

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → Atten2MultiHeadApply

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.

class deepmd.pd.model.descriptor.repformer_layer.Atten2EquiVarApply(input_dim: int, head_num: int, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: paddle.nn.Layer

input_dim

head_num

head_map

precision = 'float64'

forward(AA: paddle.Tensor, h2: paddle.Tensor) → paddle.Tensor

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → Atten2EquiVarApply

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.

class deepmd.pd.model.descriptor.repformer_layer.LocalAtten(input_dim: int, hidden_dim: int, head_num: int, smooth: bool = True, attnw_shift: float = 20.0, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: paddle.nn.Layer

input_dim

hidden_dim

head_num

mapq

mapkv

head_map

smooth = True

attnw_shift = 20.0

precision = 'float64'

forward(g1: paddle.Tensor, gg1: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor) → paddle.Tensor

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → LocalAtten

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.

class deepmd.pd.model.descriptor.repformer_layer.RepformerLayer(rcut, rcut_smth, sel: int, ntypes: int, g1_dim=128, g2_dim=16, axis_neuron: int = 4, update_chnnl_2: bool = True, 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', smooth: bool = True, precision: str = 'float64', trainable_ln: bool = True, ln_eps: float | None = 1e-05, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, seed: int | list[int] | None = None)

Bases: paddle.nn.Layer

epsilon = 0.0001

rcut

rcut_smth

ntypes

nnei

sel

sec

axis_neuron = 4

activation_function = 'tanh'

act

update_g1_has_grrg = True

update_g1_has_drrd = True

update_g1_has_conv = True

update_g1_has_attn = True

update_chnnl_2 = True

update_g2_has_g1g1 = True

update_g2_has_attn = True

update_h2 = False

attn1_hidden = 64

attn1_nhead = 4

attn2_hidden = 16

attn2_nhead = 4

attn2_has_gate = False

update_style = 'res_avg'

update_residual = 0.001

update_residual_init = 'norm'

smooth = True

g1_dim = 128

g2_dim = 16

trainable_ln = True

ln_eps = 1e-05

precision = 'float64'

seed = None

use_sqrt_nnei = True

g1_out_conv = True

g1_out_mlp = True

g1_residual = []

g2_residual = []

h2_residual = []

linear1

linear2 = None

proj_g1g2 = None

proj_g1g1g2 = None

attn2g_map = None

attn2_mh_apply = None

attn2_lm = None

attn2_ev_apply = None

loc_attn = None

cal_1_dim(g1d: int, g2d: int, ax: int) → int

_update_h2(h2: paddle.Tensor, attn: paddle.Tensor) → paddle.Tensor

Calculate the attention weights update for pair-wise equivariant rep.

Parameters:

h2

Pair-wise equivariant rep tensors, with shape nf x nloc x nnei x 3.

attn

Attention weights from g2 attention, with shape nf x nloc x nnei x nnei x nh2.

_update_g1_conv(gg1: paddle.Tensor, g2: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor) → paddle.Tensor

Calculate the convolution update for atomic invariant rep.

Parameters:

gg1

Neighbor-wise atomic invariant rep, with shape nb x nloc x nnei x ng1.

g2

Pair invariant rep, with shape nb x nloc x nnei x ng2.

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nb x nloc x nnei.

sw

The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, and remains 0 beyond rcut, with shape nb x nloc x nnei.

static _cal_hg(g2: paddle.Tensor, h2: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) → paddle.Tensor

Calculate the transposed rotation matrix.

Parameters:

g2

Neighbor-wise/Pair-wise invariant rep tensors, with shape nb x nloc x nnei x ng2.

h2

Neighbor-wise/Pair-wise equivariant rep tensors, with shape nb x nloc x nnei x 3.

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nb x nloc x nnei.

sw

The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, and remains 0 beyond rcut, with shape nb x nloc x nnei.

smooth

Whether to use smoothness in processes such as attention weights calculation.

epsilon

Protection of 1./nnei.

Returns:

hg

The transposed rotation matrix, with shape nb x nloc x 3 x ng2.

static _cal_grrg(h2g2: paddle.Tensor, axis_neuron: int) → paddle.Tensor

Calculate the atomic invariant rep.

Parameters:

h2g2

The transposed rotation matrix, with shape nb x nloc x 3 x ng2.

axis_neuron

Size of the submatrix.

Returns:

grrg

Atomic invariant rep, with shape nb x nloc x (axis_neuron x ng2)

symmetrization_op(g2: paddle.Tensor, h2: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor, axis_neuron: int, smooth: bool = True, epsilon: float = 0.0001) → paddle.Tensor

Symmetrization operator to obtain atomic invariant rep.

Parameters:

g2

Neighbor-wise/Pair-wise invariant rep tensors, with shape nb x nloc x nnei x ng2.

h2

Neighbor-wise/Pair-wise equivariant rep tensors, with shape nb x nloc x nnei x 3.

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nb x nloc x nnei.

sw

The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, and remains 0 beyond rcut, with shape nb x nloc x nnei.

axis_neuron

Size of the submatrix.

smooth

Whether to use smoothness in processes such as attention weights calculation.

epsilon

Protection of 1./nnei.

Returns:

grrg

Atomic invariant rep, with shape nb x nloc x (axis_neuron x ng2)

_update_g2_g1g1(g1: paddle.Tensor, gg1: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor) → paddle.Tensor

Update the g2 using element-wise dot g1_i * g1_j.

Parameters:

g1

Atomic invariant rep, with shape nb x nloc x ng1.

gg1

Neighbor-wise atomic invariant rep, with shape nb x nloc x nnei x ng1.

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nb x nloc x nnei.

sw

The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, and remains 0 beyond rcut, with shape nb x nloc x nnei.

forward(g1_ext: paddle.Tensor, g2: paddle.Tensor, h2: paddle.Tensor, nlist: paddle.Tensor, nlist_mask: paddle.Tensor, sw: paddle.Tensor)

Parameters:

g1_extnf x nall x ng1 extended single-atom channel

g2nf x nloc x nnei x ng2 pair-atom channel, invariant

h2nf x nloc x nnei x 3 pair-atom channel, equivariant

nlistnf x nloc x nnei neighbor list (padded neis are set to 0)

nlist_masknf x nloc x nnei masks of the neighbor list. real nei 1 otherwise 0

swnf x nloc x nnei switch function

Returns:

g1: nf x nloc x ng1 updated single-atom channel

g2: nf x nloc x nnei x ng2 updated pair-atom channel, invariant

h2: nf x nloc x nnei x 3 updated pair-atom channel, equivariant

list_update_res_avg(update_list: list[paddle.Tensor]) → paddle.Tensor

list_update_res_incr(update_list: list[paddle.Tensor]) → paddle.Tensor

list_update_res_residual(update_list: list[paddle.Tensor], update_name: str = 'g1') → paddle.Tensor

list_update(update_list: list[paddle.Tensor], update_name: str = 'g1') → paddle.Tensor

serialize() → dict

Serialize the networks to a dict.

Returns:

dict

The serialized networks.

classmethod deserialize(data: dict) → RepformerLayer

Deserialize the networks from a dict.

Parameters:

datadict

The dict to deserialize from.