Skip to content

Add an O(3)-symmetrized atomistic model wrapper - #294

Open
MichelangeloDomina wants to merge 6 commits into
metatensor:mainfrom
MichelangeloDomina:review/symmetrized-model-progressive
Open

Add an O(3)-symmetrized atomistic model wrapper#294
MichelangeloDomina wants to merge 6 commits into
metatensor:mainfrom
MichelangeloDomina:review/symmetrized-model-progressive

Conversation

@MichelangeloDomina

@MichelangeloDomina MichelangeloDomina commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

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?

@ppegolo
ppegolo force-pushed the review/symmetrized-model-progressive branch from c2acacf to 7e23329 Compare July 28, 2026 09:06
@ppegolo
ppegolo marked this pull request as ready for review July 28, 2026 12:42
@ppegolo
ppegolo requested review from Luthaf and removed request for ppegolo July 28, 2026 13:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants