Add an O(3)-symmetrized atomistic model wrapper

docs/src/torch/reference/index.rst

@@ -11,6 +11,7 @@ API reference
     units
     wrappers
     o3
+    symmetrized-model
     ase
     misc
 

docs/src/torch/reference/o3.rst

@@ -11,9 +11,10 @@ augmentation.
 Conventions
 -----------
 
-To transform a :py:class:`~metatensor.torch.TensorBlock`, :py:func:`transform_block`
-and :py:func:`transform_tensor` need to know, for each component axis, whether it
-carries a Cartesian or a spherical tensor. This is inferred from the axis name:
+To transform a :py:class:`~metatensor.torch.TensorMap`,
+:py:meth:`O3Transformation.transform_tensormap` needs to know, for each component
+axis, whether it carries a Cartesian or a spherical tensor. This is inferred from
+the axis name:
 
 - Cartesian axes are named ``xyz``, or ``xyz_1``, ``xyz_2``, ... for blocks with
   several Cartesian axes (e.g. rank-2 Cartesian tensors). These are rotated directly
@@ -82,9 +83,3 @@ Reference
     :members:
 
 .. autofunction:: metatomic.torch.o3.random_transformations
-
-.. autofunction:: metatomic.torch.o3.transform_system
-
-.. autofunction:: metatomic.torch.o3.transform_tensor
-
-.. autofunction:: metatomic.torch.o3.transform_block

docs/src/torch/reference/symmetrized-model.rst

@@ -0,0 +1,157 @@
+.. _symmetrized-model:
+
+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.
+
+Output requests
+---------------
+
+The requested output name selects both the source output and the calculation:
+
+.. list-table::
+   :header-rows: 1
+
+   * - Requested and returned name
+     - Result
+   * - ``<name>``
+     - O(3) average of the underlying ``<name>`` output
+   * - ``o3::variance::<name>``
+     - component-averaged equivariance variance of ``<name>``
+   * - ``o3::character_projection::<name>``
+     - unnormalized squared character-projection contributions of ``<name>``
+
+``<name>`` 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.
+
+Average and variance
+--------------------
+
+For an input :math:`x`, an O(3) operation :math:`g`, and the target
+representation :math:`\rho_\alpha`, define the response transformed back to the
+input frame as
+
+.. math::
+
+    z_\alpha(g;x) = \rho_\alpha(g^{-1}) f(gx).
+
+The ordinary result is the normalized Haar average
+
+.. math::
+
+    \Pi_\alpha(f,x)
+    = \int_{\mathrm{O}(3)} z_\alpha(g;x)\,\mathrm{d}\mu(g).
+
+For a TensorMap block with component multiplicity :math:`d`, the corresponding
+variance output contains
+
+.. math::
+
+    v_\alpha(f,x)
+    = \frac{1}{d}\left[
+      \int_{\mathrm{O}(3)} \lVert z_\alpha(g;x) \rVert_2^2\,
+      \mathrm{d}\mu(g)
+      - \lVert \Pi_\alpha(f,x) \rVert_2^2
+      \right].
+
+This value is returned separately for every sample and property. It has no
+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
+
+   * - 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.
+
+Character projections
+---------------------
+
+Character projections analyze the direct response :math:`u(g;x)=f(gx)`, rather
+than the back-transformed response used for averaging. For the character sector
+:math:`\beta=(\lambda,\sigma)` with :math:`d_\beta=2\lambda+1`, the squared
+projection norm is
+
+.. math::
+
+    B_\beta(u,x)
+    = d_\beta \iint_{\mathrm{O}(3)}
+      u(g_1;x)^\dagger
+      \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
+equation above.
+
+Quadrature
+----------
+
+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.
+
+Reference
+---------
+
+.. py:currentmodule:: metatomic.torch
+
+.. autoclass:: SymmetrizedModel
+    :members:

metatomic-torch/CHANGELOG.md

@@ -16,17 +16,33 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows
 ### Removed
 -->
 
+### Added
+
+- Added `metatomic.torch.SymmetrizedModel` for finite-quadrature O(3)
+  averaging, equivariance variances, and character projections of existing
+  atomistic models.
+
 ### Changed
 
+- `O3Transformation` now holds a batch of one or more operations, is usable
+  inside scripted models, and gained `inverse`, `with_inversion`,
+  `transform_systems`, and `transform_tensormap`.
 - Renamed `O3Transformation.is_inverted` to `is_improper`.
 - `wigners >= 0.4.0` is now required.
 
+### Removed
+
+- Removed the `transform_system`, `transform_block`, and `transform_tensor`
+  free functions from `metatomic.torch.o3`; use the `O3Transformation`
+  methods `transform_systems` and `transform_tensormap` instead.
+
 ### Fixed
 
 - `O3Transformation.transform_spherical` no longer applies the `(-1)^ell`
   parity factor for proper transformations with `sigma = -1`.
 - Wigner-D evaluation is now stable near the ZYZ Euler-angle poles.
-- `transform_system` now preserves autograd for registered neighbor lists.
+- Transforming a `System` now preserves autograd for registered neighbor
+  lists.
 
 ## [Version 0.1.16](https://github.com/metatensor/metatomic/releases/tag/metatomic-torch-v0.1.16) - 2026-07-13
 

python/metatomic_torch/metatomic/torch/__init__.py

@@ -61,6 +61,7 @@
     is_atomistic_model,
     load_atomistic_model,
 )
+from .o3._symmetrized import SymmetrizedModel  # noqa: F401
 from .serialization import (  # noqa: F401
     load_system,
     load_system_buffer,

python/metatomic_torch/metatomic/torch/_quantities.py

@@ -0,0 +1,88 @@
+"""Cartesian layout and spherical character for standard quantities."""
+
+from typing import Dict
+
+
+def standard_quantity_categories() -> Dict[str, str]:
+    """Return the Cartesian layout and spherical character for standard quantities.
+
+    This is the single source of truth for which outputs and inputs are
+    decomposed; it mirrors ``KNOWN_QUANTITIES`` in
+    ``metatomic-torch/src/quantities.cpp``, minus ``feature``. Only the current
+    (singular) spellings appear here: normalize deprecated aliases with
+    :py:func:`current_quantity_name` before looking them up.
+
+    TorchScript cannot read a module-level dictionary from a compiled function,
+    so the table is built by this function and bound to
+    :py:data:`STANDARD_QUANTITY_CATEGORIES` for Python callers.
+    """
+    return {
+        # scalars: l = 0
+        "charge": "scalar",
+        "energy": "scalar",
+        "energy_ensemble": "scalar",
+        "energy_uncertainty": "scalar",
+        "mass": "scalar",
+        "spin_multiplicity": "scalar",
+        # Cartesian vectors: l = 1
+        "heat_flux": "cartesian_vector",
+        "momentum": "cartesian_vector",
+        "non_conservative_force": "cartesian_vector",
+        "position": "cartesian_vector",
+        "velocity": "cartesian_vector",
+        # symmetric 3x3 matrices: l = 0 and l = 2
+        "non_conservative_stress": "symmetric_matrix",
+    }
+
+
+STANDARD_QUANTITY_CATEGORIES: Dict[str, str] = standard_quantity_categories()
+
+#: maximum angular momentum carried by each category above
+MAX_ANGULAR_MOMENTUM_PER_CATEGORY: Dict[str, int] = {
+    "scalar": 0,
+    "cartesian_vector": 1,
+    "symmetric_matrix": 2,
+}
+
+
+def new_quantity_names() -> Dict[str, str]:
+    """Return the mapping from deprecated quantity names to their current name.
+
+    TorchScript cannot read a module-level dictionary from a compiled function,
+    so the table is built by this function.
+    """
+    return {
+        "features": "feature",
+        "non_conservative_forces": "non_conservative_force",
+        "positions": "position",
+        "momenta": "momentum",
+        "masses": "mass",
+        "velocities": "velocity",
+        "charges": "charge",
+    }
+
+
+def deprecated_quantity_names() -> Dict[str, str]:
+    """Return the mapping from current quantity names to their deprecated name."""
+    result: Dict[str, str] = {}
+    for deprecated, new in new_quantity_names().items():
+        result[new] = deprecated
+    return result
+
+
+def current_quantity_name(name: str) -> str:
+    """Replace a deprecated base quantity in ``name`` with its current name."""
+    base = name.split("/")[0]
+    names = new_quantity_names()
+    if base in names:
+        return name.replace(base, names[base], 1)
+    return name
+
+
+def deprecated_quantity_name(name: str) -> str:
+    """Replace a current base quantity in ``name`` with its deprecated name."""
+    base = name.split("/")[0]
+    names = deprecated_quantity_names()
+    if base in names:
+        return name.replace(base, names[base], 1)
+    return name

python/metatomic_torch/metatomic/torch/model.py

@@ -25,6 +25,7 @@
 )
 from . import __version__ as metatomic_version
 from ._extensions import _collect_extensions
+from ._quantities import current_quantity_name, deprecated_quantity_name
 
 
 def load_atomistic_model(path, extensions_directory=None) -> "AtomisticModel":
@@ -394,28 +395,6 @@ def __init__(
         else:
             raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}")
 
-        # mapping from deprecated output/input names to their new name
-        self._new_names = {
-            "features": "feature",
-            "non_conservative_forces": "non_conservative_force",
-            "positions": "position",
-            "momenta": "momentum",
-            "masses": "mass",
-            "velocities": "velocity",
-            "charges": "charge",
-        }
-
-        # mapping from new names to the corresponding deprecated name
-        self._deprecated_names = {
-            "feature": "features",
-            "non_conservative_force": "non_conservative_forces",
-            "position": "positions",
-            "momentum": "momenta",
-            "mass": "masses",
-            "velocity": "velocities",
-            "charge": "charges",
-        }
-
         # Pretend that the model can output either the new or deprecated names
         new_outputs = {}
         for name in self._model_capabilities_outputs_names:
@@ -521,14 +500,10 @@ def requested_inputs(self, use_new_names: bool = False) -> Dict[str, ModelOutput
         return inputs
 
     def _get_new_name(self, name: str) -> str:
-        base = name.split("/")[0]
-        new_base = self._new_names.get(base, base)
-        return name.replace(base, new_base, 1)
+        return current_quantity_name(name)
 
     def _get_deprecated_name(self, name: str) -> str:
-        base = name.split("/")[0]
-        deprecated_base = self._deprecated_names.get(base, base)
-        return name.replace(base, deprecated_base, 1)
+        return deprecated_quantity_name(name)
 
     def forward(
         self,

python/metatomic_torch/metatomic/torch/o3/__init__.py

@@ -7,19 +7,10 @@
 spherical components in a :py:class:`~metatensor.torch.TensorBlock`.
 """
 
-from ._tranformations import (
-    O3Transformation,
-    random_transformations,
-    transform_block,
-    transform_system,
-    transform_tensor,
-)
+from ._transformations import O3Transformation, random_transformations
 
 
 __all__ = [
     "O3Transformation",
     "random_transformations",
-    "transform_system",
-    "transform_tensor",
-    "transform_block",
 ]