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