deepmd.tf2.descriptor

deepmd.tf2.descriptor#

Submodules#

Classes#

`DescrptDPA1`	Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.
`DescrptDPA2`	The DPA-2 descriptor[Rf0ed3afb961b-1]_.
`DescrptDPA3`	The DPA3 descriptor[Rfbfb621257be-1]_.
`DescrptDPA4`	SeZM descriptor.
`DescrptHybrid`	Concatenate a list of descriptors to form a new descriptor.
`DescrptSeAttenV2`	Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.
`DescrptSeA`	DeepPot-SE constructed from all information (both angular and radial) of
`DescrptSeR`	DeepPot-SE_R constructed from only the radial information of atomic configurations.
`DescrptSeT`	DeepPot-SE constructed from all information (both angular and radial) of atomic
`DescrptSeTTebd`	Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.

Package Contents#

class deepmd.tf2.descriptor.DescrptDPA1(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = None, trainable_ln: bool = True, ln_eps: float | None = 1e-05, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, spin: None = None, stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None, lmax: int = 1)[source]#

Bases: deepmd.dpmodel.descriptor.dpa1.DescrptDPA1

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

This descriptor, \(\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}\), is given by

\[\mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<,\]

where \(\hat{\mathcal{G}}^i\) represents the embedding matrix:math:mathcal{G}^i after additional self-attention mechanism and \(\mathcal{R}^i\) is defined by the full case in the se_e2_a descriptor. Note that we obtain \(\mathcal{G}^i\) using the type embedding method by default in this descriptor.

To perform the self-attention mechanism, the queries \(\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), keys \(\mathcal{K}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), and values \(\mathcal{V}^{i,l} \in \mathbb{R}^{N_c\times d_v}\) are first obtained:

\[\left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

\[\left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

\[\left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

where \(Q_{l}\), \(K_{l}\), \(V_{l}\) represent three trainable linear transformations that output the queries and keys of dimension \(d_k\) and values of dimension \(d_v\), and \(l\) is the index of the attention layer. The input embedding matrix to the attention layers, denoted by \(\mathcal{G}^{i,0}\), is chosen as the two-body embedding matrix.

Then the scaled dot-product attention method is adopted:

\[A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l},\]

where \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) \in \mathbb{R}^{N_c\times N_c}\) is attention weights. In the original attention method, one typically has \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}\right)=\mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)\), with \(\sqrt{d_{k}}\) being the normalization temperature. This is slightly modified to incorporate the angular information:

\[\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) = \mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right) \odot \hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T},\]

where \(\hat{\mathcal{R}}^{i} \in \mathbb{R}^{N_c\times 3}\) denotes normalized relative coordinates,: \(\hat{\mathcal{R}}^{i}_{j} = \frac{\boldsymbol{r}_{ij}}{\lVert \boldsymbol{r}_{ij} \lVert}\) and \(\odot\) means element-wise multiplication.
Then layer normalization is added in a residual way to finally obtain the self-attention local embedding matrix: \(\hat{\mathcal{G}}^{i} = \mathcal{G}^{i,L_a}\) after \(L_a\) attention layers:[^1]

\[\mathcal{G}^{i,l} = \mathcal{G}^{i,l-1} + \mathrm{LayerNorm}(A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})).\]

Parameters:

rcut: float: The cut-off radius \(r_c\)
rcut_smth: float: From where the environment matrix should be smoothed \(r_s\)
sellist[int], int: list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius
ntypesint: Number of element types
neuronlist[int]: Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)
axis_neuron: int: Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)
tebd_dim: int: Dimension of the type embedding
tebd_input_mode: str: The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the radial embedding network output.
resnet_dt: bool: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: bool: If the weights of this descriptors are trainable.
trainable_ln: bool: Whether to use trainable shift and scale weights in layer normalization.
ln_eps: float, Optional: The epsilon value for layer normalization.
type_one_side: bool: If ‘False’, type embeddings of both neighbor and central atoms are considered. If ‘True’, only type embeddings of neighbor atoms are considered. Default is ‘False’.
attn: int: Hidden dimension of the attention vectors
attn_layer: int: Number of attention layers
attn_dotr: bool: If dot the angular gate to the attention weights
attn_mask: bool: (Only support False to keep consistent with other backend references.) (Not used in this version. True option is not implemented.) If mask the diagonal of attention weights
exclude_typeslist[list[int]]: 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_protection: float: Protection parameter to prevent division by zero errors during environment matrix calculations.
set_davg_zero: bool: Set the shift of embedding net input to zero.
activation_function: str: The activation function in the embedding net. Supported options are “linear”, “gelu_erf”, “relu6”, “silut”, “tanh”, “sigmoid”, “gelu”, “softplus”, “relu”, “none”, “gelu_tf”, “silu”.
precision: str: The precision of the embedding net parameters. Supported options are “default”, “float64”, “float32”, “bfloat16”, “float16”.
scaling_factor: float: The scaling factor of normalization in calculations of attention weights. If temperature is None, the scaling of attention weights is (N_dim * scaling_factor)**0.5
normalize: bool: Whether to normalize the hidden vectors in attention weights calculation.
temperature: float: If not None, the scaling of attention weights is temperature itself.
smooth_type_embedding: bool: Whether to use smooth process in attention weights calculation.
concat_output_tebd: bool: Whether to concat type embedding at the output of the descriptor.
stripped_type_embedding: bool, Optional: (Deprecated, kept only for compatibility.) Whether to strip the type embedding into a separate embedding network. Setting this parameter to True is equivalent to setting tebd_input_mode to ‘strip’. Setting it to False is equivalent to setting tebd_input_mode to ‘concat’. The default value is None, which means the tebd_input_mode setting will be used instead.
use_econf_tebd: bool, Optional: Whether to use electronic configuration type embedding.
use_tebd_biasbool, Optional: Whether to use bias in the type embedding layer.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.
lmax: int: Maximum angular degree of the Cartesian moment basis. Supported values are 1 through 4.
spin: (Only support None to keep consistent with other backend references.) (Not used in this version. Not-none option is not implemented.) The old implementation of deepspin.

References

[1]

Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. 2022. DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. arXiv preprint arXiv:2208.08236.

class deepmd.tf2.descriptor.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)[source]#

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

class deepmd.tf2.descriptor.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, default_chg_spin: list[float] | None = None)[source]#

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

class deepmd.tf2.descriptor.DescrptDPA4(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.descriptor.dpa4.DescrptDPA4

SeZM descriptor.

DPA4 stores the state of atom \(i\) at layer \(l\) as SO(3) coefficients \(\mathbf h_i^{(l,\ell,m)}\). For an edge \(j\to i\), the source state is rotated into an edge-aligned frame, processed by an SO(2)-equivariant convolution, and rotated back:

\[\mathbf q_{ji}^{(l)} = \mathbf D(\hat{\mathbf r}_{ji})^{-1}\mathbf h_j^{(l)}, \qquad \mathbf m_{ji}^{(l)} = \mathbf D(\hat{\mathbf r}_{ji}) \operatorname{SO2Conv}\!\left( \mathbf q_{ji}^{(l)},\boldsymbol\rho(r_{ji})\right),\]

where \(\mathbf D\) contains Wigner-D rotation blocks and \(\boldsymbol\rho\) is the radial embedding multiplied by a smooth cutoff envelope. In the baseline residual path, the aggregated message is first added directly to the node state, after which every equivariant FFN subblock applies its own residual update:

\[\mathbf M_i^{(l)}=\sum_{j\in\mathcal N(i)} w_{ji}\mathbf m_{ji}^{(l)},\qquad \mathbf u_i^{(l,0)}=\mathbf h_i^{(l)}+\mathbf M_i^{(l)},\]

\[\mathbf u_i^{(l,r)}=\mathbf u_i^{(l,r-1)}+ \operatorname{FFN}_{\mathrm{eq},r} \!\left(\mathbf u_i^{(l,r-1)}\right),\qquad \mathbf h_i^{(l+1)}=\mathbf u_i^{(l,B)}.\]

Consequently, one FFN subblock gives \(\mathbf h_i^{(l+1)}=\mathbf h_i^{(l)}+\mathbf M_i^{(l)}+ \operatorname{FFN}_{\mathrm{eq}}(\mathbf h_i^{(l)}+\mathbf M_i^{(l)})\). The AttnRes modes replace these baseline shortcuts with selective depth-wise aggregation before the SO(2) and/or FFN units.

The final read-out applies the configured scalar/equivariant read-out to the last interaction state and then keeps its invariant scalar output:

\[\mathcal D_i = \operatorname{ScalarReadout}_{\mathrm{mode}} \left(\mathbf h_i^{(L)}\right).\]

In so3_readout="none" mode, coefficients with \(\ell>0\) are discarded and the \(\ell=0\) slice is processed by the configured learned scalar residual FFN stack. In "glu" and "mlp" modes, equivariant residual read-out blocks first fold higher-degree coefficients into the scalar channel before extraction.

The weights \(w_{ji}\) are either cutoff-envelope weights or normalized attention weights, depending on n_atten_head.

Parameters:

ntypes

Number of element types.

sel

Maximum number of neighbors per type within rcut. - int: broadcast to all types, e.g. sel=100 with ntypes=2 → [100, 100] - list[int]: sel[i] is the maximum number of type i atoms within rcut

rcut

Cutoff radius in Å.

env_exp

C^3 cutoff envelope exponents. A list [rbf_env_exp, edge_env_exp] specifies the radial-basis and message-passing envelopes separately. A zero radial-basis exponent disables that envelope. An integer specifies only the message-passing envelope exponent and disables the radial-basis envelope. Larger values give weaker suppression (values stay near 1.0 longer).

channels

Total channels per (l,m) coefficient.

basis_type

Radial basis type. Supported values are "bessel", "gaussian", "bessel/fix" and "gaussian/fix"; the /fix forms keep the frequencies or centres fixed during training.

n_radial

Number of radial basis functions.

radial_mlp

Hidden layer sizes for radial networks. An output layer of size (l_schedule[0]+extra_node_l+1)*channels will be automatically appended.

edge_norm

Channel RMSNorm on the descriptor’s cutoff-vanishing branches: the radial network hidden layers, the environment-seed FiLM scale/shift logits, and the cross-focus competition scalars. A bool switches all three together; a list of three bools [radial, film, focus] switches them individually. Disabled norms are identity pass-throughs. The post-SO(2) residual scaling follows the radial entry: with it disabled the norm uses the unit floor sqrt(1 + variance), so small messages retain their cutoff envelope instead of receiving the standard 1/sqrt(eps) small-signal gain.

use_env_seed

If True, seed the initial node state with local-environment information: apply environment matrix FiLM conditioning on l=0 features using 4D [s, s*r_hat] representation, and enable the non-scalar geometric initial embedding when l_schedule[0] + extra_node_l > 0. If False, the initial state contains only atom-local scalar features before message passing. FiLM deltas are normalized and scaled with learnable strengths initialized to small values. Internal dimensions are derived from channels: embed_dim=min(channels, 128), axis_dim=min(4 if embed_dim < 64 else 8, embed_dim-1), type_dim=clamp(channels//4, 8, 32), rbf_out_dim=max(32, embed_dim-2*type_dim), hidden_dim=min(256, max(2*embed_dim, rbf_out_dim+2*type_dim)).

random_gamma

If True, apply a random roll about the edge-aligned local +Z axis before building the Wigner-D blocks. The roll is sampled independently per edge and per forward call.

edge_cartesian

If True, every block whose message-passing degree is 1 or 2 replaces its per-edge SO(2) rotation-frame tensor product with the equivalent global-frame Cartesian rank-2 tensor product, removing the two per-edge Wigner-D rotations. Blocks with degree 0 or >= 3 keep the SO(2) path. When every block takes the Cartesian path the full Wigner-D construction is skipped automatically, and the geometric initial embedding falls back to the zonal coupling.

node_cartesian

Per-node global-frame Cartesian rank-2 tensor product on the aggregated message, applied in every block whose message-passing degree is 1 or 2. Configured by a "<mode>:<layers>" string where mode is "default" (one-sided product) or "parity" (symmetrized product) and layers is the stack depth; a bare integer N is shorthand for "default:N", and "none" disables it. Orthogonal to edge_cartesian: either, both, or neither may be set. Unlike edge_cartesian it does not affect the Wigner-D construction, since the per-edge message path is left unchanged.

lmax

Maximum degree, only used when l_schedule is None.

l_schedule

Pyramid schedule of lmax per block, e.g. [3, 3, 2]. Must be non-increasing. If set, lmax and n_blocks will be ignored.

mmax

Maximum SO(2) order (|m|), only used when m_schedule is None. If None, defaults to the per-block lmax (i.e. m_schedule = l_schedule).

kmax

Maximum Wigner-D frame order (|k|) used by SO(3) grid nets. The frame set is built as [0, -1, 1, ..., -kmax, kmax]. kmax=0 recovers the S2-like k=0 slice, while kmax=1 is the default low-cost setting that opens odd/antisymmetric coupling paths.

m_schedule

Schedule of mmax per block, e.g. [2, 2, 1, 0]. Must satisfy m_schedule[i] <= l_schedule[i] for every block. A non-increasing schedule is recommended but not required. If set, mmax will be ignored.

extra_node_l

Extra node representation degree above each message-passing degree. The node degree of block i is l_schedule[i] + extra_node_l, while SO(2) message passing still uses l_schedule[i].

n_blocks

Number of blocks (only used when l_schedule is None). 0 disables the interaction blocks and builds the zero-block descriptor: type embedding, optional env FiLM and geometric initial embedding, then the final SO(3) read-out. The backbone degree is taken from lmax (plus extra_node_l). Geometry then enters only through the GIE, which is active when use_env_seed=True and lmax + extra_node_l > 0.

so2_norm

If True, apply intermediate ReducedEquivariantRMSNorm between SO(2) mixing layers. When False (default), no normalization is applied between layers.

mixing_layers

Number of learnable mixing layers in the per-edge message core of each block (legacy alias: so2_layers). 0 applies only the edge-condition modulation: the rotation-free per-degree radial scaling on the SO(2) path, or a single x @ T_e when edge_cartesian applies. The per-node node_cartesian stack carries its own independent depth.

so2_attn_res

SO(2)-internal depth-wise attention residual mode inside each interaction block. Must be one of "none", "independent", or "dependent".

radial_so2_mode

Dynamic radial degree mixer mode inside SO(2) convolution. "none" applies elementwise radial modulation, "degree" uses a channel-shared edge-conditioned cross-degree kernel, and "degree_channel" uses a per-channel cross-degree kernel. Has no effect on blocks taking the Cartesian path (edge_cartesian with degree 1 or 2), where the dynamic radial degree mixer is bypassed.

radial_so2_rank

Low-rank channel factorization rank for radial_so2_mode="degree_channel". 0 uses the full per-channel dynamic degree kernel.

n_focus

Number of parallel focus streams used only inside the SO(2) convolution. Node-level backbone tensors still keep a singleton focus axis.

focus_dim

Hidden width per focus stream inside the SO(2) convolution. focus_dim=0 means using channels.

n_atten_head

Number of attention heads when aggregating messages in SO(2) convolution. 0 applies a plain envelope-weighted scatter-sum; >0 enables envelope-gated grouped softmax attention with output-side head gate. Attention uses w**2 * exp(logit) in the numerator and zeta + sum(w**2 * exp(logit)) in the denominator.

atten_f_mix

If True, merge all SO(2) focus streams into one attention stream after rotate-back. Attention heads split n_focus * focus_dim instead of each focus stream independently.

atten_v_proj

If True, apply an explicit degree-aware value projection inside SO(2) attention.

atten_o_proj

If True, apply an explicit degree-aware output projection inside SO(2) attention.

ffn_neurons

Hidden width for block FFNs and the final scalar output FFN. If >0, both paths use this width. If =0, each path resolves its own width from channels and its effective GLU setting: 4 * channels without GLU, (8 / 3) * channels with GLU, then round up to a multiple of 32.

grid_mlp

Either one boolean applied to every grid path, or three booleans [node_wise, message_node, ffn] selecting the polynomial point-wise grid MLP operation per grid path. On any path whose grid_branch entry is positive it is overridden by branch mixing, and it has no effect on the final l=0 output head.

grid_branch

Either one non-negative integer applied to every grid path, or three integers [node_wise, message_node, ffn] setting the number of scalar-routed polynomial product branches per grid path. 0 disables branch mixing on that path; positive values select branch mixing and take precedence over grid_mlp. Branch weights are computed from l=0 scalar features only, while each branch is a quadratic product of channel-mixed grid fields. The node_wise and message_node entries control the SO(2) convolution cross-grid paths, and the ffn entry controls the block-internal FFN grid path.

ffn_blocks

Number of FFN subblocks per interaction block.

sandwich_norm

Pre/post-norm switches for [SO(2), FFN] residual branches in order: [so2_pre, so2_post, ffn_pre, ffn_post], shared across all blocks.

mlp_bias

Whether to use bias in equivariant layers. When False, removes bias from: - SO3Linear: l=0 bias - SO2Linear: l=0 bias - GatedActivation: gate linear bias - DepthAttnRes: input-dependent query projection - EnvironmentInitialEmbedding:

rbf_proj_layer1/2 and g_layer1/2

Attention logit and output-gate parameters in SO(2) convolution are always bias-free.

layer_scale

If True, apply learnable LayerScale (init 1e-3) on residual branches: - SO(2) branch: per-focus-channel scales (n_focus, focus_dim)

on each SO(2) mixing layer.

FFN branch: per-channel scales (channels,) on each FFN subblock.

full_attn_res

Descriptor-level full attention residual mode over the unit history [x0, so2_0, ffn_0_0, ffn_0_1, …, so2_1, ffn_1_0, ffn_1_1, …], where each FFN subblock contributes its own completed unit representation. independent uses learned query vectors, while dependent derives queries from the current SeZM state before the SO(2) unit, before each FFN unit, and before the final aggregation. Must be one of "none", "independent", or "dependent".

block_attn_res

Descriptor-level block attention residual mode over the block history [x0, b1, b2, …], where each b_i is the sum of all unit outputs inside one SeZMInteractionBlock. independent uses learned query vectors, while dependent derives queries from the current SeZM state before the SO(2) unit, before each FFN unit, and before the final block aggregation. Must be one of "none", "independent", or "dependent". Cannot be enabled together with full_attn_res.

s2_activation

Two booleans [so2_enabled, ffn_enabled]. so2_enabled=True makes the SO(2) gated activation path use activation_function="silu". ffn_enabled=True makes the block-internal FFN path use activation_function="silu" and glu_activation=True. S2-grid resolutions are resolved automatically per block. The tensor-product grid uses [2 * mmax + 4, ceil_even(3 * lmax + 2)] in the SO(2) branch, and the FFN branch lifts it to a square [max(R_phi, R_theta), max(R_phi, R_theta)] grid. Lebedev branches use the smallest packaged rule with precision at least 3 * lmax. The final l=0 output FFN is unchanged.

ffn_so3_grid

If True, use the SO(3) Wigner-D grid in the block-internal FFN. This option takes precedence over the FFN grid path and ignores s2_activation[1]. The final l=0 output FFN is unchanged.

node_wise_s2

If True, add an edge-local S2 product branch between source and destination node features inside the SO(2) convolution.

node_wise_so3

If True, use the corresponding edge-local SO(3) Wigner-D grid-net branch. The source side is the query and the destination side is the context.

message_node_s2

If True, add a post-aggregation S2 product branch between hidden messages and destination node features before the SO(2) output projection.

message_node_so3

If True, use the corresponding post-aggregation SO(3) Wigner-D grid-net branch. The message is the query and the node state is the context.

so3_readout

Read-out FFN mode for the final l=0 descriptor. "none" applies a degree-0 scalar FFN to the l=0 slice only; l>0 coefficients are discarded before the read-out. "glu" and "mlp" apply a full equivariant FFN whose degree equals the node degree of the last interaction block, driven by the SO(3) Wigner-D grid, so l>0 geometry is folded into l=0 before the scalar is extracted. The value selects the quadratic grid product ("glu") or the polynomial point-wise grid MLP ("mlp"). The Wigner-D frame order follows kmax.

readout_layers

Number of stacked equivariant residual read-out FFNs (default 1). Every layer is an x + FFN(x) residual block sharing the read-out degree; intermediate layers keep the full SO(3) tensor so high-degree geometry is folded into l=0 repeatedly, and only the final layer slices the l=0 channel from its residual sum. With so3_readout of "none" the stack is a degree-0 scalar residual MLP on the l=0 slice.

lebedev_quadrature

Either one boolean applied to both S2 branches, or two booleans [so2_enabled, ffn_enabled] aligned with s2_activation. If enabled for a branch, that branch uses packaged Lebedev quadrature instead of the tensor-product sphere grid in its S2 projector.

activation_function

Base activation function for helper MLPs, the SO(2) gated activation path, and the final l=0 output FFN. It is overridden to "silu" only on paths whose s2_activation switch is enabled.

glu_activation

Base GLU switch for FFN. The block-internal FFN path overrides it to True only when s2_activation[1]=True. The final l=0 output FFN always keeps this user-provided value.

use_amp

If True, use automatic mixed precision (AMP) with bfloat16 on CUDA during training. This can improve speed and reduce memory usage. Enabling this option is recommended on GPUs with native bfloat16 support. Disable it on GPUs without native bfloat16 support to avoid runtime errors or additional conversion overhead.

exclude_types

List of excluded type pairs.

precision

Precision for neural network parameters and computations. Geometry computations (edge distances, Wigner-D matrices, rotations, and enabled env seeds) always run in fp32+ to provide accurate geometric information for better convergence. Only the interaction blocks use this precision.

eps

Small epsilon for numerical stability in division and normalization.

trainable

Whether parameters are trainable.

seed

Random seed(s).

type_map

Type names.

inner_clamp_f_inner

Inner radius of the zone-bridging window, as a fraction of each pair’s own length scale. When both fractions are set, the descriptor freezes short-range geometry inside the window. Given together with inner_clamp_f_outer or not at all.

inner_clamp_f_outer

Outer radius of the zone-bridging window, as a fraction of each pair’s own length scale. At or beyond it an edge contributes exactly as it does without a window.

inner_clamp_scale

Length scale the two fractions measure against. "covalent" sizes every pair by the sum of the covalent radii of its two elements, so one window shape serves every element combination and a type_map is required; "absolute" gives every element the unit radius, which makes the pair scale 1 Å and the two fractions radii in Å.

add_chg_spin_ebd

If True, add frame-level charge/spin condition embedding to scalar type features before edge features are built.

default_chg_spin

Default frame-level charge/spin condition [charge, spin]. This value is used when add_chg_spin_ebd=True and no explicit charge_spin tensor is provided at the descriptor or SeZM model boundary.

Notes

SeZM does not use the traditional environment matrix (r, a_x, a_y, a_z). Instead, it uses radial basis functions and spherical harmonics directly. The mean/stddev statistics are kept for interface compatibility but are not actively used in the forward pass.

_tf2_training_mode = False[source]#

classmethod deserialize(data: dict) → DescrptDPA4[source]#

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

BD: The deserialized descriptor

set_training_mode(training: bool) → None[source]#: Select train/eval graph tracing for training-only augmentation.

_in_training_mode() → bool[source]#: Return the Python trace-time training state used by TF2 graphs.

class deepmd.tf2.descriptor.DescrptHybrid(list: DescrptHybrid.__init__.list[deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor | dict[str, Any]], type_map: DescrptHybrid.__init__.list[str] | None = None, ntypes: int | None = None)[source]#

Bases: deepmd.dpmodel.descriptor.hybrid.DescrptHybrid

Concatenate a list of descriptors to form a new descriptor.

The hybrid descriptor combines multiple descriptors by concatenation:

\[\mathcal{D}^i = [\mathcal{D}^i_1, \mathcal{D}^i_2, ..., \mathcal{D}^i_n],\]

where \(\mathcal{D}^i_k\) is the descriptor computed by the \(k\)-th sub-descriptor for atom \(i\).

The output dimension is the sum of all sub-descriptor dimensions:

\[\dim(\mathcal{D}^i) = \sum_{k=1}^{n} \dim(\mathcal{D}^i_k).\]

Parameters:

listlistlist[Union[BaseDescriptor, dict[str, Any]]]: Build a descriptor from the concatenation of the list of descriptors. The descriptor can be either an object or a dictionary.

class deepmd.tf2.descriptor.DescrptSeAttenV2(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = DEFAULT_PRECISION, scaling_factor: float = 1.0, normalize: bool = True, temperature: float | None = None, trainable_ln: bool = True, ln_eps: float | None = 1e-05, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, spin: None = None, stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None, lmax: int = 1)[source]#

Bases: deepmd.tf2.descriptor.dpa1.DescrptDPA1, deepmd.dpmodel.descriptor.se_atten_v2.DescrptSeAttenV2