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/symmetrized-model.rst

@@ -0,0 +1,145 @@
+.. _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.
+
+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)`` 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`. Models should provide
+symmetric ``non_conservative_stress`` tensors. Stress diagnostics retain only
+the scalar trace and symmetric-traceless sectors, silently discarding any
+antisymmetric part.
+
+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,6 +16,12 @@ 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 exported
+  atomistic models.
+
 ### Changed
 
 - Renamed `O3Transformation.is_inverted` to `is_improper`.

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,79 @@
+"""
+Python-side mirror of ``metatomic-torch/src/quantities.cpp`` (``KNOWN_QUANTITIES``
+and the per-quantity checks), holding the metadata for standard quantities: their
+category (Cartesian layout and spherical character) and their deprecated-name
+aliases.
+
+This module must not import anything from ``metatomic``, so that any other module
+can import it without creating an import cycle.
+"""
+
+from typing import Dict
+
+
+def standard_quantity_categories() -> Dict[str, str]:
+    """Return the Cartesian layout of every decomposable standard quantity.
+
+    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: deprecated names are normalized before
+    they reach the code using this table.
+
+    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 map 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 and bound to
+    :py:data:`NEW_QUANTITY_NAMES` for Python callers.
+    """
+    return {
+        "features": "feature",
+        "non_conservative_forces": "non_conservative_force",
+        "positions": "position",
+        "momenta": "momentum",
+        "masses": "mass",
+        "velocities": "velocity",
+        "charges": "charge",
+    }
+
+
+NEW_QUANTITY_NAMES: Dict[str, str] = _new_quantity_names()
+
+#: mapping from current quantity names to the corresponding deprecated name
+DEPRECATED_QUANTITY_NAMES: Dict[str, str] = {
+    new: deprecated for deprecated, new in NEW_QUANTITY_NAMES.items()
+}

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 DEPRECATED_QUANTITY_NAMES, NEW_QUANTITY_NAMES
 
 
 def load_atomistic_model(path, extensions_directory=None) -> "AtomisticModel":
@@ -395,26 +396,10 @@ def __init__(
             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",
-        }
+        self._new_names = dict(NEW_QUANTITY_NAMES)
 
         # 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",
-        }
+        self._deprecated_names = dict(DEPRECATED_QUANTITY_NAMES)
 
         # Pretend that the model can output either the new or deprecated names
         new_outputs = {}