deepmd.pt_expt.descriptor.dpa1

Attributes

_DESTINATION_CSR_FIELDS
_DUAL_CSR_FIELDS

Classes

DescrptDPA1

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

Functions

_has_graph_fields(→ bool)
_env_mat(→ tuple)	Environment-matrix prologue shared by every fused path (strip / concat).
_strip_pair_index(→ torch.Tensor)	Strip-mode per-edge row index into the type-pair embedding table.
_concat_embedding_input(→ torch.Tensor)	Concat-mode embedding input: radial channel concatenated with the
_grrg_from_moment(→ Any)	Strip-mode epilogue: symmetry-invariant contraction of the moment.
_type_pair_table(→ torch.Tensor)	Type-pair embedding table (P, ng) gathered per edge by tebd_idx.

Module Contents

deepmd.pt_expt.descriptor.dpa1._has_graph_fields(graph: Any, fields: tuple[str, Ellipsis]) → bool

deepmd.pt_expt.descriptor.dpa1._DESTINATION_CSR_FIELDS = ('destination_order', 'destination_row_ptr')

deepmd.pt_expt.descriptor.dpa1._DUAL_CSR_FIELDS = ('destination_order', 'destination_row_ptr', 'source_order', 'source_row_ptr')

deepmd.pt_expt.descriptor.dpa1._env_mat(desc: Any, coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → tuple

Environment-matrix prologue shared by every fused path (strip / concat).

Returns (nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding): rr the (nfnl, nnei, 4) environment matrix (excluded edges zeroed), ss its radial channel (nfnl, nnei, 1), sw the (nfnl, nnei, 1) smooth cutoff (zeroed on excluded/padding edges), and nlist_masked the neighbor indices with excluded/padding entries mapped to 0 (for downstream gathers).

deepmd.pt_expt.descriptor.dpa1._strip_pair_index(desc: Any, atype_ext: torch.Tensor, nlist_masked: torch.Tensor, type_embedding: torch.Tensor, nf: int, nloc: int, nnei: int) → torch.Tensor

Strip-mode per-edge row index into the type-pair embedding table.

One-side uses the neighbor type; two-side folds the (center, neighbor) pair into a flat index, mirroring the dense reference. Shape (nfnl*nnei,).

deepmd.pt_expt.descriptor.dpa1._concat_embedding_input(desc: Any, ss: torch.Tensor, atype_ext: torch.Tensor, nlist_masked: torch.Tensor, type_embedding: torch.Tensor, nf: int, nloc: int, nnei: int) → torch.Tensor

Concat-mode embedding input: radial channel concatenated with the neighbor (and, two-side, center) type embeddings.

Returns (nfnl, nnei, 1 + k * tebd_dim) with k in {1, 2} (one-side / two-side), matching the dense concat reference exactly.

deepmd.pt_expt.descriptor.dpa1._grrg_from_moment(desc: Any, xyz_scatter: torch.Tensor, type_embedding: torch.Tensor, coord_ext: torch.Tensor, atype_ext: torch.Tensor, nf: int, nloc: int, nnei: int, ng: int, sw: torch.Tensor) → Any

Strip-mode epilogue: symmetry-invariant contraction of the moment.

Consumes the unnormalized moment xyz_scatter (nfnl, 4, ng), applies the 1 / nnei normalization, forms the G^T G descriptor and the rotation matrix, and appends the center type embedding when concat_output_tebd.

deepmd.pt_expt.descriptor.dpa1._type_pair_table(desc: Any, type_embedding: torch.Tensor) → torch.Tensor

Type-pair embedding table (P, ng) gathered per edge by tebd_idx.

One-side keeps the neighbor-type rows; two-side forms every (neighbor, center) pair, mirroring the dense reference.

class deepmd.pt_expt.descriptor.dpa1.DescrptDPA1(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = None, trainable_ln: bool = True, ln_eps: float | None = 1e-05, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, spin: None = None, stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.descriptor.dpa1.DescrptDPA1

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

This descriptor, \(\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}\), is given by

\[\mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<,\]

where \(\hat{\mathcal{G}}^i\) represents the embedding matrix:math:mathcal{G}^i after additional self-attention mechanism and \(\mathcal{R}^i\) is defined by the full case in the se_e2_a descriptor. Note that we obtain \(\mathcal{G}^i\) using the type embedding method by default in this descriptor.

To perform the self-attention mechanism, the queries \(\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), keys \(\mathcal{K}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), and values \(\mathcal{V}^{i,l} \in \mathbb{R}^{N_c\times d_v}\) are first obtained:

\[\left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

\[\left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

\[\left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

where \(Q_{l}\), \(K_{l}\), \(V_{l}\) represent three trainable linear transformations that output the queries and keys of dimension \(d_k\) and values of dimension \(d_v\), and \(l\) is the index of the attention layer. The input embedding matrix to the attention layers, denoted by \(\mathcal{G}^{i,0}\), is chosen as the two-body embedding matrix.

Then the scaled dot-product attention method is adopted:

\[A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l},\]

where \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) \in \mathbb{R}^{N_c\times N_c}\) is attention weights. In the original attention method, one typically has \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}\right)=\mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)\), with \(\sqrt{d_{k}}\) being the normalization temperature. This is slightly modified to incorporate the angular information:

\[\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) = \mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right) \odot \hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T},\]

where \(\hat{\mathcal{R}}^{i} \in \mathbb{R}^{N_c\times 3}\) denotes normalized relative coordinates,

\(\hat{\mathcal{R}}^{i}_{j} = \frac{\boldsymbol{r}_{ij}}{\lVert \boldsymbol{r}_{ij} \lVert}\) and \(\odot\) means element-wise multiplication.

Then layer normalization is added in a residual way to finally obtain the self-attention local embedding matrix

\(\hat{\mathcal{G}}^{i} = \mathcal{G}^{i,L_a}\) after \(L_a\) attention layers:[^1]

\[\mathcal{G}^{i,l} = \mathcal{G}^{i,l-1} + \mathrm{LayerNorm}(A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})).\]

Parameters:

rcut: float

The cut-off radius \(r_c\)

rcut_smth: float

From where the environment matrix should be smoothed \(r_s\)

sellist[int], int

list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius

ntypesint

Number of element types

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron: int

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

tebd_dim: int

Dimension of the type embedding

tebd_input_mode: str

The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the radial embedding network output.

resnet_dt: bool

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable: bool

If the weights of this descriptors are trainable.

trainable_ln: bool

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

ln_eps: float, Optional

The epsilon value for layer normalization.

type_one_side: bool

If ‘False’, type embeddings of both neighbor and central atoms are considered. If ‘True’, only type embeddings of neighbor atoms are considered. Default is ‘False’.

attn: int

Hidden dimension of the attention vectors

attn_layer: int

Number of attention layers

attn_dotr: bool

If dot the angular gate to the attention weights

attn_mask: bool

(Only support False to keep consistent with other backend references.) (Not used in this version. True option is not implemented.) If mask the diagonal of attention weights

exclude_typeslist[list[int]]

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

Protection parameter to prevent division by zero errors during environment matrix calculations.

set_davg_zero: bool

Set the shift of embedding net input to zero.

activation_function: str

The activation function in the embedding net. Supported options are “silut”, “softplus”, “relu”, “sigmoid”, “gelu_tf”, “silu”, “linear”, “none”, “tanh”, “gelu”, “relu6”.

precision: str

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float64”, “bfloat16”, “float32”.

scaling_factor: float

The scaling factor of normalization in calculations of attention weights. If temperature is None, the scaling of attention weights is (N_dim * scaling_factor)**0.5

normalize: bool

Whether to normalize the hidden vectors in attention weights calculation.

temperature: float

If not None, the scaling of attention weights is temperature itself.

smooth_type_embedding: bool

Whether to use smooth process in attention weights calculation.

concat_output_tebd: bool

Whether to concat type embedding at the output of the descriptor.

stripped_type_embedding: bool, Optional

(Deprecated, kept only for compatibility.) Whether to strip the type embedding into a separate embedding network. Setting this parameter to True is equivalent to setting tebd_input_mode to ‘strip’. Setting it to False is equivalent to setting tebd_input_mode to ‘concat’. The default value is None, which means the tebd_input_mode setting will be used instead.

use_econf_tebd: bool, Optional

Whether to use electronic configuration type embedding.

use_tebd_biasbool, Optional

Whether to use bias in the type embedding layer.

type_map: list[str], Optional

A list of strings. Give the name to each type of atoms.

spin

(Only support None to keep consistent with other backend references.) (Not used in this version. Not-none option is not implemented.) The old implementation of deepspin.

References

[1]

Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. 2022. DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. arXiv preprint arXiv:2208.08236.

_update_sel_cls

share_params(base_class: Any, shared_level: int, model_prob: float = 1.0, resume: bool = False) → None

Share parameters with base_class for multi-task training.

Level 0: share type_embedding and se_atten (all modules and buffers). Level 1: share type_embedding only.

enable_compression(min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1) → None

Enable compression for the DPA1 descriptor.

Parameters:

min_nbor_dist

The nearest distance between atoms

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

_store_compress_data() → None

Store tabulated data as buffers for the compressed geometric embedding.

_store_type_embd_data() → None

Precompute type embedding outputs and store as a buffer.

call(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, comm_dict: dict | None = None, charge_spin: torch.Tensor | None = None) → Any

Compute the descriptor.

Parameters:

coord_ext

The extended coordinates of atoms. shape: nf x (nallx3)

atype_ext

The extended aotm types. shape: nf x nall

nlist

The neighbor list. shape: nf x nloc x nnei

mapping

The index mapping from extended to local region. not used by this descriptor.

Returns:

descriptor

The descriptor. shape: nf x nloc x (ng x axis_neuron)

gr

The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3

g2

The rotationally invariant pair-partical representation. this descriptor returns None

h2

The rotationally equivariant pair-partical representation. this descriptor returns None

sw

The smooth switch function.

call_graph(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor | None = None, static_nnei: int | None = None) → Any

Graph-native forward, routed through the fused edge kernel when eligible.

Concat and strip attention-free blocks are all graph-eligible. A geometrically compressed descriptor takes the tabulated route (_call_graph_cuda_compress() at DP_CUDA_INFER >= 1, the original fused table operator through _call_graph_compress_reference() at level 0); an uncompressed descriptor takes the embedding route (_call_graph_cuda() mega kernel, the Triton edge convolution for concat, or the dpmodel reference). static_nnei is the shape-static neighbor count the dense call() adapter supplies for the attention edge-pair enumeration; it is forwarded to the dpmodel reference and unused by the attention-free fused paths.

Unlike the dense call(), the graph lower is not @cast_precision-wrapped: edge_vec arrives in the model-agnostic fp64 .pt2 ABI. The fused CUDA kernels consume it in that leaf dtype directly (casting to the model precision in-kernel); the Triton and reference paths need the Python-side alignment below so their environment matrix runs at the model precision. Force / virial precision is decided downstream by the model output transform, not here.

_fused_eligible(backend: str) → bool

Whether a fused descriptor kernel can serve this block.

Parameters:

backendstr

"cuda" for the mega kernel (cuda.dpa1.graph_descriptor) or "triton" for the environment convolution (triton.dpa1). Both require the attention-free path, both tebd input modes (concat / strip) and an activation the kernel inlines (tanh / silu; see activation.ACT_CODES). cuda additionally requires no excluded types, fp32 statistics, the compiled operator library, and – since its MLP runs entirely in registers with template-specialized widths – a three-layer embedding where each layer keeps or doubles the previous width, with neuron[0] in {8, 16, 32, 64}, neuron[1] <= 64 and neuron[2] <= 128, and one inlined activation on every layer (the strip gate network is unconstrained: its table is precomputed outside the kernel). triton serves any layer stack (the head layers run on cuBLAS), so only the last activation must be inlined.

_call_triton(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → Any

Fused Triton environment convolution (attn-free strip / concat path).

The two embedding GEMMs stay on cuBLAS; the last layer’s activation and residual, the smooth cutoff and the moment reduction collapse into se_conv(), which never materializes an (E, ng) tensor. Strip additionally folds the type-pair gate (gathered inline); concat instead feeds the type feature through the embedding input and applies no gate. Composes under make_fx / torch.export so the operator is baked into the pt_expt .pt2.

_call_compressed(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → Any

Compressed forward for DPA1 descriptor (strip only).

_call_graph_compress_reference(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Reference geo-compressed strip descriptor on the edge stream (attn-free).

The DP_CUDA_INFER == 0 graph path and the CPU / trace-time fallback. Evaluates the tabulated geometric embedding with the original fused table operator deepmd::tabulate_fusion_se_atten, treating each edge as a one-neighbor block (nloc = E, nnei = 1) so the operator returns the per-edge moment outer product (E, 4, ng); a segment_sum over edge centers then forms the per-node moment. This matches the dense DescrptDPA1._call_compressed() (same table, same gate) to the fp32 summation-order floor, and composes under autograd so the force / virial assembly differentiates through it. Not make_fx-traceable (the table operator has no meta kernel); the exportable graph is the level >= 1 CUDA kernel.

Parameters:

graphNeighborGraph

Lowered neighbor graph; edge_vec is in the model precision.

atypetorch.Tensor

Flat node atom types with shape (N,), int64.

type_embeddingtorch.Tensor

Type embedding table with shape (ntypes + 1, tebd_dim).

Returns:

grrgtorch.Tensor

Descriptor with shape (N, ng * axis [+ tebd_dim]).

rot_mattorch.Tensor

Equivariant rotation matrix with shape (N, ng, 3).

_call_graph_cuda_compress(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Fused CUDA graph-native geo-compressed strip descriptor (attn-free).

Numerically equivalent to DescrptDPA1._call_compressed() through the dpa1_graph_compress() operator: the environment matrix, quintic table lookup, strip type-pair gate, moment reduction and G^T G contraction collapse into one CUDA mega kernel whose registered backward exposes the edge_vec gradient for the analytic force / virial assembly.

_call_graph_cuda(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Fused CUDA graph-native descriptor (concat, attn-free).

Numerically equivalent to DescrptDPA1DP.call_graph() through the dpa1_graph_descriptor() operator: the environment matrix, embedding MLP, moment reduction and G^T G contraction collapse into one CUDA mega kernel whose registered backward exposes the edge_vec gradient for the analytic force / virial assembly. Composes under make_fx / torch.export so the operator is baked into the pt_expt graph-form .pt2.

fused_energy_force_graph(fit: Any, graph: Any, atype: torch.Tensor, ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool) → tuple[torch.Tensor, Ellipsis] | None

End-to-end fused energy / force / virial from the edge stream.

Collapses this descriptor, the energy fitting and the analytic force / virial assembly into one value-returning CUDA operator (no autograd tape). Returns (energy, atom_energy, force, virial, atom_virial), or None when the descriptor or fitting is not fused-eligible or the operator library is unavailable – the caller then uses the autograd lower. The geo-compressed descriptor dispatches to its tabulated operator (dpa1_graph_compress_energy_force()); the embedding-MLP descriptor to dpa1_graph_energy_force().

Parameters:

fitEnergyFittingNet

The energy fitting module fused with this descriptor.

graphNeighborGraph

The lowered neighbor graph (edge_vec, edge_index, edge_mask, n_node).

atypetorch.Tensor

Flat node atom types with shape (N,), int64.

ownershiptorch.Tensor

Energy-contributing node mask with shape (N,), bool.

atom_biastorch.Tensor

Combined fitting and atomic-model bias with shape (ntypes,).

do_atomic_virialbool

Whether to also assemble the per-atom virial.

Returns:

tuple[torch.Tensor, …] or None

(energy, atom_energy, force, virial, atom_virial), or None.

_call_graph_triton(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Fused graph-native environment convolution (concat / strip, attn-free).

Bit-exact analogue of DescrptDPA1DP.call_graph(): builds the per-edge environment matrix and embedding input, runs the two embedding GEMMs on cuBLAS, then folds the last layer, the type-pair gate (strip), the edge mask, the outer product and the segment_sum scatter into edge_conv(). The two tebd-input modes differ only in the embedding input (concat folds the type feature in; strip runs on the radial channel alone) and the gate (strip multiplies by the type-pair table, concat does not). Composes under make_fx / torch.export so the operator is baked into the pt_expt graph-form .pt2.