diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 0dcf6358f9..b6db2272b7 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -14,6 +14,12 @@ # use "class" to resolve "Variable not allowed in type expression" class DPModelCommon: + r"""Common methods for DP models. + + This class provides common functionality for DeepPot models, including + neighbor selection updates and fitting network access. + """ + @classmethod def update_sel( cls, diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index ac90d94fc5..85fe98c7d8 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -31,6 +31,31 @@ @BaseModel.register("ener") class EnergyModel(DPModelCommon, DPEnergyModel_): + r"""Energy model that predicts total energy and derived quantities. + + The model takes atomic energies from the atomic model and computes + global properties by reduction and differentiation: + + **Reduction** (total energy): + + .. math:: + E = \sum_{i=1}^{N} E^i, + + where :math:`E^i` is the atomic energy from the atomic model. + + **Differentiation** (forces and virials): + + .. math:: + \mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i}, + + .. math:: + \boldsymbol{\Xi} = -\sum_{i=1}^{N} \frac{\partial E}{\partial \mathbf{r}_i} \otimes \mathbf{r}_i + = \sum_{i=1}^{N} \mathbf{r}_i \otimes \mathbf{F}_i, + + where :math:`\mathbf{F}_i` is the force on atom :math:`i` and + :math:`\boldsymbol{\Xi}` is the virial tensor. + """ + def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 85e23df3cc..e62442887c 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -29,7 +29,24 @@ class SpinModel(NativeOP): - """A spin model wrapper, with spin input preprocess and output split.""" + r"""A spin model wrapper, with spin input preprocess and output split. + + This model extends a backbone DP model to handle magnetic spin degrees of freedom. + Virtual atoms are created at positions offset from real atoms by their spin vectors: + + .. math:: + \mathbf{r}_i^{\mathrm{virtual}} = \mathbf{r}_i^{\mathrm{real}} + s_i \cdot \boldsymbol{\sigma}_i, + + where :math:`s_i` is a scaling factor and :math:`\boldsymbol{\sigma}_i` is the spin vector. + + The model then computes interactions between real atoms, virtual atoms, and between + real and virtual atoms, enabling the prediction of spin-dependent properties. + + The output forces on virtual atoms are converted to magnetic torques: + + .. math:: + \boldsymbol{\tau}_i = \mathbf{F}_i^{\mathrm{virtual}} \times \boldsymbol{\sigma}_i. + """ def __init__( self,