docs(dpmodel): document mathematical operations

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -53,6 +53,15 @@
 
 
 class BaseAtomicModel(BaseAtomicModel_, NativeOP):
+    """Base interface mapping local atomic environments to per-atom outputs.
+
+    The local environment includes the selected neighbor indices and the
+    corresponding coordinates and atom types, together with optional frame,
+    atomic, or descriptor-specific conditioning inputs.  Concrete subclasses
+    may learn a descriptor-plus-fitting map, interpolate a fixed pair table, or
+    combine outputs from existing atomic models.
+    """
+
     def __init__(
         self,
         type_map: list[str],

deepmd/dpmodel/atomic_model/dipole_atomic_model.py

@@ -22,6 +22,22 @@
 
 
 class DPDipoleAtomicModel(DPAtomicModel):
+    r"""Atomic dipole model reconstructed from descriptor rotation matrices.
+
+    The fitting network predicts local coefficients and contracts them with
+    the equivariant descriptor output:
+
+    .. math::
+
+       \mathbf M_i=F_\theta(\mathcal D_i),\qquad
+       \boldsymbol\mu_i=\mathbf M_i\mathbf R_i,
+
+    where :math:`\mathbf R_i\in\mathbb R^{m_1\times3}` is the descriptor
+    rotation matrix.  This contraction produces the lab-frame vector.
+
+    Frame dipoles are additive: :math:`\boldsymbol\mu=\sum_i\boldsymbol\mu_i`.
+    """
+
     def __init__(
         self,
         descriptor: BaseDescriptor,

deepmd/dpmodel/atomic_model/dos_atomic_model.py

@@ -19,6 +19,11 @@
 
 
 class DPDOSAtomicModel(DPAtomicModel):
+    r"""Atomic DOS model predicting :math:`D_{ik}=F_{\theta,k}(\mathcal D_i)`.
+
+    The global DOS is :math:`D_k=\sum_iD_{ik}`.
+    """
+
     def __init__(
         self,
         descriptor: BaseDescriptor,

deepmd/dpmodel/atomic_model/energy_atomic_model.py

@@ -14,6 +14,11 @@
 
 
 class DPEnergyAtomicModel(DPAtomicModel):
+    r"""Atomic energy model with :math:`E_i=F_\theta(\mathcal D_i)`.
+
+    The frame energy is :math:`E=\sum_iE_i`.
+    """
+
     def __init__(
         self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any
     ) -> None:

deepmd/dpmodel/atomic_model/polar_atomic_model.py

@@ -18,6 +18,31 @@
 
 
 class DPPolarAtomicModel(DPAtomicModel):
+    r"""Atomic polarizability model reconstructed in the laboratory frame.
+
+    Let :math:`\mathbf R_i\in\mathbb R^{m_1\times3}` be the descriptor
+    rotation matrix.  In diagonal fitting mode the network predicts
+    :math:`\mathbf p_i=F_\theta(\mathcal D_i)` and reconstructs
+
+    .. math::
+
+       \boldsymbol\alpha_i=\mathbf R_i^T
+       \operatorname{diag}(\mathbf p_i)\mathbf R_i.
+
+    In full-matrix mode it predicts :math:`\widehat{\mathbf P}_i`, symmetrizes
+    :math:`\mathbf P_i=(\widehat{\mathbf P}_i+
+    \widehat{\mathbf P}_i^T)/2`, and reconstructs
+
+    .. math::
+
+       \boldsymbol\alpha_i=\mathbf R_i^T\mathbf P_i\mathbf R_i.
+
+    Type-dependent scaling is applied to the predicted local coefficients, and
+    an optional isotropic shift :math:`c_{t_i}\mathbf I` is added after the
+    reconstruction.  The frame tensor is additive:
+    :math:`\boldsymbol\alpha=\sum_i\boldsymbol\alpha_i`.
+    """
+
     def __init__(
         self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any
     ) -> None:

deepmd/dpmodel/atomic_model/property_atomic_model.py

@@ -16,6 +16,11 @@
 
 
 class DPPropertyAtomicModel(DPAtomicModel):
+    r"""Generic atomic property map :math:`p_i=F_\theta(\mathcal D_i)`.
+
+    Extensive frame properties use :math:`p=\sum_i p_i`.
+    """
+
     def __init__(
         self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any
     ) -> None:

deepmd/dpmodel/descriptor/dpa1.py

@@ -2153,6 +2153,8 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
 
 
 class NeighborGatedAttention(NativeOP):
+    r"""Gated neighbor aggregation :math:`h_i'=h_i+\sum_j a_{ij}v_{ij}`."""
+
     def __init__(
         self,
         layer_num: int,
@@ -2286,6 +2288,21 @@ def deserialize(cls, data: dict) -> "NeighborGatedAttention":
 
 
 class NeighborGatedAttentionLayer(NativeOP):
+    r"""Single gated neighbor-attention residual layer.
+
+    For neighbor features :math:`\mathbf X`, the layer applies gated attention,
+    adds a residual connection, and normalizes the result:
+
+    .. math::
+        \mathbf X' = \operatorname{LayerNorm}\!\left(
+        \mathbf X + \operatorname{GatedAttention}
+        (\mathbf X, \mathbf M, \mathbf R, \mathbf S)\right),
+
+    where :math:`\mathbf M` is the neighbor mask and the optional
+    :math:`\mathbf R` and :math:`\mathbf S` supply directional and switching
+    information.
+    """
+
     def __init__(
         self,
         nnei: int,
@@ -2345,6 +2362,12 @@ def call(
         input_r: Array | None = None,
         sw: Array | None = None,
     ) -> Array:
+        r"""Apply attention, its residual connection, and layer normalization.
+
+        .. math::
+            H_{\mathrm{out}}=\operatorname{LayerNorm}
+            \left(H+\operatorname{GatedAttention}(H,M,R,S)\right).
+        """
         residual = x
         x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw)
         x = residual + x
@@ -2394,6 +2417,35 @@ def deserialize(cls, data: dict) -> "NeighborGatedAttentionLayer":
 
 
 class GatedAttentionLayer(NativeOP):
+    r"""Projected gated self-attention output.
+
+    With projected queries, keys, and values, the layer returns only the
+    attention output (the residual connection is applied by
+    :class:`NeighborGatedAttentionLayer`):
+
+    .. math::
+        Q,K,V=\operatorname{split}(H W_{\mathrm{in}}),\qquad
+        L=\alpha\,\widetilde Q\widetilde K^T,\qquad
+        S_{ij}=s_i s_j,
+
+    .. math::
+
+        \overline L_{ij}=(L_{ij}+c)S_{ij}-c,\qquad
+        \overline A_{ij}=\operatorname{softmax}_{j}(\overline L_{ij}),
+        \qquad A_{ij}=S_{ij}\overline A_{ij},
+
+        O=\operatorname{reshape}((A\odot R)V)W_{\mathrm{out}}.
+
+    Here the tildes denote optional per-vector normalization of :math:`Q`,
+    :math:`K`, and :math:`V`.  The implementation uses
+    :math:`\alpha=(d\,s)^{-1/2}` for ``scaling_factor`` :math:`s`, or the
+    configured ``temperature`` value when it is provided.  Neighbor masks and
+    cutoff smoothing modifies both the logits before softmax and, through
+    :math:`S`, the attention amplitude afterward.  Without smoothing, invalid
+    keys are masked before softmax and invalid query rows are zeroed.  The
+    optional angular matrix :math:`R` is applied after these operations.
+    """
+
     def __init__(
         self,
         nnei: int,

deepmd/dpmodel/descriptor/dpa2.py

@@ -75,6 +75,8 @@
 
 
 class RepinitArgs:
+    r"""Representation initialization parameters for :math:`\mathcal G`."""
+
     def __init__(
         self,
         rcut: float,
@@ -185,6 +187,8 @@ def deserialize(cls, data: dict) -> "RepinitArgs":
 
 
 class RepformerArgs:
+    r"""Representation update parameters for :math:`\mathcal G^{l+1}=\Phi_l(\mathcal G^l)`."""
+
     def __init__(
         self,
         rcut: float,

deepmd/dpmodel/descriptor/dpa3.py

@@ -64,21 +64,54 @@ class RepFlowArgs:
     The DPA-3 descriptor uses a repflow architecture that maintains and updates three types
     of representations: node (:math:`\mathbf{n}`), edge (:math:`\mathbf{e}`), and angle (:math:`\mathbf{a}`).
 
-    The update equations for each layer are:
+    DPA3 applies message passing to the first two graphs of the line-graph
+    series.  Writing node, edge, and angle features as
+    :math:`\mathbf n_i^l`, :math:`\mathbf e_{ij}^l`, and
+    :math:`\mathbf a_{ij,ik}^l`, respectively, the default parallel layer with
+    angle updates forms
 
     .. math::
-        \mathbf{n}^{l+1} = \text{UpdateNode}(\mathbf{n}^l, \mathbf{e}^l, \mathbf{a}^l),
+
+        \mathbf m_{ij}^{E,\mathrm{self}}=
+        U_E(\mathbf e_{ij}^l,\mathbf n_i^l,\mathbf n_j^l),
+        \qquad
+        \mathbf m_{ij}^{A\to E}=\operatorname{Reduce}_k
+        U_{A\to E}(\mathbf a_{ij,ik}^l,\mathbf n_i^l,
+        \mathbf e_{ij}^l,\mathbf e_{ik}^l),
 
     .. math::
-        \mathbf{e}^{l+1} = \text{UpdateEdge}(\mathbf{n}^l, \mathbf{e}^l, \mathbf{a}^l),
+
+        \mathbf e_{ij}^{l+1}=\operatorname{Combine}_E
+        (\mathbf e_{ij}^l,\mathbf m_{ij}^{E,\mathrm{self}},
+        \mathbf m_{ij}^{A\to E}),
 
     .. math::
-        \mathbf{a}^{l+1} = \text{UpdateAngle}(\mathbf{n}^l, \mathbf{e}^l, \mathbf{a}^l).
 
-    The final descriptor is obtained by symmetrization:
+        \mathbf a_{ij,ik}^{l+1}=\operatorname{Combine}_A\!\left(
+        \mathbf a_{ij,ik}^l,
+        U_A(\mathbf a_{ij,ik}^l,\mathbf n_i^l,
+        \mathbf e_{ij}^l,\mathbf e_{ik}^l)\right),
+
+    .. math::
+
+        \mathbf n_i^{l+1}=\operatorname{Combine}_N\!\left(
+        \mathbf n_i^l,U_N^{\mathrm{self}}(\mathbf n_i^l),
+        U_N^{\mathrm{sym}}(\{\mathbf e_{ij}^l,\mathbf n_j^l\}_j),
+        \operatorname{Reduce}_j U_{E\to N}
+        (\mathbf n_i^l,\mathbf n_j^l,\mathbf e_{ij}^l)\right).
+
+    The ``Combine`` operation is selected by ``update_style``; its default
+    ``res_residual`` form adds each message with a learned residual weight.
+    The angle-to-edge reduction is switch-weighted over :math:`k`.  In
+    sequential mode the same dependencies are evaluated with the most recently
+    updated edge and angle features.
+
+    Here the vertices of the second line graph are the edges of the first, so
+    :math:`\mathbf v_{ij}^{(2,l)} \equiv \mathbf e_{ij}^l`.  The invariant
+    atomic descriptor is the final first-graph node representation:
 
     .. math::
-        \mathcal{D}^i = \text{Symmetrize}(\mathbf{n}^L, \mathbf{e}^L),
+        \mathcal D^i = \mathbf n_i^L,
 
     where :math:`L` is the number of repflow layers.
 

deepmd/dpmodel/descriptor/dpa4.py

@@ -142,9 +142,64 @@
 @BaseDescriptor.register("DPA4")
 @BaseDescriptor.register("dpa4")
 class DescrptDPA4(NativeOP, BaseDescriptor):
-    """
+    r"""
     SeZM descriptor.
 
+    DPA4 stores the state of atom :math:`i` at layer :math:`l` as SO(3)
+    coefficients :math:`\mathbf h_i^{(l,\ell,m)}`.  For an edge :math:`j\to i`,
+    the source state is rotated into an edge-aligned frame, processed by an
+    SO(2)-equivariant convolution, and rotated back:
+
+    .. math::
+        \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 :math:`\mathbf D` contains Wigner-D rotation blocks and
+    :math:`\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:
+
+    .. math::
+
+        \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)},
+
+    .. math::
+
+        \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
+    :math:`\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:
+
+    .. math::
+        \mathcal D_i = \operatorname{ScalarReadout}_{\mathrm{mode}}
+        \left(\mathbf h_i^{(L)}\right).
+
+    In ``so3_readout="none"`` mode, coefficients with :math:`\ell>0` are
+    discarded and the :math:`\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 :math:`w_{ji}` are either cutoff-envelope weights or normalized
+    attention weights, depending on ``n_atten_head``.
+
     Execution outline
     -----------------
     1. Build a per-forward `EdgeFeatureCache` (geometry, envelope, Wigner-D).

deepmd/dpmodel/descriptor/repflows.py

@@ -1070,6 +1070,8 @@ def symmetrization_op_dynamic(
 
 
 class RepFlowLayer(NativeOP):
+    r"""Residual node/edge/angle update :math:`(n,e,a)^{l+1}=\Phi_l(n,e,a)`."""
+
     # Mirrors the descriptor-block internal switch.  The owning block writes the
     # instance value during construction/deserialization.
     _use_static_dynamic_sel: bool = False