deepmd.pt_expt.descriptor.dpa2

Classes

DescrptDPA2

The DPA-2 descriptor[Rf0ed3afb961b-1]_.

Module Contents

class deepmd.pt_expt.descriptor.dpa2.DescrptDPA2(ntypes: int, repinit: RepinitArgs | dict, repformer: RepformerArgs | dict, concat_output_tebd: bool = True, precision: str = 'float64', smooth: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None)

Bases: deepmd.dpmodel.descriptor.dpa2.DescrptDPA2

The DPA-2 descriptor[Rf0ed3afb961b-1]_.

The DPA-2 descriptor combines a repinit block and a repformer block to extract atomic representations. The overall descriptor is computed as:

\[\mathcal{D}^i = \mathrm{Repformer}(\mathrm{Linear}(\mathrm{Repinit}(\mathcal{R}^i, \mathcal{T}^i))),\]

where \(\mathcal{R}^i\) is the environment matrix and \(\mathcal{T}^i\) is the type embedding.

The repinit block computes initial node and edge representations using attention-based message passing. The repformer block further refines these representations through multiple layers of graph convolution and attention mechanisms.

The final output dimension is:

\[\dim(\mathcal{D}^i) = \text{g1\_dim} + \text{tebd\_dim} \quad (\text{if concat\_output\_tebd}).\]

Parameters:

repinitUnion[RepinitArgs, dict]

The arguments used to initialize the repinit block, see docstr in RepinitArgs for details information.

repformerUnion[RepformerArgs, dict]

The arguments used to initialize the repformer block, see docstr in RepformerArgs for details information.

concat_output_tebdbool, optional

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

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.

trainablebool, optional

If the parameters are trainable.

seedint, optional

(Unused yet) Random seed for parameter initialization.

add_tebd_to_repinit_outbool, optional

Whether to add type embedding to the output representation from repinit before inputting it into repformer.

use_econf_tebdbool, Optional

Whether to use electronic configuration type embedding.

use_tebd_biasbool, Optional

Whether to use bias in the type embedding layer.

type_maplist[str], Optional

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

Returns:

descriptor: torch.Tensor

the descriptor of shape nf x nloc x g1_dim. invariant single-atom representation.

g2: torch.Tensor

invariant pair-atom representation.

h2: torch.Tensor

equivariant pair-atom representation.

rot_mat: torch.Tensor

rotation matrix for equivariant fittings

sw: torch.Tensor

The switch function for decaying inverse distance.

References

[1]

Zhang, D., Liu, X., Zhang, X. et al. DPA-2: a large atomic model as a multi-task learner. npj Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2

_update_sel_cls

share_params(base_class: DescrptDPA2, 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, repinit, repinit_three_body,

g1_shape_tranform, and repformers.

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 DPA2 descriptor.

Compression applies to the repinit block (DescrptBlockSeAtten). When attn_layer == 0, the geometric embedding is tabulated. Type embedding outputs are always precomputed.

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 for repinit 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, maps extended region index to local region.

comm_dict

MPI communication metadata for parallel inference. Forwarded to the repformer block (the message-passing part). The repinit sub-block does no message passing and does not receive it. None for non-parallel inference (default).

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. shape: nf x nloc x nnei x ng

h2

The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3

sw

The smooth switch function. shape: nf x nloc x nnei

_call_compressed(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None) → Any

Compressed forward for DPA2 descriptor.

The repinit forward is done inline with compressed ops, then the rest (g1_shape_transform, repformers, etc.) proceeds normally.

_compressed_repinit_forward(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, nframes: int, nloc: int, nall: int, type_embedding: torch.Tensor) → torch.Tensor

Compressed forward for the repinit block.

Same logic as DPA1’s _call_compressed but only produces the g1 output (nf x nloc x (ng x axis_neuron)), without rot_mat/sw returns.

Parameters:

coord_ext

Extended coordinates. shape: nf x (nall x 3)

atype_ext

Extended atom types. shape: nf x nall

nlist

Neighbor list for repinit. shape: nf x nloc x nnei_repinit

nframes

Number of frames.

nloc

Number of local atoms.

nall

Number of all atoms (local + ghost).

type_embedding

Full type embedding. shape: (ntypes+1) x tebd_dim

Returns:

torch.Tensor

Repinit output. shape: nf x nloc x (ng x axis_neuron)