deepmd.dpmodel.descriptor.repformers

Classes

DescrptBlockRepformers	The repformer descriptor block.
Atten2Map	The unit operation of a native model.
Atten2MultiHeadApply	The unit operation of a native model.
Atten2EquiVarApply	The unit operation of a native model.
LocalAtten	The unit operation of a native model.
RepformerLayer	The unit operation of a native model.

Functions

xp_transpose_01423(→ deepmd.dpmodel.array_api.Array)
xp_transpose_01342(→ deepmd.dpmodel.array_api.Array)
get_residual(→ deepmd.dpmodel.array_api.Array)	Get residual tensor for one update vector.
_make_nei_g1(→ deepmd.dpmodel.array_api.Array)	Make neighbor-wise atomic invariant rep.
_apply_nlist_mask(→ deepmd.dpmodel.array_api.Array)	Apply nlist mask to neighbor-wise rep tensors.
_apply_switch(→ deepmd.dpmodel.array_api.Array)	Apply switch function to neighbor-wise rep tensors.
_cal_hg(→ deepmd.dpmodel.array_api.Array)	Calculate the transposed rotation matrix.
_cal_grrg(→ deepmd.dpmodel.array_api.Array)	Calculate the atomic invariant rep.
symmetrization_op(→ deepmd.dpmodel.array_api.Array)	Symmetrization operator to obtain atomic invariant rep.

Module Contents

deepmd.dpmodel.descriptor.repformers.xp_transpose_01423(x: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

deepmd.dpmodel.descriptor.repformers.xp_transpose_01342(x: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

class deepmd.dpmodel.descriptor.repformers.DescrptBlockRepformers(rcut: float, rcut_smth: float, sel: int, ntypes: int, nlayers: int = 3, g1_dim: int = 128, g2_dim: int = 16, axis_neuron: int = 4, direct_dist: bool = False, 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', set_davg_zero: bool = True, smooth: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, precision: str = 'float64', trainable_ln: bool = True, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, ln_eps: float | None = 1e-05, seed: int | list[int] | None = None, trainable: bool = True)

Bases: deepmd.dpmodel.NativeOP, deepmd.dpmodel.descriptor.descriptor.DescriptorBlock

The repformer descriptor block.

The repformer block iteratively updates single-atom (\(\mathcal{G}_1\)), pair-atom (\(\mathcal{G}_2\)), and equivariant pair-atom (\(\mathcal{H}_2\)) representations through multiple layers:

Update of :math:`mathcal{G}_1` (single-atom representation):

The update can include multiple terms:

• Convolution term: \(\mathcal{G}_1^{i,l+1} \leftarrow \mathcal{G}_1^{i,l} + \mathrm{MLP}(\sum_j \mathcal{G}_2^{ij,l} \odot \mathcal{G}_1^{j,l})\)

• GRRG term: \(\mathcal{G}_1^{i,l+1} \leftarrow \mathcal{G}_1^{i,l} + \mathrm{MLP}((\mathcal{G}_2^{i,l})^T \mathcal{H}_2^{i,l} (\mathcal{H}_2^{i,l})^T \mathcal{G}_{2,<}^{i,l})\)

• DRRD term: \(\mathcal{G}_1^{i,l+1} \leftarrow \mathcal{G}_1^{i,l} + \mathrm{MLP}((\mathcal{G}_1^{j,l})^T \mathcal{H}_2^{i,l} (\mathcal{H}_2^{i,l})^T \mathcal{G}_{1,<}^{j,l})\)

• Attention term: \(\mathcal{G}_1^{i,l+1} \leftarrow \mathcal{G}_1^{i,l} + \mathrm{SelfAttention}(\mathcal{G}_1^{i,l}, \mathcal{G}_1^{j,l})\)

Update of :math:`mathcal{G}_2` (pair-atom representation):

• G1xG1 term: \(\mathcal{G}_2^{ij,l+1} \leftarrow \mathcal{G}_2^{ij,l} + \mathrm{MLP}(\mathcal{G}_1^{i,l} \otimes \mathcal{G}_1^{j,l})\)

• Attention term: \(\mathcal{G}_2^{ij,l+1} \leftarrow \mathcal{G}_2^{ij,l} + \mathrm{GatedSelfAttention}(\mathcal{G}_2^{ij,l})\)

Update of :math:`mathcal{H}_2` (equivariant pair-atom representation):

\[\mathcal{H}_2^{ij,l+1} = \mathcal{H}_2^{ij,l} + \mathrm{MLP}(\mathcal{G}_2^{ij,l}) \odot \mathcal{R}^{ij}.\]

The final descriptor is the iteratively updated single-atom representation:

\[\mathcal{D}^i = \mathcal{G}_1^{i,L},\]

where \(L\) is the number of repformer layers.

Parameters:

rcutfloat

The cut-off radius.

rcut_smthfloat

Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth.

selint

Maximally possible number of selected neighbors.

ntypesint

Number of element types

nlayersint, optional

Number of repformer layers.

g1_dimint, optional

Dimension of the first graph convolution layer.

g2_dimint, optional

Dimension of the second graph convolution layer.

axis_neuronint, optional

Size of the submatrix of G (embedding matrix).

direct_distbool, optional

Whether to use direct distance information (1/r term) in the repformer block.

update_g1_has_convbool, optional

Whether to update the g1 rep with convolution term.

update_g1_has_drrdbool, optional

Whether to update the g1 rep with the drrd term.

update_g1_has_grrgbool, optional

Whether to update the g1 rep with the grrg term.

update_g1_has_attnbool, optional

Whether to update the g1 rep with the localized self-attention.

update_g2_has_g1g1bool, optional

Whether to update the g2 rep with the g1xg1 term.

update_g2_has_attnbool, optional

Whether to update the g2 rep with the gated self-attention.

update_h2bool, optional

Whether to update the h2 rep.

attn1_hiddenint, optional

The hidden dimension of localized self-attention to update the g1 rep.

attn1_nheadint, optional

The number of heads in localized self-attention to update the g1 rep.

attn2_hiddenint, optional

The hidden dimension of gated self-attention to update the g2 rep.

attn2_nheadint, optional

The number of heads in gated self-attention to update the g2 rep.

attn2_has_gatebool, optional

Whether to use gate in the gated self-attention to update the g2 rep.

activation_functionstr, optional

The activation function in the embedding net.

update_stylestr, optional

Style to update a representation. Supported options are: -‘res_avg’: Updates a rep u with: u = 1/\sqrt{n+1} (u + u_1 + u_2 + … + u_n) -‘res_incr’: Updates a rep u with: u = u + 1/\sqrt{n} (u_1 + u_2 + … + u_n) -‘res_residual’: Updates a rep u with: u = u + (r1*u_1 + r2*u_2 + … + r3*u_n) where r1, r2 … r3 are residual weights defined by update_residual and update_residual_init.

update_residualfloat, optional

When update using residual mode, the initial std of residual vector weights.

update_residual_initstr, optional

When update using residual mode, the initialization mode of residual vector weights.

set_davg_zerobool, optional

Set the normalization average to zero.

precisionstr, optional

The precision of the embedding net parameters.

smoothbool, optional

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

exclude_typeslist[list[int]], optional

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_protectionfloat, optional

Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection.

trainable_lnbool, optional

Whether to use trainable shift and scale weights in layer normalization.

use_sqrt_nneibool, optional

Whether to use the square root of the number of neighbors for symmetrization_op normalization instead of using the number of neighbors directly.

g1_out_convbool, optional

Whether to put the convolutional update of g1 separately outside the concatenated MLP update.

g1_out_mlpbool, optional

Whether to put the self MLP update of g1 separately outside the concatenated MLP update.

ln_epsfloat, optional

The epsilon value for layer normalization.

seedint, optional

The random seed for initialization.

trainablebool, default: True

Whether the block is trainable

rcut

rcut_smth

ntypes

nlayers = 3

nnei

ndescrpt

sel

sec

split_sel

axis_neuron = 4

set_davg_zero = True

g1_dim = 128

g2_dim = 16

update_g1_has_conv = True

update_g1_has_drrd = True

update_g1_has_grrg = True

update_g1_has_attn = True

update_g2_has_g1g1 = True

update_g2_has_attn = True

update_h2 = False

attn1_hidden = 64

attn1_nhead = 4

attn2_has_gate = False

attn2_hidden = 16

attn2_nhead = 4

activation_function = 'tanh'

update_style = 'res_avg'

update_residual = 0.001

update_residual_init = 'norm'

direct_dist = False

act

smooth = True

env_protection = 0.0

precision = 'float64'

trainable_ln = True

use_sqrt_nnei = True

g1_out_conv = True

g1_out_mlp = True

ln_eps = 1e-05

epsilon = 0.0001

g2_embd

layers = []

env_mat

mean

stddev

orig_sel

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_env_protection() → float

Returns the protection of building environment matrix.

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_emb() → int

Returns the embedding dimension g2.

__setitem__(key: str, value: deepmd.dpmodel.array_api.Array) → None

__getitem__(key: str) → deepmd.dpmodel.array_api.Array

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.

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 embedding dimension g2.

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

_exchange_ghosts(g1: deepmd.dpmodel.array_api.Array, mapping_tiled: deepmd.dpmodel.array_api.Array | None, comm_dict: dict | None, nall: int, nloc: int) → deepmd.dpmodel.array_api.Array

Build g1_ext (the ghost-aware single-atom embedding) for the per-layer loop.

Default: array-api gather via the pre-tiled mapping_tiled. comm_dict, nall, nloc are unused in this default impl; they exist so the pt_expt subclass can perform the per-layer MPI ghost exchange (deepmd_export::border_op) when comm_dict is not None.

call(nlist: deepmd.dpmodel.array_api.Array, coord_ext: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array, atype_embd_ext: deepmd.dpmodel.array_api.Array | None = None, mapping: deepmd.dpmodel.array_api.Array | None = None, type_embedding: deepmd.dpmodel.array_api.Array | None = None, comm_dict: dict | None = None) → deepmd.dpmodel.array_api.Array

Forward pass in NumPy implementation.

has_message_passing() → bool

Returns whether the descriptor block has message passing.

has_message_passing_across_ranks() → bool

Returns whether per-layer g1 needs MPI ghost exchange.

Repformers has no use_loc_mapping opt-out; it always passes g1 in [nb, nall, n_dim] layout, so multi-rank always needs cross-rank exchange of the per-atom feature tensor.

need_sorted_nlist_for_lower() → bool

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

classmethod deserialize(data: dict[str, Any]) → DescrptBlockRepformers

Deserialize the descriptor block.

serialize() → dict[str, Any]

Serialize the descriptor block.

deepmd.dpmodel.descriptor.repformers.get_residual(_dim: int, _scale: float, _mode: str = 'norm', trainable: bool = True, precision: str = 'float64', seed: int | list[int] | None = None) → deepmd.dpmodel.array_api.Array

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.

deepmd.dpmodel.descriptor.repformers._make_nei_g1(g1_ext: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Make neighbor-wise atomic invariant rep.

Parameters:

g1_ext

Extended atomic invariant rep, with shape [nf, nall, ng1].

nlist

Neighbor list, with shape [nf, nloc, nnei].

Returns:

gg1: Array

Neighbor-wise atomic invariant rep, with shape [nf, nloc, nnei, ng1].

deepmd.dpmodel.descriptor.repformers._apply_nlist_mask(gg: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Apply nlist mask to neighbor-wise rep tensors.

Parameters:

gg

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

nlist_mask

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

deepmd.dpmodel.descriptor.repformers._apply_switch(gg: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Apply switch function to neighbor-wise rep tensors.

Parameters:

gg

Neighbor-wise rep tensors, with shape [nf, nloc, nnei, 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, nloc, nnei].

deepmd.dpmodel.descriptor.repformers._cal_hg(g: deepmd.dpmodel.array_api.Array, h: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) → deepmd.dpmodel.array_api.Array

Calculate the transposed rotation matrix.

Parameters:

g

Neighbor-wise/Pair-wise invariant rep tensors, with shape [nf, nloc, nnei, ng].

h

Neighbor-wise/Pair-wise equivariant rep tensors, with shape [nf, nloc, nnei, 3].

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape [nf, nloc, 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 [nf, nloc, nnei].

smooth

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

epsilon

Protection of 1./nnei.

use_sqrt_nneibool, optional

Whether to use the square root of the number of neighbors for symmetrization_op normalization instead of using the number of neighbors directly.

Returns:

hg

The transposed rotation matrix, with shape [nf, nloc, 3, ng].

deepmd.dpmodel.descriptor.repformers._cal_grrg(hg: deepmd.dpmodel.array_api.Array, axis_neuron: int) → deepmd.dpmodel.array_api.Array

Calculate the atomic invariant rep.

Parameters:

hg

The transposed rotation matrix, with shape [nf, nloc, 3, ng].

axis_neuron

Size of the submatrix.

Returns:

grrg

Atomic invariant rep, with shape [nf, nloc, (axis_neuron * ng)].

deepmd.dpmodel.descriptor.repformers.symmetrization_op(g: deepmd.dpmodel.array_api.Array, h: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, axis_neuron: int, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) → deepmd.dpmodel.array_api.Array

Symmetrization operator to obtain atomic invariant rep.

Parameters:

g

Neighbor-wise/Pair-wise invariant rep tensors, with shape [nf, nloc, nnei, ng].

h

Neighbor-wise/Pair-wise equivariant rep tensors, with shape [nf, nloc, nnei, 3].

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape [nf, nloc, 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 [nf, nloc, nnei].

axis_neuron

Size of the submatrix.

smooth

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

epsilon

Protection of 1./nnei.

use_sqrt_nneibool, optional

Whether to use the square root of the number of neighbors for symmetrization_op normalization instead of using the number of neighbors directly.

Returns:

grrg

Atomic invariant rep, with shape [nf, nloc, (axis_neuron * ng)].

class deepmd.dpmodel.descriptor.repformers.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, trainable: bool = True)

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim

hidden_dim

head_num

mapqk

has_gate = False

smooth = True

attnw_shift = 20.0

precision = 'float64'

call(g2: deepmd.dpmodel.array_api.Array, h2: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Forward pass in NumPy implementation.

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.dpmodel.descriptor.repformers.Atten2MultiHeadApply(input_dim: int, head_num: int, precision: str = 'float64', seed: int | list[int] | None = None, trainable: bool = True)

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim

head_num

mapv

head_map

precision = 'float64'

call(AA: deepmd.dpmodel.array_api.Array, g2: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Forward pass in NumPy implementation.

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.dpmodel.descriptor.repformers.Atten2EquiVarApply(input_dim: int, head_num: int, precision: str = 'float64', seed: int | list[int] | None = None, trainable: bool = True)

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim

head_num

head_map

precision = 'float64'

call(AA: deepmd.dpmodel.array_api.Array, h2: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Forward pass in NumPy implementation.

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.dpmodel.descriptor.repformers.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, trainable: bool = True)

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim

hidden_dim

head_num

mapq

mapkv

head_map

smooth = True

attnw_shift = 20.0

precision = 'float64'

call(g1: deepmd.dpmodel.array_api.Array, gg1: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Forward pass in NumPy implementation.

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.dpmodel.descriptor.repformers.RepformerLayer(rcut: float, rcut_smth: float, sel: int, ntypes: int, g1_dim: int = 128, g2_dim: int = 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, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, ln_eps: float | None = 1e-05, seed: int | list[int] | None = None, trainable: bool = True)

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

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

use_sqrt_nnei = True

g1_out_conv = True

g1_out_mlp = True

ln_eps = 1e-05

precision = 'float64'

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

g1_residual = []

g2_residual = []

h2_residual = []

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

_update_h2(h2: deepmd.dpmodel.array_api.Array, attn: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

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: deepmd.dpmodel.array_api.Array, g2: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

Calculate the convolution update for atomic invariant rep.

Parameters:

gg1

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

g2

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

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nf 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 nf x nloc x nnei.

_update_g2_g1g1(g1: deepmd.dpmodel.array_api.Array, gg1: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array

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

Parameters:

g1

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

gg1

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

nlist_mask

Neighbor list mask, where zero means no neighbor, with shape nf 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 nf x nloc x nnei.

call(g1_ext: deepmd.dpmodel.array_api.Array, g2: deepmd.dpmodel.array_api.Array, h2: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, nlist_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array]

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[deepmd.dpmodel.array_api.Array]) → deepmd.dpmodel.array_api.Array

list_update_res_incr(update_list: list[deepmd.dpmodel.array_api.Array]) → deepmd.dpmodel.array_api.Array

list_update_res_residual(update_list: list[deepmd.dpmodel.array_api.Array], update_name: str = 'g1') → deepmd.dpmodel.array_api.Array

list_update(update_list: list[deepmd.dpmodel.array_api.Array], update_name: str = 'g1') → deepmd.dpmodel.array_api.Array

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.