Add an O(3)-symmetrized atomistic model wrapper by MichelangeloDomina · Pull Request #294 · metatensor/metatomic

MichelangeloDomina · 2026-07-23T15:48:24Z

Summary

Adds metatomic.torch.symmetrized_model: wrap an exported AtomisticModel to
average its predictions over O(3) with a finite rotation quadrature, and to
measure how far the underlying model is from being equivariant.

Everything is selected through regular output requests:

Requested output	Result
`<name>`	O(3) average of the underlying output
`o3::variance::<name>`	per-sample equivariance variance
`o3::character_projection::<name>`	squared character-projection contribution per irrep

SymmetrizedModel.wrap() returns a regular exportable AtomisticModel
(metadata, capabilities, neighbor lists, and custom inputs are carried over),
so the wrapped model can be used anywhere the original could, including
TorchScript export.

The conventions (quadrature, variance definition, irrep bookkeeping in
TensorMap keys) are documented in docs/src/torch/reference/symmetrized-model.rst.

Limitations

the quadrature is exact only up to a maximal angular momentum; convergence
should be checked by increasing max_o3_lambda_grid;
CPU and CUDA only, float32/float64;
explicit TensorBlock gradients are not supported (PyTorch autograd through
returned values works).

Contributor (creator of pull-request) checklist

Tests updated (for new features and bugfixes)?
Documentation updated (for new features)?
~~[ ] Issue referenced (for PRs that solve an issue)?~~

Reviewer checklist

CHANGELOG updated with public API or any other important changes?

Luthaf

I only checked the code, I'll try to do a pass on the tests ASAP

Luthaf · 2026-07-29T14:33:53Z

+O(3)-symmetrized models
+=======================
+
+The :py:mod:`metatomic.torch.symmetrized_model` module wraps an exported


Let's export it as metatomic.torch.SymmetrizedModel, and make the module private

Luthaf · 2026-07-29T14:34:04Z

+O(3)-symmetrized models
+=======================
+
+The :py:mod:`metatomic.torch.symmetrized_model` module wraps an exported


Suggested change

The :py:mod:`metatomic.torch.symmetrized_model` module wraps an exported

The :py:mod:`metatomic.torch.symmetrized_model` module wraps an existing

Luthaf · 2026-07-29T14:35:11Z

+averaging and equivariance diagnostics. Ordinary outputs are averaged over
+rotated and inverted copies of each input. Additional output names request an
+equivariance variance or squared character-projection contributions of the
+model response.


Suggested change

averaging and equivariance diagnostics. Ordinary outputs are averaged over

rotated and inverted copies of each input. Additional output names request an

equivariance variance or squared character-projection contributions of the

model response.

averaging and equivariance diagnostics. Pre-existing outputs of the model are

averaged over rotated and inverted copies of each input.

:py:class:`SymmetrizedModel` also add extra outputs to compute the equivariance

variance or squared character-projection contributions of the model response.

Luthaf · 2026-07-29T14:35:40Z

+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.


This is not very clear, but I can rephrase when everything else is done here =)

Luthaf · 2026-07-29T14:36:49Z

+.. autoclass:: SymmetrizedModel
+    :members:
+
+.. autofunction:: get_rotation_quadrature


why is this part of the public API?

Luthaf · 2026-07-29T14:45:36Z

+import torch
+from metatensor.torch import Labels, TensorBlock, TensorMap
+
+from ._utils import (


functions in ._utils (or this file) that are intended to be imported by other files should not start by an underscore.

Function in a class starting with underscore => the function is private to the class/the module that defines the class

Function in a module starting with underscore => the function is private to the module and should not be imported outside of it.

The module itself is private and should not be used by code outside of metatomic, but it can still provide public functions to be used outside of it.

This way, we easily see if a function is just a current implementation detail or something used outside of the current module/file.

_quadrature.py does this well

Luthaf · 2026-07-29T14:46:04Z

+    return TensorMap(Labels(key_names + ["chi_lambda"], values), blocks)
+
+
+def _character_projection_tensormap_from_cosets(


Could you explain to me what a coset is?

They are in general subsets of groups that can be used to split a group in some interesting way. In this case into proper (SO(3)) and improper (SO(3) * inversion) rotations. Coset is the correct word here (subgroup would be wrong and apply to the SO(3) case only). We could call it "subset", but we would be using a generic word instead of the specific one. I can add a comment or update the docs to explain it

cosets is fine, just might be worth a small comment for people like me who forgot all their group theory

Luthaf · 2026-07-29T14:48:38Z

+from ..o3 import O3Transformation
+
+
+def _build_packed_wigner_matrices(


I don't love that we have a different representation here and in O3Transformation. This is just putting all matrices in the same tensor, right? I assume this is for efficiency purposes?

Could we change the format used by O3Transformation to be the same as this one?

Luthaf · 2026-07-29T14:51:03Z

+        max_o3_lambda_target: int,
+        max_o3_lambda_input: int = 0,
+        max_o3_lambda_character: Optional[int] = None,


Can we auto-guess these from the model by default, by checking the capabilities and extra inputs against the list of standard quantities? And then allow to override for custom inputs/outputs?

For non-standard outputs/more complex things down the line, this would be another use case for #271

MichelangeloDomina requested a review from ppegolo July 23, 2026 15:48

ppegolo force-pushed the review/symmetrized-model-progressive branch from c2acacf to 7e23329 Compare July 28, 2026 09:06

ppegolo marked this pull request as ready for review July 28, 2026 12:42

ppegolo requested review from Luthaf and removed request for ppegolo July 28, 2026 13:20

Luthaf reviewed Jul 29, 2026

View reviewed changes

MichelangeloDomina and others added 8 commits July 30, 2026 13:26

feat(torch): add reviewed SymmetrizedModel core

9a4bcdb

test(torch): cover SymmetrizedModel forward

5ebf4d5

feat(torch): wrap symmetrized models for export

0056d07

fix(torch): harden SymmetrizedModel contracts

1daf33a

Apply o3 review conventions to symmetrized model

5058bd4

Close coverage gaps and use physics terminology in docstrings

6a3a4bb

Address review

5ff2605

Move quantity metadata to _quantities.py

8f87505

ppegolo force-pushed the review/symmetrized-model-progressive branch from f551943 to 8f87505 Compare July 30, 2026 11:26

ppegolo added 3 commits July 30, 2026 14:52

Demote decompose shape validation to asserts

17ebabc

Relax angular momentum inference, rename max_o3_lambda parameters

824277b

Polish SymmetrizedModel API: group parameters, keyword-only init

b97b2e6

	The :py:mod:`metatomic.torch.symmetrized_model` module wraps an exported
	The :py:mod:`metatomic.torch.symmetrized_model` module wraps an existing

		return TensorMap(Labels(key_names + ["chi_lambda"], values), blocks)


		def _character_projection_tensormap_from_cosets(

		from ..o3 import O3Transformation


		def _build_packed_wigner_matrices(

Uh oh!

Conversation

MichelangeloDomina commented Jul 23, 2026 • edited by ppegolo Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Limitations

Contributor (creator of pull-request) checklist

Reviewer checklist

Uh oh!

Luthaf left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ppegolo Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

MichelangeloDomina commented Jul 23, 2026 •

edited by ppegolo

Loading

ppegolo Jul 29, 2026 •

edited

Loading