deepmd.dpmodel.descriptor.dpa2

Classes

RepinitArgs
RepformerArgs
DescrptDPA2	The DPA-2 descriptor[Rf0ed3afb961b-1]_.

Module Contents

class deepmd.dpmodel.descriptor.dpa2.RepinitArgs(rcut: float, rcut_smth: float, nsel: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, tebd_input_mode: str = 'concat', set_davg_zero: bool = True, activation_function: str = 'tanh', resnet_dt: bool = False, type_one_side: bool = False, use_three_body: bool = False, three_body_neuron: list[int] = [2, 4, 8], three_body_sel: int = 40, three_body_rcut: float = 4.0, three_body_rcut_smth: float = 0.5)

rcut

rcut_smth

nsel

neuron = [25, 50, 100]

axis_neuron = 16

tebd_dim = 8

tebd_input_mode = 'concat'

set_davg_zero = True

activation_function = 'tanh'

resnet_dt = False

type_one_side = False

use_three_body = False

three_body_neuron = [2, 4, 8]

three_body_sel = 40

three_body_rcut = 4.0

three_body_rcut_smth = 0.5

__getitem__(key: str) → Any

serialize() → dict

classmethod deserialize(data: dict) → RepinitArgs

class deepmd.dpmodel.descriptor.dpa2.RepformerArgs(rcut: float, rcut_smth: float, nsel: 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, 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)

rcut

rcut_smth

nsel

nlayers = 3

g1_dim = 128

g2_dim = 16

axis_neuron = 4

direct_dist = False

update_g1_has_conv = True

update_g1_has_drrd = True

update_g1_has_grrg = True

update_g1_has_attn = True

update_g2_has_g1g1 = True

update_g2_has_attn = True

update_h2 = False

attn1_hidden = 64

attn1_nhead = 4

attn2_hidden = 16

attn2_nhead = 4

attn2_has_gate = False

activation_function = 'tanh'

update_style = 'res_avg'

update_residual = 0.001

update_residual_init = 'norm'

set_davg_zero = True

trainable_ln = True

use_sqrt_nnei = True

g1_out_conv = True

g1_out_mlp = True

ln_eps = 1e-05

__getitem__(key: str) → Any

serialize() → dict

classmethod deserialize(data: dict) → RepformerArgs

class deepmd.dpmodel.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.NativeOP, deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor

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

repinit_args

repformer_args

repinit

use_three_body = False

repformers

rcsl_list

rcut_list

nsel_list

use_econf_tebd = False

use_tebd_bias = False

type_map = None

type_embedding

concat_output_tebd = True

precision = 'float64'

smooth = True

exclude_types = []

env_protection = 0.0

rcut_smth

trainable = True

add_tebd_to_repinit_out = False

repinit_out_dim

tebd_transform = None

tebd_dim = 8

rcut

ntypes

sel

get_rcut() → float

Returns the cut-off radius.

get_rcut_smth() → float

Returns the radius where the neighbor information starts to smoothly decay to 0.

get_nsel() → int

Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int]

Returns the number of selected atoms for each type.

get_ntypes() → int

Returns the number of element types.

get_type_map() → list[str]

Get the name to each type of atoms.

get_dim_out() → int

Returns the output dimension of this descriptor.

get_dim_emb() → int

Returns the embedding dimension of this descriptor.

mixed_types() → bool

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.

has_message_passing() → bool

Returns whether the descriptor has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the descriptor needs sorted nlist when using forward_lower.

get_env_protection() → float

Returns the protection of building environment matrix.

abstractmethod share_params(base_class: Any, shared_level: int, resume: bool = False) → NoReturn

Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.

change_type_map(type_map: list[str], model_with_new_type_stat: Any = None) → None

Change the type related params to new ones, according to type_map and the original one in the model. If there are new types in type_map, statistics will be updated accordingly to model_with_new_type_stat for these new types.

property dim_out: int

property dim_emb: int

Returns the embedding dimension g2.

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

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: torch.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.

set_stat_mean_and_stddev(mean: list[deepmd.dpmodel.array_api.Array], stddev: list[deepmd.dpmodel.array_api.Array]) → None

Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[deepmd.dpmodel.array_api.Array], list[deepmd.dpmodel.array_api.Array]]

Get mean and stddev for descriptor.

call(coord_ext: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | 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]

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.

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

serialize() → dict

Serialize the obj to dict.

classmethod deserialize(data: dict) → DescrptDPA2

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BD

The deserialized descriptor

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array]

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms