deepmd.jax.descriptor.dpa3

Classes

DescrptDPA3

The DPA3 descriptor[Rfbfb621257be-1]_.

Module Contents

class deepmd.jax.descriptor.dpa3.DescrptDPA3(ntypes: int, repflow: 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, add_chg_spin_ebd: bool = False)

Bases: deepmd.dpmodel.descriptor.dpa3.DescrptDPA3

The DPA3 descriptor[Rfbfb621257be-1]_.

The DPA-3 descriptor uses a repflow block to iteratively update node, edge, and angle representations. The descriptor is computed as:

\[\mathcal{D}^i = \mathrm{RepFlow}(\mathcal{N}^i, \mathcal{E}^i, \mathcal{A}^i),\]

where \(\mathcal{N}^i\), \(\mathcal{E}^i\), and \(\mathcal{A}^i\) are the initial node, edge, and angle representations respectively.

The repflow block performs iterative updates through multiple layers:

\[\mathcal{N}^{i,l+1} = \mathrm{UpdateNode}(\mathcal{N}^{i,l}, \mathcal{E}^{i,l}, \mathcal{A}^{i,l}),\]

\[\mathcal{E}^{i,l+1} = \mathrm{UpdateEdge}(\mathcal{N}^{i,l}, \mathcal{E}^{i,l}, \mathcal{A}^{i,l}),\]

\[\mathcal{A}^{i,l+1} = \mathrm{UpdateAngle}(\mathcal{N}^{i,l}, \mathcal{E}^{i,l}, \mathcal{A}^{i,l}).\]

The final descriptor output dimension is:

\[\dim(\mathcal{D}^i) = \text{n\_dim} \times \text{axis\_neuron} \quad (\text{after symmetrization}).\]

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).

__setattr__(name: str, value: Any) → None