Equivariant layers in three dimensions

[laserkelvin]

PhysicsNeMo Pull Request

Description

This PR extends #1367 by adding 3D equivariance preserving linear transformations: SO3LinearGrid in isolation, and composed together with nonlinearities with SO3ConvolutionBlock. Accompanying these changes are also routines to compute 3D rotations, both as part of testing as well as for the actual workflow itself.

The full intended workflow, now with these layers is given some set of edge-wise spherical harmonic embeddings with shape [num_edges, lmax + 1, mmax + 1, 2, num_channels], the Wigner rotation aligns the embeddings along the direction of their edges, reduces an SO3 problem into an SO2, i.e. the hard part is done in SO2¹:

rotate(x) |> so2_convolution() |> inv_rotate() |> so3_block()

The end result should be transformed edge embeddings, and an aggregation operation will get you the updated node embeddings.

Rotations, through the EdgeRotation class, can be performed at arbitrary precision as needed: for example, when training at lower precision, we can choose to perform rotations at higher precision to maintain information. Conversely, maybe the rotation could be done at lower precision without too much degradation.

One potential gotcha is that the Wigner matrix is cached: the idea is that, an architecture will comprise multiple SO2/SO3 blocks, and the rotation matrix only needs to be computed once per batch. I've left notes in the docstrings to remind users to manually clear the cache after backprop/optimization steps are done as to not have dangling graphs. I've thought about writing a hash-based caching system (i.e. hash the batch, and reuse the matrix only if the hash matches), but part of my design philosophy was to minimize branching withtorch.compile in mind.

The PR also adds unit tests that more or less mirror the SO2 ones (also updating them to homogenize), with a little more coverage due to the complexity of 3D operations.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

None

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

Footnotes

1. This is to avoid the expensive tensor product ↩

[laserkelvin]

/blossom-ci

[greptile-apps]

Greptile Overview

Greptile Summary

This PR extends the symmetry module with SO(3) equivariant layers, building on the existing SO(2) foundation from PR #1367. The implementation adds three key components:

Core implementations:

• SO3LinearGrid: Degree-wise linear transformations preserving SO(3) equivariance
• SO3ConvolutionBlock: Feed-forward block combining SO3Linear layers with gated activation
• EdgeRotation & Wigner D-matrix computation: On-the-fly rotation matrix generation with caching

Key strengths:

• Comprehensive test coverage with regression tests against SymPy reference implementations
• Multi-precision support (float16/bfloat16/float32/float64) with appropriate tolerance handling
• Extensive documentation following repository standards with detailed docstrings, examples, and math notation
• Efficient vectorized implementations using einsum operations
• Proper masking for invalid (l,m) positions and numerical stability features

Architecture:
The workflow aligns edge-wise spherical harmonic embeddings using Wigner rotations, reducing SO(3) problems to SO(2) before applying convolutions, then transforming back through SO(3) blocks.

The code is well-structured, follows coding standards (MOD-000a placement in experimental, MOD-001 Module inheritance, MOD-003 documentation), and includes appropriate validation. The previous review comments about docstring inconsistencies and shape validation have been noted but not yet addressed in this iteration.

Important Files Changed

Filename	Overview
physicsnemo/experimental/nn/symmetry/wigner.py	Implements Wigner D-matrix computation for 3D rotations with comprehensive caching, on-the-fly J-matrix computation, and safe gradient handling
physicsnemo/experimental/nn/symmetry/so3_linear.py	SO(3) equivariant linear layer using degree-wise weight sharing with proper masking for invalid (l,m) positions
physicsnemo/experimental/nn/symmetry/so3_block.py	SO(3) block combining two SO3Linear layers with GateActivation, using scalar MLP for gate generation
test/experimental/nn/symmetry/test_wigner.py	Comprehensive tests for Wigner D-matrices with regression tests against sympy reference implementations
test/experimental/nn/symmetry/test_so3_linear.py	Tests for SO3LinearGrid covering shapes, masking, equivariance, and gradient flow

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[laserkelvin]

/blossom-ci

[laserkelvin]

/blossom-ci

[laserkelvin]

/blossom-ci

[laserkelvin]

/blossom-ci

[laserkelvin]

/blossom-ci

[laserkelvin]

@peterdsharpe @CharlelieLrt ready for review (finally)!

[laserkelvin]

/blossom-ci

[laserkelvin]

/blossom-ci

[peterdsharpe]

Looks good overall! Left some (optional) comments, and we might want to revisit a few things when this module graduates to nn (e.g., can SO3 linear inherit from Linear, so that we don't have to write our own initialization? Can we generalize the SO3 and SO2 layers to be dimensionally-generic?), but no need to hold this up.

Nice work, looking forward to giving these a spin!

[laserkelvin]

/blossom-ci

[laserkelvin]

e.g., can SO3 linear inherit from Linear, so that we don't have to write our own initialization? Can we generalize the SO3 and SO2 layers to be dimensionally-generic?

Yeah let's definitely touch base when it moves out of experimental: I need to play around with composing the layers and running it on real data first, and once I've dogfood'd it a little bit might have a more informed decision to make about these things. If you get the chance, too, let me know if there are pain points