diff --git a/docs/src/torch/reference/o3.rst b/docs/src/torch/reference/o3.rst index 0177b821..d9dce3c6 100644 --- a/docs/src/torch/reference/o3.rst +++ b/docs/src/torch/reference/o3.rst @@ -8,8 +8,8 @@ augmentation. .. _o3-conventions: -Conventions ------------ +Conventions for O(3) representations +------------------------------------ To transform a :py:class:`~metatensor.torch.TensorMap` (or a :py:class:`~metatensor.torch.TensorBlock`), the diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst index 4cd1056a..4da87e2b 100644 --- a/docs/src/torch/reference/symmetrized-model.rst +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -3,45 +3,46 @@ O(3)-symmetrized models ======================= -The :py:class:`metatomic.torch.SymmetrizedModel` class wraps an exported -:py:class:`~metatomic.torch.AtomisticModel` with finite-quadrature O(3) -averaging and equivariance diagnostics. Pre-existing outputs of the model are -averaged over rotated and inverted copies of each input. -:py:class:`~metatomic.torch.SymmetrizedModel` also adds extra outputs to -compute the equivariance variance or squared character-projection -contributions of the model response. +The :py:class:`metatomic.torch.o3.SymmetrizedModel` class wraps an existing +:py:class:`metatomic.torch.AtomisticModel` with finite-quadrature O(3) averaging +and equivariance diagnostics. Pre-existing outputs of the model are averaged +over rotated and inverted copies of each input. -Output requests ---------------- +Models outputs +-------------- -The requested output name selects both the source output and the calculation: +:py:class:`metatomic.torch.o3.SymmetrizedModel` adds extra outputs to the model, +computing the equivariance variance or squared character-projection +contributions of the model response. .. list-table:: :header-rows: 1 - * - Requested and returned name + * - output name - Result * - ```` - - O(3) average of the underlying ```` output + - O(3) average of the wrapped model's ```` output * - ``o3::variance::`` - component-averaged equivariance variance of ```` * - ``o3::character_projection::`` - unnormalized squared character-projection contributions of ```` -```` is preserved verbatim. It can therefore be a standard quantity, a -variant such as ``energy/pbe``, or a custom name such as -``mtt::feature::node``. For example, -``o3::variance::energy/pbe`` evaluates the underlying ``energy/pbe`` output. - -The meaning of the variance depends on the metadata of the underlying output. -When its blocks carry recognized component labels (``o3_mu``-style spherical -or ``xyz``-style Cartesian axes), each response is rotated back to the input -frame first, and the variance measures the breaking of *equivariance*. Outputs -without such labels cannot be rotated back: their responses are compared as-is -across the quadrature, so their variance measures the deviation from -*invariance* only. An equivariant but unlabelled output — for example an -internal feature vector — reports a large variance even when it transforms -correctly. +The above outputs are added for every output of the wrapped model, including +variants (such as ``energy/pbe``) and custom outputs (such as +``custom::feature::node``). For example, ``o3::variance::energy/pbe`` would +compute the equivariance variance of the ``energy/pbe`` output. + +Quadrature +---------- + +The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced +in-plane rotations, and both parities: O(3) splits into two cosets of SO(3), +the proper rotations, and the improper ones (a rotation composed with +inversion). Its weights are normalized to sum to one. A general machine-learning +model need not be band-limited, so a finite grid is not automatically exact. +``max_angular_momentum_grid`` controls the quadrature resolution, not the +representation: increase it until the averages, variances, and character +projections of interest converge. Average and variance -------------------- @@ -78,40 +79,72 @@ component axes, and it is not reduced across samples or square-rooted. A weighted mean of these values over a group of samples, followed by a square root, gives a block-wise equivariance RMSE. -TensorMap representation ------------------------- - -An averaged output retains the physical schema declared by the source model. -For diagnostics, the standard quantities are represented as follows: - -.. list-table:: - :header-rows: 1 +The meaning of the variance depends on the structure of the wrapped model's +output. When its blocks carry recognized component labels (``o3_mu``-style +spherical or ``xyz``-style Cartesian axes, see the :ref:`o3-conventions` +documentation), each output is rotated back to the input frame first, and the +variance measures the breaking of *equivariance*. Outputs without such +components cannot be rotated back: their responses are compared as-is across the +quadrature, so their variance measures the deviation from *invariance* only. An +equivariant but unlabelled output --- for example an internal equivariant +feature vector --- can thus report a large variance even when it transforms +correctly. - * - Source quantity - - Diagnostic target keys - * - ``energy``, ``energy_ensemble``, ``energy_uncertainty`` - - ``o3_lambda=0``, ``o3_sigma=1`` - * - ``non_conservative_force`` - - ``o3_lambda=1``, ``o3_sigma=1`` - * - ``non_conservative_stress`` - - ``(o3_lambda, o3_sigma)=(0,1)``, ``(1,-1)``, and ``(2,1)`` - -Variants after ``/`` use the same representation as their base quantity. - -Energy-like scalars acquire an ``o3_mu`` component of size one for diagnostics. -Cartesian force components are reordered into the real spherical -:math:`\ell=1` basis described in :ref:`o3-conventions`. Stress diagnostics -cover the full matrix: the scalar trace, the antisymmetric (axial pseudovector, -:math:`\ell=1` with ``o3_sigma=-1``) part, and the symmetric-traceless sector. -For a symmetric stress the pseudovector sector is exactly zero; a model -producing a non-symmetric stress (before any downstream symmetrization) sees -its antisymmetric response in this sector. - -Already-spherical outputs retain their ``o3_lambda`` and ``o3_sigma`` keys and -``o3_mu`` components, and other semantic source keys are preserved. The wrapper -does not infer the physical meaning of a custom output from its shape; in -particular, a custom Cartesian :math:`3\times3` output is not treated as a -symmetric stress. +Variance metadata +~~~~~~~~~~~~~~~~~ + +The ``o3::variance::`` outputs produced by +:py:class:`metatomic.torch.o3.SymmetrizedModel` have the following metadata +structure: + +.. list-table:: Metadata for ``"o3::variance::"`` + :widths: 2 3 7 + :header-rows: 1 + + * - Metadata + - Names + - Description + + * - keys + - ``[, "o3_lambda", "o3_sigma"]`` + - The keys are the same as the original ```` output, with + ``"o3_lambda"`` and ``"o3_sigma"`` dimensions added if they are not + already present, and the ``_ = 0`` dummy key removed if present. The + ``"o3_lambda"`` dimension contains the angular momentum of each block, and + the ``"o3_sigma"`` dimension contains its parity under inversion. + + * - samples + - ``[]`` + - the samples are the same as the original ```` output. + + * - components + - + - Since the variance is computed over the ``o3_mu`` components of each + block, the resulting TensorMap does not have any component axes. + + * - properties + - ``[]`` + - the properties are the same as the original ```` output. + +When computing the variance of :ref:`standard quantities `, +the data is first converted to spherical representation as follow: + +- Scalar quantities (such as :ref:`energy `, :ref:`charge + `, *etc.*) gain a ``o3_lambda=0, o3_sigma=1`` key, as well as + an ``o3_mu`` component of size one. +- Cartesian vector quantities (such as :ref:`non-conservative force + `) gain a ``o3_lambda=1, o3_sigma=1`` key, + and their ``xyz`` components are replaced by an ``o3_mu`` component of size + three. +- Cartesian rank-2 tensor quantities (such as :ref:`non-conservative stress + `) gain ``o3_lambda=0, o3_sigma=1``, + ``o3_lambda=1, o3_sigma=-1``, and ``o3_lambda=2, o3_sigma=1`` keys, and their + ``xyz_1`` and ``xyz_2`` components are replaced by an ``o3_mu`` component of + size one, three, and five, respectively. +- Already-spherical outputs retain their ``o3_lambda`` and ``o3_sigma`` keys and + ``o3_mu`` components, and other keys are preserved. +- Custom outputs with components that do not match the :ref:`convention ` are + not supported, and will raise an error when the variance is requested. Character projections --------------------- @@ -129,29 +162,57 @@ projection norm is \chi_\beta(g_1g_2^{-1})u(g_2;x)\, \mathrm{d}\mu(g_1)\,\mathrm{d}\mu(g_2). -Character results append ``chi_lambda`` and ``chi_sigma`` to the TensorMap -keys. These labels describe the O(3) dependence of the response over the -rotation orbit. They are distinct from ``o3_lambda`` and ``o3_sigma``, which -describe the target representation of the output itself. Target component axes -are retained; summing over them gives the complete component norm in the +Character results append ``chi_lambda`` and ``chi_sigma`` to the TensorMap keys. +These labels describe the O(3) dependence of the response over the rotation +orbit. They are distinct from ``o3_lambda`` and ``o3_sigma``, which describe the +target representation of the output itself. Any other pre-existing component +axes are retained; summing over them gives the complete component norm in the equation above. -Quadrature ----------- +Character projections metadata +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced -in-plane rotations, and both O(3) cosets: O(3) splits into two cosets of SO(3), -the proper rotations, and the improper ones (a rotation composed with -inversion). Its weights are normalized to sum to one. A general -machine-learning model need not be band-limited, so a finite grid is not -automatically exact. ``max_angular_momentum_grid`` controls the quadrature -resolution, not the representation: increase it until the averages, variances, -and character projections of interest converge. +The ``o3::character_projection::`` outputs produced by +:py:class:`metatomic.torch.o3.SymmetrizedModel` have the following metadata +structure: + +.. list-table:: Metadata for ``"o3::variance::"`` + :widths: 2 3 7 + :header-rows: 1 + + * - Metadata + - Names + - Description + + * - keys + - ``[, "o3_lambda", "o3_sigma", "chi_lambda", "chi_sigma"]`` + - The keys are the same as the original ```` output, with + ``"o3_lambda"`` and ``"o3_sigma"`` dimensions added if they are not + already present, the ``_ = 0`` dummy key removed if present, and + ``"chi_lambda"`` and ``"chi_sigma"`` keys added. + + * - samples + - ``[]`` + - the samples are the same as the original ```` output. + + * - components + - ``[, "o3_mu"]`` + - The components are the same as the original ```` output, with an + additional ``"o3_mu"`` component of size :math:`2\lambda+1` added for each + block. This ``"o3_mu"`` component replaces existing ``"xyz"`` or + ``"o3_mu"`` components, if present. + + * - properties + - ``[]`` + - the properties are the same as the original ```` output. -Reference ---------- +When computing character projections, the data is first converted to spherical +representation like for the variance, with the exception that arbitrary +pre-existing component axes are allowed and will not raise an exception. These +are retained as-is in the output. -.. py:currentmodule:: metatomic.torch +API reference +------------- -.. autoclass:: SymmetrizedModel +.. autoclass:: metatomic.torch.o3.SymmetrizedModel :members: diff --git a/python/metatomic_torch/metatomic/torch/__init__.py b/python/metatomic_torch/metatomic/torch/__init__.py index bcc12ec7..06a9ae9c 100644 --- a/python/metatomic_torch/metatomic/torch/__init__.py +++ b/python/metatomic_torch/metatomic/torch/__init__.py @@ -61,7 +61,6 @@ is_atomistic_model, load_atomistic_model, ) -from .o3 import SymmetrizedModel # noqa: F401 from .serialization import ( # noqa: F401 load_system, load_system_buffer, diff --git a/python/metatomic_torch/tests/o3/symmetrized.py b/python/metatomic_torch/tests/o3/symmetrized.py index 077f48e3..e26949ed 100644 --- a/python/metatomic_torch/tests/o3/symmetrized.py +++ b/python/metatomic_torch/tests/o3/symmetrized.py @@ -14,11 +14,10 @@ ModelMetadata, ModelOutput, NeighborListOptions, - SymmetrizedModel, System, load_atomistic_model, ) -from metatomic.torch.o3 import O3Transformations +from metatomic.torch.o3 import O3Transformations, SymmetrizedModel # These helpers back the exported SymmetrizedModel and have no public entry # point: decompose_quantity is its Cartesian-to-spherical boundary (tested here