deepmd.dpmodel.descriptor.repformers

deepmd.dpmodel.descriptor.repformers#

Classes#

`DescrptBlockRepformers`	The repformer descriptor block.
`Atten2Map`	Masked and smoothed angular attention map.
`Atten2MultiHeadApply`	Multi-head attention application.
`Atten2EquiVarApply`	Equivariant attention application preserving rotation laws.
`LocalAtten`	Local attention with masked, smoothed softmax weights.
`RepformerLayer`	Residual representation update \((G_1,G_2,H_2)^{l+1}=\Phi_l(...)\).

Functions#

`xp_transpose_01423`(→ deepmd.dpmodel.array_api.Array)
`xp_transpose_01342`(→ deepmd.dpmodel.array_api.Array)
`_statically_compatible_shape`(→ bool)
`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.
`_cal_hg_graph`(→ deepmd.dpmodel.array_api.Array)	Graph twin of `_cal_hg()`: hg[n, a, b] = sum_{e: dst(e)=n} h_e[a] g_e[b].
`_cal_grrg_graph`(→ deepmd.dpmodel.array_api.Array)	Graph twin of `_cal_grrg()` (node-local, no neighbor axis).
`symmetrization_op_graph`(→ deepmd.dpmodel.array_api.Array)	Graph twin of `symmetrization_op()`.

Module Contents#

deepmd.dpmodel.descriptor.repformers.xp_transpose_01423(x: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

deepmd.dpmodel.descriptor.repformers.xp_transpose_01342(x: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

deepmd.dpmodel.descriptor.repformers._statically_compatible_shape(actual: tuple, expected: tuple) → bool[source]#

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

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

rcut_smth[source]#

ntypes[source]#

nlayers = 3[source]#

nnei[source]#

ndescrpt[source]#

sel[source]#

sec[source]#

split_sel[source]#

axis_neuron = 4[source]#

set_davg_zero = True[source]#

g1_dim = 128[source]#

g2_dim = 16[source]#

update_g1_has_conv = True[source]#

update_g1_has_drrd = True[source]#

update_g1_has_grrg = True[source]#

update_g1_has_attn = True[source]#

update_g2_has_g1g1 = True[source]#

update_g2_has_attn = True[source]#

update_h2 = False[source]#

attn1_hidden = 64[source]#

attn1_nhead = 4[source]#

attn2_has_gate = False[source]#

attn2_hidden = 16[source]#

attn2_nhead = 4[source]#

activation_function = 'tanh'[source]#

update_style = 'res_avg'[source]#

update_residual = 0.001[source]#

update_residual_init = 'norm'[source]#

direct_dist = False[source]#

act[source]#

smooth = True[source]#

env_protection = 0.0[source]#

precision = 'float64'[source]#

trainable_ln = True[source]#

use_sqrt_nnei = True[source]#

g1_out_conv = True[source]#

g1_out_mlp = True[source]#

ln_eps = 1e-05[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_rcut_smth() → float[source]#: Returns the radius where the neighbor information starts to smoothly decay to 0.

get_env_protection() → float[source]#: Returns the protection of building environment matrix.

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: str, value: deepmd.dpmodel.array_api.Array) → None[source]#

__getitem__(key: str) → deepmd.dpmodel.array_api.Array[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: int[source]#: Returns the output dimension of this descriptor.

property dim_in: int[source]#: Returns the atomic input dimension of this descriptor.

property dim_emb: int[source]#: Returns the embedding dimension g2.

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

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][source]#: Get the statistics of the descriptor.

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

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

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.

_exchange_ghosts_graph(g1: deepmd.dpmodel.array_api.Array, comm_dict: dict | None, n_total: int) → deepmd.dpmodel.array_api.Array[source]#

Per-layer node-channel refresh on the graph path.

Ghost-free graphs (Python single-rank; src IS the local owner) and extended single-process graphs need NO exchange – identity. The pt_expt subclass overrides this to overwrite ghost rows via deepmd_export::border_op when comm_dict is provided (C++ multi-rank extended-region graphs).

Parameters:

g1: Flat node-wise atomic invariant rep, with shape [n_total, ng1].
comm_dict: MPI communication metadata; must be None on this (dpmodel) path.
n_total: Total number of nodes (unused here; the pt_expt override needs it).

Returns:

g1Array: The (unchanged) node channel, with shape [n_total, ng1].

Raises:

NotImplementedError: If comm_dict is not None.

call_graph(graph: Any, atype: deepmd.dpmodel.array_api.Array, g1_input: deepmd.dpmodel.array_api.Array, comm_dict: dict | None = None, static_nnei: int | None = None) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#

Graph twin of call() on the flat node/edge axes.

Parameters:

graph: A NeighborGraph whose edge_index = [src, dst] (src = neighbor local owner, dst = center), edge_vec = r_src - r_dst and edge_mask marks real edges.
atype: (N,) flat node atom types (N = sum(graph.n_node)).
g1_input: (N, g1_dim) node channel BEFORE the block activation – the flat twin of the dense atype_embd (atype_embd_ext sliced to the first nloc atoms); the block applies self.act to it, exactly like the dense body.
comm_dict: MPI communication metadata forwarded to _exchange_ghosts_graph(). None for non-parallel inference (dpmodel default: identity).
static_nnei: When the graph uses the shape-static center-major layout (graph_from_dense_quartet / from_dense_quartet(compact=False), E = n_center * nnei), pass nnei so the attention edge-pair enumeration stays jit/export-traceable (no nonzero). None (carry-all / compact graphs) selects the dynamic eager form.

Returns:

g1Array: (N, dim_out) updated node channel.
g2Array: (E, ng2) updated edge channel, invariant.
h2Array: (E, 3) updated edge channel, equivariant.
rot_matArray: (N, dim_emb, 3) rotationally equivariant single-particle representation.
swArray: (E,) smooth switch, zeroed on padding edges.

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[source]#: Forward pass in NumPy implementation.

has_message_passing() → bool[source]#: Returns whether the descriptor block has message passing.

has_message_passing_across_ranks() → bool[source]#

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[source]#: Returns whether the descriptor block needs sorted nlist when using forward_lower.

classmethod deserialize(data: dict[str, Any]) → DescrptBlockRepformers[source]#: Deserialize the descriptor block.

serialize() → dict[str, Any][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) → deepmd.dpmodel.array_api.Array[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: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[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: 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[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: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[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: 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[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: deepmd.dpmodel.array_api.Array, axis_neuron: int) → deepmd.dpmodel.array_api.Array[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: 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[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)].

deepmd.dpmodel.descriptor.repformers._cal_hg_graph(g: deepmd.dpmodel.array_api.Array, h: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, dst: deepmd.dpmodel.array_api.Array, n_total: int, nnei: int, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of _cal_hg(): hg[n, a, b] = sum_{e: dst(e)=n} h_e[a] g_e[b].

Parameters:

g: Flat edge-wise invariant rep, with shape [n_edge, ng].
h: Flat edge-wise equivariant rep, with shape [n_edge, 3].
edge_mask: Edge mask, where zero means no edge, with shape [n_edge].
sw: The switch function per edge, with shape [n_edge].
dst: Destination (center) node index of each edge, with shape [n_edge].
n_total: Total number of nodes.
nnei: The block sel (the smooth-branch normalization constant, spec decision #9: sel = normalization only).
smooth: Whether to use the smooth (sw-weighted, sel-normalized) branch.
epsilon: Degree-normalization protection for the non-smooth branch.
use_sqrt_nnei: Whether to normalize by sqrt(nnei) instead of nnei.

Returns:

hgArray: Transposed rotation matrix per node, with shape [n_total, 3, ng].

deepmd.dpmodel.descriptor.repformers._cal_grrg_graph(hg: deepmd.dpmodel.array_api.Array, axis_neuron: int) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of _cal_grrg() (node-local, no neighbor axis).

Parameters:

hg: Transposed rotation matrix per node, with shape [n_total, 3, ng].
axis_neuron: Size of the submatrix of hg (embedding matrix).

Returns:

grrgArray: Atomic invariant rep, with shape [n_total, axis_neuron * ng].

deepmd.dpmodel.descriptor.repformers.symmetrization_op_graph(g: deepmd.dpmodel.array_api.Array, h: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, dst: deepmd.dpmodel.array_api.Array, n_total: int, nnei: int, axis_neuron: int, smooth: bool = True, epsilon: float = 0.0001, use_sqrt_nnei: bool = True) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of symmetrization_op().

Parameters:

g: Flat edge-wise invariant rep, with shape [n_edge, ng].
h: Flat edge-wise equivariant rep, with shape [n_edge, 3].
edge_mask: Edge mask, where zero means no edge, with shape [n_edge].
sw: The switch function per edge, with shape [n_edge].
dst: Destination (center) node index of each edge, with shape [n_edge].
n_total: Total number of nodes.
nnei: The block sel (the smooth-branch normalization constant).
axis_neuron: Size of the submatrix of hg (embedding matrix).
smooth: Whether to use the smooth (sw-weighted, sel-normalized) branch.
epsilon: Degree-normalization protection for the non-smooth branch.
use_sqrt_nnei: Whether to normalize by sqrt(nnei) instead of nnei.

Returns:

grrgArray: Atomic invariant rep, with shape [n_total, 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)[source]#

Bases: deepmd.dpmodel.NativeOP

Masked and smoothed angular attention map.

Define the raw logits and the (not necessarily unit-vector) angular factor by

\[L_{ijk}^{(h)}=\frac{q_{ij}^{(h)}\cdot k_{ik}^{(h)}}{\sqrt d}, \qquad G_{ijk}=\mathbf h_{ij}\cdot\mathbf h_{ik}.\]

When has_gate=True, \(G_{ijk}\) also gates the logits, so let \(\widetilde L=L G\); otherwise let \(\widetilde L=L\). In the smooth path the normalized attention and returned map are

\[\overline A_{ijk}^{(h)}=\operatorname{softmax}_k\!\left( (\widetilde L_{ijk}^{(h)}+c)s_{ij}s_{ik}-c\right),\qquad M_{ijk}^{(h)}=s_{ij}s_{ik}\overline A_{ijk}^{(h)} \frac{G_{ijk}}{\sqrt 3}.\]

Without smoothing, masked softmax replaces the first expression and the cutoff factors are omitted. Invalid query and key entries are zeroed. The factor \(G_{ijk}/\sqrt 3\) multiplies the returned map unconditionally, including when has_gate=False.

input_dim[source]#

hidden_dim[source]#

head_num[source]#

mapqk[source]#

has_gate = False[source]#

smooth = True[source]#

attnw_shift = 20.0[source]#

precision = 'float64'[source]#

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[source]#: Forward pass in NumPy implementation.

call_graph(g2: deepmd.dpmodel.array_api.Array, h2: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, q_e: deepmd.dpmodel.array_api.Array, k_e: deepmd.dpmodel.array_api.Array, pair_mask: deepmd.dpmodel.array_api.Array, e_tot: int, sel: int) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of call(): per-pair attention map over edges sharing a center.

The dense (nnei, nnei) square becomes ordered+self edge pairs; the softmax over the key axis becomes a segment_softmax grouped by the query edge. The dense post-softmax query-row AND key-column zeroing both collapse to pair_mask (a pair is real iff BOTH edges are real).

Parameters:

g2: Flat edge-wise pair invariant rep, with shape [n_edge, ng2].
h2: Flat edge-wise pair equivariant rep, with shape [n_edge, 3].
sw: The switch function per edge, with shape [n_edge].
q_e: Query edge index of each pair, with shape [n_pair].
k_e: Key edge index of each pair, with shape [n_pair].
pair_mask: Pair validity (both edges real), with shape [n_pair].
e_tot: Total number of edges (the number of softmax segments).
sel: The block sel: the fixed dense softmax width used for the phantom-count compensation of the smooth branch.

Returns:

attnArray: Attention map on the flat pair axis, with shape [n_pair, nh].

Notes

The smooth branch reproduces dense’s FIXED-WIDTH softmax denominator exactly: dense keeps every one of its sel key slots in the denominator, the non-real ones each at exp(-attnw_shift). Here the masked pairs are excluded from the softmax and replaced by sel - n_real SIGNED phantom denominator terms per query edge, so the phantom COUNT is geometry-independent. Without this, one phantom per PRESENT pair would make the denominator term count change when an edge enters/leaves the graph at the model cutoff – an O(exp(-attnw_shift)) energy/force discontinuity – and differ from dense by O(sel * exp(-attnw_shift)).

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, trainable: bool = True)[source]#

Bases: deepmd.dpmodel.NativeOP

Multi-head attention application.

Atten2Map has already applied masking, smoothing, and softmax. This stage therefore computes the value aggregation and head projection only:

\[O_{ij}^{(h)}=\sum_k A_{ijk}^{(h)}V_{ik}^{(h)},\qquad O_{ij}=\operatorname{HeadMap}\!\left(\operatorname{concat}_h O_{ij}^{(h)}\right).\]

input_dim[source]#

head_num[source]#

mapv[source]#

head_map[source]#

precision = 'float64'[source]#

call(AA: deepmd.dpmodel.array_api.Array, g2: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#: Forward pass in NumPy implementation.

call_graph(AA: deepmd.dpmodel.array_api.Array, g2: deepmd.dpmodel.array_api.Array, q_e: deepmd.dpmodel.array_api.Array, k_e: deepmd.dpmodel.array_api.Array, e_tot: int) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of call(): out[q] = sum_k AA[q,k] g2v[k] per head.

Parameters:

AA: Attention map on the flat pair axis, with shape [n_pair, nh].
g2: Flat edge-wise pair invariant rep, with shape [n_edge, ng2].
q_e: Query edge index of each pair, with shape [n_pair].
k_e: Key edge index of each pair, with shape [n_pair].
e_tot: Total number of edges.

Returns:

g2_newArray: Attention-updated edge channel, with shape [n_edge, ng2].

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, trainable: bool = True)[source]#

Bases: deepmd.dpmodel.NativeOP

Equivariant attention application preserving rotation laws.

For attention matrix \(A_{ijk}^{(h)}\) and equivariant neighbor vectors \(\mathbf h_{ik}\), each head produces

\[\mathbf u_{ij}^{(h)} = \sum_k A_{ijk}^{(h)}\mathbf h_{ik}, \qquad \mathbf u_{ij} = \sum_h w_h\mathbf u_{ij}^{(h)}.\]

The attention weights and head mixing coefficients are scalars, so the output follows the same rotation law as \(\mathbf h\).

input_dim[source]#

head_num[source]#

head_map[source]#

precision = 'float64'[source]#

call(AA: deepmd.dpmodel.array_api.Array, h2: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#: Forward pass in NumPy implementation.

call_graph(AA: deepmd.dpmodel.array_api.Array, h2: deepmd.dpmodel.array_api.Array, q_e: deepmd.dpmodel.array_api.Array, k_e: deepmd.dpmodel.array_api.Array, e_tot: int) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of call() (heads applied to the equivariant channel).

Parameters:

AA: Attention map on the flat pair axis, with shape [n_pair, nh].
h2: Flat edge-wise pair equivariant rep, with shape [n_edge, 3].
q_e: Query edge index of each pair, with shape [n_pair].
k_e: Key edge index of each pair, with shape [n_pair].
e_tot: Total number of edges.

Returns:

h2_newArray: Attention-updated equivariant edge channel, with shape [n_edge, 3].

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, trainable: bool = True)[source]#

Bases: deepmd.dpmodel.NativeOP

Local attention with masked, smoothed softmax weights.

For each head, the query is formed from the destination feature and keys and values from its neighbors. The logits use the standard scaling and the cutoff-smoothed softmax when enabled:

\[L_{ij}^{(h)}=q_i^{(h)}\cdot k_{ij}^{(h)}/\sqrt{d},\qquad A_{ij}^{(h)}=s_{ij}\operatorname{softmax}_j\!\left( (L_{ij}^{(h)}+c)s_{ij}-c\right),\qquad O_i=\operatorname{HeadMap}\!\left(\operatorname{concat}_h \sum_j A_{ij}^{(h)}v_{ij}^{(h)}\right).\]

input_dim[source]#

hidden_dim[source]#

head_num[source]#

mapq[source]#

mapkv[source]#

head_map[source]#

smooth = True[source]#

attnw_shift = 20.0[source]#

precision = 'float64'[source]#

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[source]#: Forward pass in NumPy implementation.

call_graph(g1: deepmd.dpmodel.array_api.Array, gg1: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, dst: deepmd.dpmodel.array_api.Array, n_total: int, sel: int) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of call(): per-center attention over its edges.

The dense softmax over the neighbor axis becomes a segment_softmax over dst.

Parameters:

g1: Flat node-wise atomic invariant rep, with shape [n_total, ng1].
gg1: Neighbor (source-node) g1 gathered per edge, with shape [n_edge, ng1].
edge_mask: Edge mask, where zero means no edge, with shape [n_edge].
sw: The switch function per edge, with shape [n_edge].
dst: Destination (center) node index of each edge, with shape [n_edge].
n_total: Total number of nodes (the number of softmax segments).
sel: The block sel: the fixed dense softmax width used for the phantom-count compensation of the smooth branch.

Returns:

g1_attnArray: Attention-updated node channel, with shape [n_total, ng1].

Notes

The smooth branch uses the dense-width phantom compensation (see Atten2Map.call_graph()): masked edges are excluded from the softmax and sel - n_real SIGNED phantom denominator terms at exp(-attnw_shift) are added per center, keeping the denominator term count geometry-independent (continuous at the cutoff, and equal to the dense softmax at non-binding sel).

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

Bases: deepmd.dpmodel.NativeOP

Residual representation update \((G_1,G_2,H_2)^{l+1}=\Phi_l(...)\).

epsilon = 0.0001[source]#

rcut[source]#

rcut_smth[source]#

ntypes[source]#

nnei[source]#

sel[source]#

sec[source]#

axis_neuron = 4[source]#

activation_function = 'tanh'[source]#

act[source]#

update_g1_has_grrg = True[source]#

update_g1_has_drrd = True[source]#

update_g1_has_conv = True[source]#

update_g1_has_attn = True[source]#

update_chnnl_2 = True[source]#

update_g2_has_g1g1 = True[source]#

update_g2_has_attn = True[source]#

update_h2 = False[source]#

attn1_hidden = 64[source]#

attn1_nhead = 4[source]#

attn2_hidden = 16[source]#

attn2_nhead = 4[source]#

attn2_has_gate = False[source]#

update_style = 'res_avg'[source]#

update_residual = 0.001[source]#

update_residual_init = 'norm'[source]#

smooth = True[source]#

g1_dim = 128[source]#

g2_dim = 16[source]#

trainable_ln = True[source]#

use_sqrt_nnei = True[source]#

g1_out_conv = True[source]#

g1_out_mlp = True[source]#

ln_eps = 1e-05[source]#

precision = 'float64'[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: deepmd.dpmodel.array_api.Array, attn: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[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: 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[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_g1_conv_graph(gg1: deepmd.dpmodel.array_api.Array, g2: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array, dst: deepmd.dpmodel.array_api.Array, n_total: int, nnei: int) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of _update_g1_conv() (segment_sum over dst).

Parameters:

gg1: Edge-wise atomic invariant rep, with shape [n_edge, ng1].
g2: Edge-wise pair invariant rep, with shape [n_edge, ng2].
edge_mask: Edge mask, where zero means no edge, with shape [n_edge].
sw: The switch function, with shape [n_edge].
dst: Destination (center) node index of each edge, with shape [n_edge].
n_total: Total number of nodes.
nnei: The block sel: the smooth-branch normalization constant AND the fixed dense softmax width for the attention phantom-count compensation (see Atten2Map.call_graph()).

Returns:

g1_convArray: Convolution-updated node channel, with shape [n_total, ng1].

_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[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.

_update_g2_g1g1_graph(g1: deepmd.dpmodel.array_api.Array, src: deepmd.dpmodel.array_api.Array, dst: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, sw: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

Graph twin of _update_g2_g1g1(): per-edge g1_center * g1_neighbor.

Parameters:

g1: Atomic invariant rep, with shape [n_total, ng1].
src: Source (neighbor) node index of each edge, with shape [n_edge].
dst: Destination (center) node index of each edge, with shape [n_edge].
edge_mask: Edge mask, where zero means no edge, with shape [n_edge].
sw: The switch function, with shape [n_edge].

Returns:

g1g1Array: Per-edge center-neighbor g1 product, with shape [n_edge, ng1].

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][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

Graph twin of call(). One residual update of (g1, g2, h2) on the flat node/edge axes; the neighbor axis of every dense reduction is replaced by segment ops over dst.

Parameters:

g1: Flat node-wise atomic invariant rep, with shape [n_total, ng1]. Unlike call(), this is the FULL node channel with no local/ghost slice: ghost-free graphs have no ghosts, and extended (multi-rank) graphs deliberately compute ghost rows here and let the later communication step overwrite them.
g2: Flat edge-wise pair invariant rep, with shape [n_edge, ng2].
h2: Flat edge-wise pair equivariant rep, with shape [n_edge, 3].
src: Source (neighbor) node index of each edge, with shape [n_edge].
dst: Destination (center) node index of each edge, with shape [n_edge].
edge_mask: Edge mask, where zero means no edge, with shape [n_edge].
sw: The switch function, with shape [n_edge].
n_total: Total number of nodes.
nnei: The block sel: the smooth-branch normalization constant AND the fixed dense softmax width for the attention phantom-count compensation (see Atten2Map.call_graph()).
pairs: (q_e, k_e, pair_mask) from center_edge_pairs(). Required iff self.update_g2_has_attn or self.update_h2.

Returns:

g1_new[n_total, ng1] updated node channel
g2_new[n_edge, ng2] updated edge channel, invariant
h2_new[n_edge, 3] updated edge channel, equivariant

list_update_res_avg(update_list: list[deepmd.dpmodel.array_api.Array]) → deepmd.dpmodel.array_api.Array[source]#

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

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

list_update(update_list: list[deepmd.dpmodel.array_api.Array], update_name: str = 'g1') → deepmd.dpmodel.array_api.Array[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.