deepmd.pt.model.descriptor.dpa3

Classes

DescrptDPA3

The DPA3 descriptor[Rfbfb621257be-1]_.

Module Contents

class deepmd.pt.model.descriptor.dpa3.DescrptDPA3(ntypes: int, repflow: deepmd.dpmodel.descriptor.dpa3.RepFlowArgs | dict, concat_output_tebd: bool = False, activation_function: str = 'silu', precision: str = 'float64', exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, type_map: list[str] | None = None)

Bases: deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.Module

The DPA3 descriptor[Rfbfb621257be-1]_.

Parameters:

repflowUnion[RepFlowArgs, dict]

The arguments used to initialize the repflow block, see docstr in RepFlowArgs for details information.

concat_output_tebdbool, optional

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

activation_functionstr, optional

The activation function in the embedding net.

precisionstr, optional

The precision of the embedding net parameters.

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

Random seed for parameter initialization.

use_econf_tebdbool, Optional

Whether to use electronic configuration type embedding.

use_tebd_biasbool, Optional

Whether to use bias in the type embedding layer.

use_loc_mappingbool, Optional

Whether to use local atom index mapping in training or non-parallel inference. When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.

type_maplist[str], Optional

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

References

[1]

Zhang, D., Peng, A., Cai, C. et al. Graph neural network model for the era of large atomistic models. arXiv preprint arXiv:2506.01686 (2025).

repflow_args

activation_function = 'silu'

repflows

use_econf_tebd = False

use_loc_mapping = True

use_tebd_bias = False

type_map = None

tebd_dim = 128

type_embedding

concat_output_tebd = False

precision = 'float64'

prec

exclude_types = []

env_protection = 0.0

trainable = True

rcut

rcut_smth

sel

ntypes

compress = False

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.

share_params(base_class, shared_level, resume=False) → None

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=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

property dim_emb

Returns the embedding dimension g2.

compute_input_stats(merged: 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[torch.Tensor], stddev: list[torch.Tensor]) → None

Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[torch.Tensor], list[torch.Tensor]]

Get mean and stddev for descriptor.

serialize() → dict

Serialize the obj to dict.

classmethod deserialize(data: dict) → DescrptDPA3

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BD

The deserialized descriptor

forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None)

Compute the descriptor.

Parameters:

extended_coord

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

extended_atype

The extended aotm types. shape: nf x nall

nlist

The neighbor list. shape: nf x nloc x nnei

mapping

The index mapping, mapps extended region index to local region.

comm_dict

The data needed for communication for parallel inference.

Returns:

node_ebd

The output descriptor. shape: nf x nloc x n_dim (or n_dim + tebd_dim)

rot_mat

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

edge_ebd

The edge embedding. shape: nf x nloc x nnei x e_dim

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

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

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

abstractmethod 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

Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.

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