Improve O(3) transformations and Wigner-D evaluation by MichelangeloDomina · Pull Request #293 · metatensor/metatomic

MichelangeloDomina · 2026-07-21T15:04:19Z

This PR extracts and improves the O(3)-only part initially developed in #119.

What changes

Correct spherical parity for improper transformations: the additional factor is o3_sigma * (-1) ** o3_lambda and is applied only to the improper O(3) component.
Make real Wigner-D evaluation stable near both ZYZ Euler-angle poles and use the structured wigner_D_array API directly.
Store an independent copy of the matrix passed to O3Transformation, and construct the Wigner cache only when spherical data is first requested.
Build randomly sampled transformations through one batched validation and a trusted internal construction path.
Preserve TensorMap information and registered-neighbor autograd while transforming Systems, blocks, values, and gradients.
Validate angular-momentum arguments, spherical/Cartesian component metadata, system assignments, and transformation dtype/device consistently.
Support all ten component-axis positions already recognized by the O(3) naming convention.
Document the resulting public contract and require wigners >= 0.4.0 for the structured Wigner-D API.

Compatibility decisions

This PR proposes renaming O3Transformation.is_inverted property to is_improper, which describes the actual condition (det(matrix) < 0) more precisely.

The following observable behaviors also become explicit:

matrices passed to the public O3Transformation constructor are copied, while matrix and wigner_D_matrix retain their released no-copy return behavior;
Wigner-D construction moves from object construction to first spherical use;
random_transformations accepts the model-capability dtypes torch.float32 and torch.float64;
structurally ambiguous component metadata and system identifiers are rejected before numerical transformation.

Performance evidence

The retained optimizations were isolated from correctness changes:

Wigner-D matrices are now built only when a spherical transformation or a direct Wigner-D request requires them. For 32 float64 CUDA transformations with max_angular_momentum=4, this reduced creation time by 72.9% and creation plus Cartesian-only use by 71.9%. The Wigner-D cost is just deferred as shown by the fact that, when a spherical transformation was requested immediately, total time was unchanged within 0.2%.
When random_transformations generated 32 to 4096 transformations, replacing separate validation of every 3×3 matrix with one batched validation reduced creation time by 82.7–97.9% on CPU float64 and 90.8–99.6% on CUDA.

The compared matrices, proper/improper classifications, and representative Wigner-D results were identical.

Validation

canonical Ruff formatting and lint checks passed;
121 focused O(3) cases passed, including CPU/CUDA and float32/float64 cases;
all 243 tests in the complete metatomic-torch Python suite passed;
the canonical documentation build passed with Sphinx warnings treated as errors.

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?

Recover ZYZ angles without amplifying roundoff at Euler poles and consume structured wigner_D_array results directly. Require wigners 0.4 or newer for the structured API.

Correct spherical parity, protect transformation state, defer Wigner caches, and batch trusted random construction. Preserve System and TensorMap metadata/autograd while validating component metadata, system assignments, and dtype/device contracts. Expand focused CPU/CUDA coverage for the released and corrected behavior.

Document component labels, parity, ownership, lazy Wigner construction, system assignment, and supported random dtypes. Record the observable API changes and correctness fixes in the unreleased changelog.

Luthaf

just a quick look at the API changes, I'll do a deeper pass for the implementation & tests after everything is converged

Luthaf · 2026-07-22T08:09:09Z

+    @classmethod
+    def _create_from_internal_matrix(
+        cls,
+        matrix: torch.Tensor,
+        max_angular_momentum: int,
+        *,
+        is_improper: bool,
+    ) -> "O3Transformation":
+        """Create a transformation from an internally generated matrix."""
+        transformation = cls.__new__(cls)
+        transformation._matrix = matrix
+        transformation._max_angular_momentum = max_angular_momentum
+        transformation._is_improper = is_improper
+        transformation._wigner_D_cache = None
+        return transformation


What's the point of this compared to the constructor above?

This private path is used only after random_transformations has generated and validated all matrices as one batch and already knows which one is proper/improper and It avoids repeating the public constructor’s per-matrix validation, copy, and determinant calculation.

Exported model wrappers can store Cartesian and Wigner-D matrices as registered tensor buffers, while the existing transform_tensor interface accepts Python O3Transformation objects and can not be compiled by TorchScript. Add a private tensor-only transformation path that consumes precomputed matrix batches while reusing the O(3) component, parity, routing, gradient, and metadata conventions. Add a private query for the largest spherical component rank present in TensorMap values or gradients, allowing callers to validate their serialized Wigner-D coverage. Keep the public O(3) API unchanged, and compare the new path directly with transform_tensor for proper and improper operations on CPU and CUDA.

Luthaf

The tests are very sloppy, checking minute details and drowning relevant information in a lot of lines and cases. This makes the test suite takes longer than needed for no reason, and makes it harder to understand the intent behind each test.

I would remove most of the tests that just check one error message in a very rare misuse of the API (i.e. sigma = SomeRandomClass()), and try to consolidate/clearly explain for each test why it is needed/what it is testing

Luthaf

Getting there, some of the tests can be made more useful by checking the actual full message instead of a substring

Luthaf · 2026-07-28T08:57:14Z

Thank @MichelangeloDomina and @ppegolo for your work here!

MichelangeloDomina added 3 commits July 21, 2026 16:26

fix(torch): stabilize real Wigner-D construction

a05f005

Recover ZYZ angles without amplifying roundoff at Euler poles and consume structured wigner_D_array results directly. Require wigners 0.4 or newer for the structured API.

docs(torch): define O(3) transformation contracts

7d93e17

Document component labels, parity, ownership, lazy Wigner construction, system assignment, and supported random dtypes. Record the observable API changes and correctness fixes in the unreleased changelog.

MichelangeloDomina requested a review from ppegolo July 21, 2026 15:04

Luthaf reviewed Jul 22, 2026

View reviewed changes

ppegolo reviewed Jul 22, 2026

View reviewed changes

Comment thread python/metatomic_torch/tests/o3.py Outdated

Comment thread python/metatomic_torch/tests/o3.py

MichelangeloDomina added 2 commits July 23, 2026 14:07

fix(torch): address O(3) review feedback

ca69e65

test(torch): move Wigner-D composition coverage upstream

ed6a7a8

MichelangeloDomina mentioned this pull request Jul 23, 2026

Test Wigner-D matrix composition Luthaf/wigners#21

Merged

MichelangeloDomina added 2 commits July 23, 2026 15:49

test(torch): remove stale Wigner test import

a3b3e97

MichelangeloDomina mentioned this pull request Jul 23, 2026

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

Open

3 tasks

ppegolo marked this pull request as ready for review July 24, 2026 14:25

Luthaf reviewed Jul 27, 2026

View reviewed changes

ppegolo force-pushed the fix/o3-transformations branch 8 times, most recently from 2b1a9fb to 4863f80 Compare July 27, 2026 14:12

Address review round

d8c324f

ppegolo force-pushed the fix/o3-transformations branch from 4863f80 to d8c324f Compare July 27, 2026 14:34

Luthaf reviewed Jul 27, 2026

View reviewed changes

Address second review round

3cef3f2

Luthaf approved these changes Jul 28, 2026

View reviewed changes

Luthaf merged commit 851890c into metatensor:main Jul 28, 2026
63 checks passed

Uh oh!

Conversation

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

Uh oh!

What changes

Compatibility decisions

Performance evidence

Validation

Contributor (creator of pull-request) checklist

Reviewer checklist

Uh oh!

Luthaf left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Luthaf Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

MichelangeloDomina Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Luthaf left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Luthaf left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Luthaf commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

MichelangeloDomina commented Jul 21, 2026 •

edited

Loading