deepmd.dpmodel.descriptor.repformers

Contents

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(x)

xp_transpose_01342(x)

get_residual(→ numpy.ndarray)

Get residual tensor for one update vector.

_make_nei_g1(→ numpy.ndarray)

Make neighbor-wise atomic invariant rep.

_apply_nlist_mask(→ numpy.ndarray)

Apply nlist mask to neighbor-wise rep tensors.

_apply_switch(→ numpy.ndarray)

Apply switch function to neighbor-wise rep tensors.

_cal_hg(→ numpy.ndarray)

Calculate the transposed rotation matrix.

_cal_grrg(→ numpy.ndarray)

Calculate the atomic invariant rep.

symmetrization_op(→ numpy.ndarray)

Symmetrization operator to obtain atomic invariant rep.

Module Contents#

deepmd.dpmodel.descriptor.repformers.xp_transpose_01423(x)[source]#
deepmd.dpmodel.descriptor.repformers.xp_transpose_01342(x)[source]#
class deepmd.dpmodel.descriptor.repformers.DescrptBlockRepformers(rcut, rcut_smth, sel: int, ntypes: int, nlayers: int = 3, g1_dim=128, g2_dim=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)[source]#

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

The repformer descriptor block.

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, r2r3 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.

rcut[source]#
rcut_smth[source]#
ntypes[source]#
nlayers[source]#
nnei[source]#
ndescrpt[source]#
sel[source]#
sec[source]#
split_sel[source]#
axis_neuron[source]#
set_davg_zero[source]#
g1_dim[source]#
g2_dim[source]#
update_g1_has_conv[source]#
update_g1_has_drrd[source]#
update_g1_has_grrg[source]#
update_g1_has_attn[source]#
update_g2_has_g1g1[source]#
update_g2_has_attn[source]#
update_h2[source]#
attn1_hidden[source]#
attn1_nhead[source]#
attn2_has_gate[source]#
attn2_hidden[source]#
attn2_nhead[source]#
activation_function[source]#
update_style[source]#
update_residual[source]#
update_residual_init[source]#
direct_dist[source]#
act[source]#
smooth[source]#
env_protection[source]#
precision[source]#
trainable_ln[source]#
use_sqrt_nnei[source]#
g1_out_conv[source]#
g1_out_mlp[source]#
ln_eps[source]#
epsilon = 0.0001[source]#
g2_embd[source]#
layers[source]#
env_mat[source]#
mean[source]#
stddev[source]#
orig_sel[source]#
get_rcut() float[source]#

Returns the cut-off radius.

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_emb() int[source]#

Returns the embedding dimension g2.

__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.

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

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

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

abstract get_stats() NoReturn[source]#

Get the statistics of the descriptor.

reinit_exclude(exclude_types: list[tuple[int, int]] = []) None[source]#
call(nlist: numpy.ndarray, coord_ext: numpy.ndarray, atype_ext: numpy.ndarray, atype_embd_ext: numpy.ndarray | None = None, mapping: numpy.ndarray | None = None, type_embedding: numpy.ndarray | None = None)[source]#

Forward pass in NumPy implementation.

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.

classmethod deserialize(data)[source]#

Deserialize the descriptor block.

serialize()[source]#

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) numpy.ndarray[source]#

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: numpy.ndarray, nlist: numpy.ndarray) numpy.ndarray[source]#

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: np.ndarray

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

deepmd.dpmodel.descriptor.repformers._apply_nlist_mask(gg: numpy.ndarray, nlist_mask: numpy.ndarray) numpy.ndarray[source]#

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: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray[source]#

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: numpy.ndarray, h: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) numpy.ndarray[source]#

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: numpy.ndarray, axis_neuron: int) numpy.ndarray[source]#

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: numpy.ndarray, h: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray, axis_neuron: int, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) numpy.ndarray[source]#

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)[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim[source]#
hidden_dim[source]#
head_num[source]#
mapqk[source]#
has_gate[source]#
smooth[source]#
attnw_shift[source]#
precision[source]#
call(g2: numpy.ndarray, h2: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray[source]#

Forward pass in NumPy implementation.

serialize() dict[source]#

Serialize the networks to a dict.

Returns:
dict

The serialized networks.

classmethod deserialize(data: dict) Atten2Map[source]#

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)[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim[source]#
head_num[source]#
mapv[source]#
head_map[source]#
precision[source]#
call(AA: numpy.ndarray, g2: numpy.ndarray) numpy.ndarray[source]#

Forward pass in NumPy implementation.

serialize() dict[source]#

Serialize the networks to a dict.

Returns:
dict

The serialized networks.

classmethod deserialize(data: dict) Atten2MultiHeadApply[source]#

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)[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim[source]#
head_num[source]#
head_map[source]#
precision[source]#
call(AA: numpy.ndarray, h2: numpy.ndarray) numpy.ndarray[source]#

Forward pass in NumPy implementation.

serialize() dict[source]#

Serialize the networks to a dict.

Returns:
dict

The serialized networks.

classmethod deserialize(data: dict) Atten2EquiVarApply[source]#

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)[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

input_dim[source]#
hidden_dim[source]#
head_num[source]#
mapq[source]#
mapkv[source]#
head_map[source]#
smooth[source]#
attnw_shift[source]#
precision[source]#
call(g1: numpy.ndarray, gg1: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray[source]#

Forward pass in NumPy implementation.

serialize() dict[source]#

Serialize the networks to a dict.

Returns:
dict

The serialized networks.

classmethod deserialize(data: dict) LocalAtten[source]#

Deserialize the networks from a dict.

Parameters:
datadict

The dict to deserialize from.

class deepmd.dpmodel.descriptor.repformers.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, 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)[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

epsilon = 0.0001[source]#
rcut[source]#
rcut_smth[source]#
ntypes[source]#
nnei[source]#
sel[source]#
sec[source]#
axis_neuron[source]#
activation_function[source]#
act[source]#
update_g1_has_grrg[source]#
update_g1_has_drrd[source]#
update_g1_has_conv[source]#
update_g1_has_attn[source]#
update_chnnl_2[source]#
update_g2_has_g1g1[source]#
update_g2_has_attn[source]#
update_h2[source]#
attn1_hidden[source]#
attn1_nhead[source]#
attn2_hidden[source]#
attn2_nhead[source]#
attn2_has_gate[source]#
update_style[source]#
update_residual[source]#
update_residual_init[source]#
smooth[source]#
g1_dim[source]#
g2_dim[source]#
trainable_ln[source]#
use_sqrt_nnei[source]#
g1_out_conv[source]#
g1_out_mlp[source]#
ln_eps[source]#
precision[source]#
linear1[source]#
linear2 = None[source]#
proj_g1g2 = None[source]#
proj_g1g1g2 = None[source]#
attn2g_map = None[source]#
attn2_mh_apply = None[source]#
attn2_lm = None[source]#
attn2_ev_apply = None[source]#
loc_attn = None[source]#
g1_residual[source]#
g2_residual[source]#
h2_residual[source]#
cal_1_dim(g1d: int, g2d: int, ax: int) int[source]#
_update_h2(h2: numpy.ndarray, attn: numpy.ndarray) numpy.ndarray[source]#

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: numpy.ndarray, g2: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray[source]#

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: numpy.ndarray, gg1: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray[source]#

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: numpy.ndarray, g2: numpy.ndarray, h2: numpy.ndarray, nlist: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray)[source]#
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[numpy.ndarray]) numpy.ndarray[source]#
list_update_res_incr(update_list: list[numpy.ndarray]) numpy.ndarray[source]#
list_update_res_residual(update_list: list[numpy.ndarray], update_name: str = 'g1') numpy.ndarray[source]#
list_update(update_list: list[numpy.ndarray], update_name: str = 'g1') numpy.ndarray[source]#
serialize() dict[source]#

Serialize the networks to a dict.

Returns:
dict

The serialized networks.

classmethod deserialize(data: dict) RepformerLayer[source]#

Deserialize the networks from a dict.

Parameters:
datadict

The dict to deserialize from.