diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index ee406c11ae..38134ed752 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/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], diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py index 7cfa24526a..c78eb44d5d 100644 --- a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py +++ b/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, diff --git a/deepmd/dpmodel/atomic_model/dos_atomic_model.py b/deepmd/dpmodel/atomic_model/dos_atomic_model.py index ce457cb472..2241fdd6f9 100644 --- a/deepmd/dpmodel/atomic_model/dos_atomic_model.py +++ b/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, diff --git a/deepmd/dpmodel/atomic_model/energy_atomic_model.py b/deepmd/dpmodel/atomic_model/energy_atomic_model.py index 6deb87662d..8981d4825e 100644 --- a/deepmd/dpmodel/atomic_model/energy_atomic_model.py +++ b/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: diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index 208a7c96be..3462ed59d2 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/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: diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index f8588f7330..fef23a66cd 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/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: diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 3a47a7319c..de967af9f6 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/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, diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 08f68849d2..696c912c63 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/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, diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 76124fa544..dab6c193c1 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/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. diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 8d59a5caba..11f96c4669 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/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). diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 5d1ee9fd24..627184dd05 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index 1207d87baa..d20f7703f7 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -965,6 +965,34 @@ def symmetrization_op( class Atten2Map(NativeOP): + r"""Masked and smoothed angular attention map. + + Define the raw logits and the (not necessarily unit-vector) angular factor + by + + .. math:: + + L_{ijk}^{(h)}=\frac{q_{ij}^{(h)}\cdot k_{ik}^{(h)}}{\sqrt d}, + \qquad + G_{ijk}=\mathbf h_{ij}\cdot\mathbf h_{ik}. + + When ``has_gate=True``, :math:`G_{ijk}` also gates the logits, so let + :math:`\widetilde L=L G`; otherwise let :math:`\widetilde L=L`. In the + smooth path the normalized attention and returned map are + + .. math:: + + \overline A_{ijk}^{(h)}=\operatorname{softmax}_k\!\left( + (\widetilde L_{ijk}^{(h)}+c)s_{ij}s_{ik}-c\right),\qquad + M_{ijk}^{(h)}=s_{ij}s_{ik}\overline A_{ijk}^{(h)} + \frac{G_{ijk}}{\sqrt 3}. + + Without smoothing, masked softmax replaces the first expression and the + cutoff factors are omitted. Invalid query and key entries are zeroed. The + factor :math:`G_{ijk}/\sqrt 3` multiplies the returned map unconditionally, + including when ``has_gate=False``. + """ + def __init__( self, input_dim: int, @@ -1094,6 +1122,17 @@ def deserialize(cls, data: dict) -> "Atten2Map": class Atten2MultiHeadApply(NativeOP): + r"""Multi-head attention application. + + ``Atten2Map`` has already applied masking, smoothing, and softmax. This + stage therefore computes the value aggregation and head projection only: + + .. math:: + O_{ij}^{(h)}=\sum_k A_{ijk}^{(h)}V_{ik}^{(h)},\qquad + O_{ij}=\operatorname{HeadMap}\!\left(\operatorname{concat}_h + O_{ij}^{(h)}\right). + """ + def __init__( self, input_dim: int, @@ -1185,6 +1224,20 @@ def deserialize(cls, data: dict) -> "Atten2MultiHeadApply": class Atten2EquiVarApply(NativeOP): + r"""Equivariant attention application preserving rotation laws. + + For attention matrix :math:`A_{ijk}^{(h)}` and equivariant neighbor vectors + :math:`\mathbf h_{ik}`, each head produces + + .. math:: + \mathbf u_{ij}^{(h)} = \sum_k A_{ijk}^{(h)}\mathbf h_{ik}, + \qquad + \mathbf u_{ij} = \sum_h w_h\mathbf u_{ij}^{(h)}. + + The attention weights and head mixing coefficients are scalars, so the + output follows the same rotation law as :math:`\mathbf h`. + """ + def __init__( self, input_dim: int, @@ -1263,6 +1316,20 @@ def deserialize(cls, data: dict) -> "Atten2EquiVarApply": class LocalAtten(NativeOP): + r"""Local attention with masked, smoothed softmax weights. + + For each head, the query is formed from the destination feature and keys + and values from its neighbors. The logits use the standard scaling and + the cutoff-smoothed softmax when enabled: + + .. math:: + L_{ij}^{(h)}=q_i^{(h)}\cdot k_{ij}^{(h)}/\sqrt{d},\qquad + A_{ij}^{(h)}=s_{ij}\operatorname{softmax}_j\!\left( + (L_{ij}^{(h)}+c)s_{ij}-c\right),\qquad + O_i=\operatorname{HeadMap}\!\left(\operatorname{concat}_h + \sum_j A_{ij}^{(h)}v_{ij}^{(h)}\right). + """ + def __init__( self, input_dim: int, @@ -1404,6 +1471,8 @@ def deserialize(cls, data: dict) -> "LocalAtten": class RepformerLayer(NativeOP): + r"""Residual representation update :math:`(G_1,G_2,H_2)^{l+1}=\Phi_l(...)`.""" + def __init__( self, rcut: float, diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index 8d088656d5..b15935602a 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -32,6 +32,12 @@ @InvarFitting.register("dos") class DOSFittingNet(InvarFitting): + # Compose rather than replace the base documentation so AutoAPI/numpydoc + # retains the fitting equations and complete Parameters section. + __doc__ = r"""Fit DOS vectors :math:`\mathbf D_i=F_\theta(\mathcal D_i)`. + + """ + (InvarFitting.__doc__ or "") + def __init__( self, ntypes: int, diff --git a/deepmd/dpmodel/fitting/dpa4_ener.py b/deepmd/dpmodel/fitting/dpa4_ener.py index f08097f032..ef3597af16 100644 --- a/deepmd/dpmodel/fitting/dpa4_ener.py +++ b/deepmd/dpmodel/fitting/dpa4_ener.py @@ -38,9 +38,25 @@ class GLUFittingNet(NativeOP): - """ + r""" GLU-based fitting network for SeZM. + Each hidden layer splits its affine output into a value and a gate. For + input :math:`\mathbf x^{(l)}`, the update is + + .. math:: + [\mathbf v^{(l)},\mathbf g^{(l)}] + = \mathbf x^{(l)}\mathbf W^{(l)}+\mathbf b^{(l)}, + \qquad + \mathbf x^{(l+1)}=\mathbf v^{(l)}\odot + \phi(\mathbf g^{(l)}), + + followed by the output projection + + .. math:: + \mathbf y=\mathbf x^{(L)}\mathbf W^{\mathrm{out}} + +\mathbf b^{\mathrm{out}}. + Parameters ---------- in_dim diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 713aef5117..44a5e2d0c4 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -22,6 +22,12 @@ @InvarFitting.register("ener") class EnergyFittingNet(InvarFitting): + # Compose rather than replace the base documentation so AutoAPI/numpydoc + # retains the fitting equations and complete Parameters section. + __doc__ = r"""Fit scalar atomic energies :math:`E_i=F_\theta(\mathcal D_i)`. + + """ + (InvarFitting.__doc__ or "") + def __init__( self, ntypes: int, diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py index 953b413dd1..15d94ef8f4 100644 --- a/deepmd/dpmodel/loss/dos.py +++ b/deepmd/dpmodel/loss/dos.py @@ -22,6 +22,25 @@ class DOSLoss(Loss): r"""Loss on DOS (density of states) for both local and global predictions. + For DOS bin :math:`k`, define its discrete cumulative distribution by + + .. math:: + + C_k=\sum_{q=0}^{k}D_q. + + The optimized objective combines global and atomic DOS and CDF errors, + + .. math:: + + L=p_D\langle(D-\hat D)^2\rangle + +p_C\langle(C-\hat C)^2\rangle + +p_{D_i}\langle(D_i-\hat D_i)^2\rangle + +p_{C_i}\langle(C_i-\hat C_i)^2\rangle. + + Padded atoms are excluded from the atomic averages. Each prefactor uses + :math:`p=p_{\mathrm{limit}}+(p_{\mathrm{start}}-p_{\mathrm{limit}}) + \eta/\eta_0`. + Parameters ---------- starter_learning_rate : float @@ -101,7 +120,11 @@ def call( label_dict: dict[str, Array], mae: bool = False, ) -> tuple[Array, dict[str, Array]]: - """Calculate loss from model results and labeled results.""" + r"""Evaluate the weighted DOS and cumulative-DOS objective. + + The cumulative terms use :math:`C_k=\sum_{q\le k}D_q`; local terms + are averaged only over real atoms when a mask is present. + """ # Get array namespace from any available tensor first_key = next(iter(model_dict)) xp = array_api_compat.array_namespace(model_dict[first_key]) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 5ee4916762..6670fb26ff 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -23,6 +23,17 @@ def custom_huber_loss(predictions: Array, targets: Array, delta: float = 1.0) -> Array: + r"""Return the mean Huber loss. + + For residual :math:`e=y-\hat y`, the elementwise loss is + + .. math:: + + H_\delta(e)=\begin{cases} + \tfrac12 e^2,& |e|\le\delta,\\ + \delta(|e|-\tfrac12\delta),& |e|>\delta. + \end{cases} + """ xp = array_api_compat.array_namespace(predictions, targets) error = targets - predictions abs_error = xp.abs(error) @@ -35,6 +46,27 @@ def custom_huber_loss(predictions: Array, targets: Array, delta: float = 1.0) -> class EnergyLoss(Loss): r"""Construct a layer to compute loss on energy, force and virial. + The total objective is a weighted sum of the enabled error terms, + + .. math:: + + L=p_E L_E+p_F L_F+p_\Xi L_\Xi+p_{E_i}L_{E_i} + +p_{PF}L_{PF}+p_{GF}L_{GF}. + + Each prefactor is interpolated using the current learning rate + :math:`\eta` as + + .. math:: + + p(\eta)=p_{\mathrm{limit}}+ + (p_{\mathrm{start}}-p_{\mathrm{limit}}) + \frac{\eta}{\eta_0}. + + The individual terms are mean squared, mean absolute, or Huber errors as + configured. In relative-force mode, each force residual is divided by + :math:`\lVert\hat{\mathbf F}_i\rVert+\nu`, where :math:`\nu` is + ``relative_f``. + Parameters ---------- starter_learning_rate : float @@ -194,7 +226,16 @@ def call( label_dict: dict[str, Array], mae: bool = False, ) -> tuple[Array, dict[str, Array]]: - """Calculate loss from model results and labeled results.""" + r"""Calculate the weighted energy-model objective. + + This evaluates the objective and learning-rate-dependent prefactors + defined in :class:`EnergyLoss`. The diagnostics contain per-term RMSE + values in MSE/Huber mode and per-term MAE values when ``loss_func`` is + ``"mae"`` or ``mae=True``. RMSE diagnostics remain ordinary residual + RMSEs when the optimized objective uses Huber loss. The aggregate + ``rmse`` entry is :math:`\sqrt{L}` for the fully weighted objective, + including all enabled prefactors and any configured Huber terms. + """ energy = model_dict["energy"] force = model_dict["force"] virial = model_dict["virial"] diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index e39e0a6a9e..d736c074d0 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -22,6 +22,32 @@ class EnergySpinLoss(Loss): r"""Loss on energy, real force, magnetic force and virial for spin models. + For mean-squared error, the objective is a weighted sum + + .. math:: + + L = p_E\left\langle\frac{(\Delta E)^2}{N^q}\right\rangle + +p_{F_r}\left\langle(\Delta F^r_{i\alpha})^2 + \right\rangle_{i,\alpha} + +p_{F_m}\left\langle(\Delta F^m_{i\alpha})^2 + \right\rangle_{i\in\mathcal M,\alpha} + +p_\Xi\left\langle\frac{(\Delta\Xi_{\alpha\beta})^2}{N^q} + \right\rangle_{\alpha,\beta} + +p_{E_i}\langle(\Delta E_i)^2\rangle, + + where :math:`q=2` for intensive energy/virial normalization and :math:`q=1` + for the legacy normalization, and :math:`\mathcal M` is the set selected by + ``mask_mag``. Thus force and virial errors are componentwise means rather + than means of squared vector or tensor norms. In MAE mode the squared + component differences are replaced by absolute differences and extensive + terms use :math:`1/N`. Every prefactor is interpolated using the current + learning rate, + + .. math:: + p_X(\eta)=p_X^{\mathrm{limit}}+ + \left(p_X^{\mathrm{start}}-p_X^{\mathrm{limit}}\right) + \frac{\eta}{\eta_0}. + Parameters ---------- starter_learning_rate : float diff --git a/deepmd/dpmodel/loss/loss.py b/deepmd/dpmodel/loss/loss.py index 7d4e052dca..3f45b12b8c 100644 --- a/deepmd/dpmodel/loss/loss.py +++ b/deepmd/dpmodel/loss/loss.py @@ -21,6 +21,12 @@ class Loss(NativeOP, ABC, make_plugin_registry("loss")): + r"""Base interface for objectives :math:`L(\hat y,y)`. + + Concrete losses map model outputs and labels to a scalar objective and + return diagnostic metrics alongside it. + """ + @abstractmethod def call( self, diff --git a/deepmd/dpmodel/loss/property.py b/deepmd/dpmodel/loss/property.py index ef649c9b2f..a25db77df8 100644 --- a/deepmd/dpmodel/loss/property.py +++ b/deepmd/dpmodel/loss/property.py @@ -22,6 +22,26 @@ class PropertyLoss(Loss): r"""Loss on property predictions. + Except for MAPE, the residual is evaluated in normalized property space, + + .. math:: + + e_k=\frac{y_k-b_k}{s_k}-\frac{\hat y_k-b_k}{s_k}. + + Extensive properties are first divided by the number of real atoms. + ``smooth_mae`` uses + + .. math:: + + \ell_\beta(e)=\begin{cases} + e^2/(2\beta),&|e|<\beta,\\ + |e|-\beta/2,&|e|\ge\beta. + \end{cases} + + The other choices are the usual :math:`\sum|e|`, :math:`\sum e^2`, + :math:`\sqrt{\langle e^2\rangle}`, and + :math:`\langle |(y-\hat y)/(y+10^{-3})|\rangle`. + Parameters ---------- task_dim : int @@ -75,7 +95,12 @@ def call( label_dict: dict[str, Array], mae: bool = False, ) -> tuple[Array, dict[str, Array]]: - """Calculate loss from model results and labeled results.""" + r"""Evaluate the configured normalized property error. + + For extensive targets, both prediction and label are divided by the + real atom count before the equations in :class:`PropertyLoss` are + applied. + """ del learning_rate, mae var_name = self.var_name pred = model_dict[var_name] diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index 4b21175185..5a036aa4a0 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -22,6 +22,19 @@ class TensorLoss(Loss): r"""Loss on local and global tensors (e.g. dipole, polarizability). + With atomic tensors :math:`T_i`, global tensor :math:`T`, and optional + atomic scalar weights :math:`a_i`, the objective is + + .. math:: + + L=p_{\mathrm{atom}}\left\langle + [a_i(T_i-\hat T_i)]^2\right\rangle + +p_{\mathrm{global}}\left\langle(T-\hat T)^2\right\rangle. + + Padded atoms are omitted from the local mean. Reported global RMSE is + divided by the number of real atoms, matching the convention for + extensive tensor labels. + Parameters ---------- tensor_name : str @@ -74,7 +87,11 @@ def call( label_dict: dict[str, Array], mae: bool = False, ) -> tuple[Array, dict[str, Array]]: - """Calculate loss from model results and labeled results.""" + r"""Evaluate the weighted local and global tensor MSE. + + This applies the objective in :class:`TensorLoss`, including optional + per-atom weights and padding masks. + """ del learning_rate, mae first_key = next(iter(model_dict)) xp = array_api_compat.array_namespace(model_dict[first_key]) diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py index 9e85403e9f..8c1242cfc1 100644 --- a/deepmd/dpmodel/model/dipole_model.py +++ b/deepmd/dpmodel/model/dipole_model.py @@ -28,6 +28,10 @@ @BaseModel.register("dipole") class DipoleModel(DPModelCommon, DPDipoleModel_): + r"""Dipole model reducing atomic contributions as + :math:`\boldsymbol\mu=\sum_i\boldsymbol\mu_i`. + """ + def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index bc9ee4fcb0..39783127ef 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -28,6 +28,8 @@ @BaseModel.register("dos") class DOSModel(DPModelCommon, DPDOSModel_): + r"""DOS model with global spectrum :math:`D_k=\sum_iD_{ik}`.""" + def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index d389c0b7b2..d4ff3ed800 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -31,6 +31,32 @@ @BaseModel.register("zbl") class DPZBLModel(DPZBLModel_): + r"""Combine DP and ZBL atomic energies by a short-range switch. + + For the soft minimum neighbor distance :math:`\sigma_i`, define + :math:`u_i=(\sigma_i-r_{\min})/(r_{\max}-r_{\min})` and + + .. math:: + + s_i=\begin{cases} + 1,&\sigma_i Array: - """Compute smooth weight for descriptor elements.""" + r"""Compute the quintic cutoff weight. + + With :math:`u=(r-r_s)/(r_c-r_s)` clipped to :math:`[0,1]`, the weight is + + .. math:: + + w(r)=1-10u^3+15u^4-6u^5. + + Thus :math:`w=1` for :math:`r\le r_s`, :math:`w=0` for + :math:`r\ge r_c`, and the value and first two derivatives are continuous + at both boundaries. + """ if rmin >= rmax: raise ValueError("rmin should be less than rmax.") xp = array_api_compat.array_namespace(distance) @@ -42,7 +53,15 @@ def compute_exp_sw( rmin: float, rmax: float, ) -> Array: - """Compute the exponential switch function for neighbor update.""" + r"""Compute the exponential switch used for neighbor updates. + + For a clipped distance :math:`\bar r=\min(\max(r,0),r_c)`, this computes + + .. math:: + + w(r)=\exp\!\left[-\exp\!\left(\frac{20}{r_s} + (\bar r-r_s)\right)\right]. + """ if rmin >= rmax: raise ValueError("rmin should be less than rmax.") xp = array_api_compat.array_namespace(distance) @@ -105,6 +124,24 @@ def _make_env_mat( class EnvMat(NativeOP): + r"""Construct radial or full local environment matrices. + + For center atom :math:`i` and neighbor :math:`j`, let + :math:`\mathbf r_{ji}=\mathbf r_j-\mathbf r_i` and + :math:`r_{ji}=\lVert\mathbf r_{ji}\rVert`. Before optional + type-dependent normalization, the full row is + + .. math:: + + \mathcal R_{ij}=w(r_{ji})\left( + \frac{1}{r_{ji}+p}, + \frac{\mathbf r_{ji}}{(r_{ji}+p)^2}\right), + + where :math:`p` is ``protection``. In radial-only mode only the first + component is returned. If statistics are supplied, the final matrix is + :math:`(\mathcal R_{ij}-\mu_{t_i j})/\sigma_{t_i j}`. + """ + def __init__( self, rcut: float, @@ -126,7 +163,11 @@ def call( dstd: Array | None = None, radial_only: bool = False, ) -> tuple[Array, Array, Array]: - """Compute the environment matrix. + r"""Compute the environment matrix. + + This evaluates the :class:`EnvMat` equation for every neighbor-list + entry, masks padded neighbors, and applies the optional normalization + :math:`(\mathcal R-\mathrm{davg})/\mathrm{dstd}`. Parameters ---------- diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 543b918b42..e713944043 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -114,6 +114,8 @@ def merge_env_stat( class EnvMatStat(BaseEnvMatStat): + r"""Environment statistics estimating :math:`\mu=\langle R\rangle` and scale.""" + def compute_stat(self, env_mat: dict[str, Array]) -> dict[str, StatItem]: """Compute the statistics of the environment matrix for a single system. diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index d26432528c..75fad6782f 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -26,6 +26,8 @@ class BaseLR(ABC, PluginVariant, make_plugin_registry("lr")): + r"""Base learning-rate schedule :math:`\eta=\eta(t)`.""" + def __new__(cls: type, *args: Any, **kwargs: Any) -> Any: if cls is BaseLR: cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__)) diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index be98f2b67b..18b89730eb 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -57,6 +57,8 @@ def softplus_t(x): # noqa: ANN001, ANN201 class Identity(NativeOP): + r"""Identity map :math:`y=x`.""" + def __init__(self) -> None: super().__init__() @@ -76,7 +78,18 @@ def deserialize(cls, data: dict) -> "Identity": class NativeLayer(NativeOP): - """Native representation of a layer. + r"""Native affine neural-network layer with optional residual connection. + + The activated branch is + + .. math:: + + \mathbf z=\boldsymbol\phi(\mathbf x\mathbf W+\mathbf b), + + optionally scaled elementwise by a trainable timestep :math:`\mathbf d`. + If ``resnet`` is enabled, the output is :math:`\mathbf z+\mathbf x` for + equal widths, or :math:`\mathbf z+[\mathbf x,\mathbf x]` when the output + width is twice the input width. Parameters ---------- @@ -266,7 +279,13 @@ def dim_out(self) -> int: return self.w.shape[1] def call(self, x): # noqa: ANN001, ANN201 - """Forward pass. + r"""Evaluate the affine, activation, timestep, and residual operations. + + In compact form, this returns + :math:`\mathbf y=\mathbf r(\mathbf x)+ + \mathbf d\odot\boldsymbol\phi(\mathbf x\mathbf W+\mathbf b)`, where + :math:`\mathbf r` is the configured identity/duplicated residual and + :math:`\mathbf d=\mathbf 1` when no timestep is used. Parameters ---------- @@ -404,7 +423,17 @@ def fn(x): # noqa: ANN001, ANN202 class LayerNorm(NativeLayer): - """Implementation of Layer Normalization layer. + r"""Implementation of layer normalization. + + For features normalized over the final axes, this layer computes + + .. math:: + + \mathbf y=\boldsymbol\gamma\odot + \frac{\mathbf x-\mu}{\sqrt{\sigma^2+\epsilon}}+\boldsymbol\beta, + + where :math:`\mu` and :math:`\sigma^2` are the mean and variance over + those axes. Parameters ---------- @@ -540,7 +569,7 @@ def dim_out(self) -> int: return self.w.shape[0] def call(self, x): # noqa: ANN001, ANN201 - """Forward pass. + r"""Normalize :math:`\mathbf x` and apply learned scale and shift. Parameters ---------- diff --git a/deepmd/dpmodel/utils/spherical_harmonics.py b/deepmd/dpmodel/utils/spherical_harmonics.py index fc45abf5f7..cc9066bb9f 100644 --- a/deepmd/dpmodel/utils/spherical_harmonics.py +++ b/deepmd/dpmodel/utils/spherical_harmonics.py @@ -38,12 +38,21 @@ def real_spherical_harmonics(vecs: np.ndarray, lmax: int) -> np.ndarray: - """ - Evaluate real spherical harmonics in the e3nn ``"norm"`` convention. + r"""Evaluate real spherical harmonics in the e3nn ``"norm"`` convention. Exactly matches ``e3nn.o3.spherical_harmonics(list(range(lmax + 1)), vecs, normalize=True, normalization="norm")`` (see module docstring). + For :math:`m>0`, using the e3nn polar-axis convention, the components are + + .. math:: + + Y_{l,-m}=\sqrt{2}\,N_{lm}P_l^m(\cos\theta)\sin(m\phi),\qquad + Y_{l,m}=\sqrt{2}\,N_{lm}P_l^m(\cos\theta)\cos(m\phi), + + with :math:`Y_{l,0}=N_{l0}P_l(\cos\theta)` and normalization chosen so + that :math:`\sum_{m=-l}^lY_{lm}(\hat{\mathbf r})^2=1`. + Parameters ---------- vecs diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index f945d0f801..7bda00081d 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -34,6 +34,17 @@ def _array_device_or_none(array: Array) -> Any: class TypeEmbedNet(NativeOP): r"""Type embedding network. + Each atom type :math:`t` is represented by a one-hot vector + :math:`\mathbf e_t` (or an electronic-configuration vector), then mapped + by an embedding network :math:`\mathcal N`: + + .. math:: + + \mathbf T_t=\mathcal N(\mathbf e_t). + + If ``padding`` is enabled, an additional all-zero row represents padded + neighbor-list entries. + Parameters ---------- ntypes : int @@ -103,7 +114,7 @@ def __init__( ) def call(self) -> Array: - """Compute the type embedding network.""" + r"""Return all type embeddings :math:`\mathbf T_t=\mathcal N(\mathbf e_t)`.""" sample_array = self.embedding_net[0]["w"] xp = array_api_compat.array_namespace(sample_array) if not self.use_econf_tebd: diff --git a/deepmd/dpmodel/utils/update_sel.py b/deepmd/dpmodel/utils/update_sel.py index 3f2900771f..a340fe10f8 100644 --- a/deepmd/dpmodel/utils/update_sel.py +++ b/deepmd/dpmodel/utils/update_sel.py @@ -9,6 +9,8 @@ class UpdateSel(BaseUpdateSel): + r"""Neighbor-selection update computing :math:`n_{sel}` from statistics.""" + @property def neighbor_stat(self) -> type[NeighborStat]: return NeighborStat