deepmd.dpmodel.descriptor.repformers#
Classes#
The repformer descriptor block. | |
The unit operation of a native model. | |
The unit operation of a native model. | |
The unit operation of a native model. | |
The unit operation of a native model. | |
The unit operation of a native model. |
Functions#
| Get residual tensor for one update vector. |
| Make neighbor-wise atomic invariant rep. |
| Apply nlist mask to neighbor-wise rep tensors. |
| Apply switch function to neighbor-wise rep tensors. |
| Calculate the transposed rotation matrix. |
| Calculate the atomic invariant rep. |
| Symmetrization operator to obtain atomic invariant rep. |
Module Contents#
- 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:
- rcut
float
The cut-off radius.
- rcut_smth
float
Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth.
- sel
int
Maximally possible number of selected neighbors.
- ntypes
int
Number of element types
- nlayers
int
,optional
Number of repformer layers.
- g1_dim
int
,optional
Dimension of the first graph convolution layer.
- g2_dim
int
,optional
Dimension of the second graph convolution layer.
- axis_neuron
int
,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_hidden
int
,optional
The hidden dimension of localized self-attention to update the g1 rep.
- attn1_nhead
int
,optional
The number of heads in localized self-attention to update the g1 rep.
- attn2_hidden
int
,optional
The hidden dimension of gated self-attention to update the g2 rep.
- attn2_nhead
int
,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_function
str
,optional
The activation function in the embedding net.
- update_style
str
,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_residual
float
,optional
When update using residual mode, the initial std of residual vector weights.
- update_residual_init
str
,optional
When update using residual mode, the initialization mode of residual vector weights.
- set_davg_zerobool,
optional
Set the normalization average to zero.
- precision
str
,optional
The precision of the embedding net parameters.
- smoothbool,
optional
Whether to use smoothness in processes such as attention weights calculation.
- exclude_types
list
[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_protection
float
,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_eps
float
,optional
The epsilon value for layer normalization.
- seed
int
,optional
The random seed for initialization.
- rcut
- 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.
- 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.
- 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.
- 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:
- _dim
int
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.
- _dim
- 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].
- gg1:
- 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.
- call(g2: numpy.ndarray, h2: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray [source]#
Forward pass in NumPy implementation.
- 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.
- call(AA: numpy.ndarray, g2: numpy.ndarray) numpy.ndarray [source]#
Forward pass in NumPy implementation.
- classmethod deserialize(data: dict) Atten2MultiHeadApply [source]#
Deserialize the networks from a dict.
- Parameters:
- data
dict
The dict to deserialize from.
- data
- 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.
- call(AA: numpy.ndarray, h2: numpy.ndarray) numpy.ndarray [source]#
Forward pass in NumPy implementation.
- classmethod deserialize(data: dict) Atten2EquiVarApply [source]#
Deserialize the networks from a dict.
- Parameters:
- data
dict
The dict to deserialize from.
- data
- 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.
- call(g1: numpy.ndarray, gg1: numpy.ndarray, nlist_mask: numpy.ndarray, sw: numpy.ndarray) numpy.ndarray [source]#
Forward pass in NumPy implementation.
- classmethod deserialize(data: dict) LocalAtten [source]#
Deserialize the networks from a dict.
- Parameters:
- data
dict
The dict to deserialize from.
- data
- 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.
- _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_ext
nf
x
nall
x
ng1
extended
single-atomchannel
- g2
nf
x
nloc
x
nnei
x
ng2
pair-atomchannel
,invariant
- h2
nf
x
nloc
x
nnei
x
3 pair-atomchannel
,equivariant
- nlist
nf
x
nloc
x
nnei
neighbor
list
(padded
neis
are
set
to
0) - nlist_mask
nf
x
nloc
x
nnei
masks
of
the
neighbor
list.real
nei
1otherwise
0 - sw
nf
x
nloc
x
nnei
switch
function
- g1_ext
- Returns:
- 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]#
- classmethod deserialize(data: dict) RepformerLayer [source]#
Deserialize the networks from a dict.
- Parameters:
- data
dict
The dict to deserialize from.
- data