O(3)-symmetrized models

The metatomic.torch.o3.SymmetrizedModel class wraps an existing 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.

Models outputs

metatomic.torch.o3.SymmetrizedModel adds extra outputs to the model, computing the equivariance variance or squared character-projection contributions of the model response.

output name	Result
<name>	O(3) average of the wrapped model’s <name> output
o3::variance::<name>	component-averaged equivariance variance of <name>
o3::character_projection::<name>	unnormalized squared character-projection contributions of <name>

The above outputs are added for every output of the wrapped model, including variants (such as energy/pbe) and custom outputs (such as custom::feature::node). For example, o3::variance::energy/pbe would compute the equivariance variance of the energy/pbe output.

Quadrature

The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced in-plane rotations, and both parities: 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.

Average and variance

For an input \(x\), an O(3) operation \(g\), and the target representation \(\rho_\alpha\), define the response transformed back to the input frame as

\[z_\alpha(g;x) = \rho_\alpha(g^{-1}) f(gx).\]

The ordinary result is the normalized Haar average

\[\Pi_\alpha(f,x) = \int_{\mathrm{O}(3)} z_\alpha(g;x)\,\mathrm{d}\mu(g).\]

For a TensorMap block with component multiplicity \(d\), the corresponding variance output contains

\[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.

The meaning of the variance depends on the structure of the wrapped model’s output. When its blocks carry recognized component labels (o3_mu-style spherical or xyz-style Cartesian axes, see the Conventions for O(3) representations documentation), each output is rotated back to the input frame first, and the variance measures the breaking of equivariance. Outputs without such components cannot be rotated back: their responses are compared as-is across the quadrature, so their variance measures the deviation from invariance only. An equivariant but unlabelled output — for example an internal equivariant feature vector — can thus report a large variance even when it transforms correctly.

Variance metadata

The o3::variance::<name> outputs produced by metatomic.torch.o3.SymmetrizedModel have the following metadata structure:

Metadata for "o3::variance::<name>"

Metadata	Names	Description
keys	[<keys...>, "o3_lambda", "o3_sigma"]	The keys are the same as the original <name> output, with "o3_lambda" and "o3_sigma" dimensions added if they are not already present, and the _ = 0 dummy key removed if present. The "o3_lambda" dimension contains the angular momentum of each block, and the "o3_sigma" dimension contains its parity under inversion.
samples	[<samples...>]	the samples are the same as the original <name> output.
components		Since the variance is computed over the o3_mu components of each block, the resulting TensorMap does not have any component axes.
properties	[<properties...>]	the properties are the same as the original <name> output.

When computing the variance of standard quantities, the data is first converted to spherical representation as follow:

• Scalar quantities (such as energy, charge, etc.) gain a o3_lambda=0, o3_sigma=1 key, as well as an o3_mu component of size one.

• Cartesian vector quantities (such as non-conservative force) gain a o3_lambda=1, o3_sigma=1 key, and their xyz components are replaced by an o3_mu component of size three.

• Cartesian rank-2 tensor quantities (such as non-conservative stress) gain o3_lambda=0, o3_sigma=1, o3_lambda=1, o3_sigma=-1, and o3_lambda=2, o3_sigma=1 keys, and their xyz_1 and xyz_2 components are replaced by an o3_mu component of size one, three, and five, respectively.

• Already-spherical outputs retain their o3_lambda and o3_sigma keys and o3_mu components, and other keys are preserved.

• Custom outputs with components that do not match the convention are not supported, and will raise an error when the variance is requested.

Character projections

Character projections analyze the direct response \(u(g;x)=f(gx)\), rather than the back-transformed response used for averaging. For the character sector \(\beta=(\lambda,\sigma)\) with \(d_\beta=2\lambda+1\), the squared projection norm is

\[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. Any other pre-existing component axes are retained; summing over them gives the complete component norm in the equation above.

Character projections metadata

The o3::character_projection::<name> outputs produced by metatomic.torch.o3.SymmetrizedModel have the following metadata structure:

Metadata for "o3::variance::<name>"

Metadata	Names	Description
keys	[<keys...>, "o3_lambda", "o3_sigma", "chi_lambda", "chi_sigma"]	The keys are the same as the original <name> output, with "o3_lambda" and "o3_sigma" dimensions added if they are not already present, the _ = 0 dummy key removed if present, and "chi_lambda" and "chi_sigma" keys added.
samples	[<samples...>]	the samples are the same as the original <name> output.
components	[<components...>, "o3_mu"]	The components are the same as the original <name> output, with an additional "o3_mu" component of size \(2\lambda+1\) added for each block. This "o3_mu" component replaces existing "xyz" or "o3_mu" components, if present.
properties	[<properties...>]	the properties are the same as the original <name> output.

When computing character projections, the data is first converted to spherical representation like for the variance, with the exception that arbitrary pre-existing component axes are allowed and will not raise an exception. These are retained as-is in the output.

API reference

class metatomic.torch.o3.SymmetrizedModel(model: ModelInterface, *, max_angular_momentum_target: int, max_angular_momentum_input: int = 0, max_angular_momentum_character: int | None = None, max_angular_momentum_grid: int | None = None, batch_size: int = 32)

Wrap a model with finite-quadrature O(3) averaging and equivariance diagnostics.

Requesting an output declared by the wrapped model returns its O(3) average, evaluated over rotated and inverted copies of the input and transformed back to the input frame. Requests named o3::variance::<name> return the component-averaged equivariance variance of the <name> output and, when max_angular_momentum_character is set, o3::character_projection::<name> requests return its unnormalized squared character-projection contributions. Outputs whose blocks carry no recognized component labels are not rotated back, so their variance measures the deviation from invariance only. The definition of these quantities, their TensorMap representation, and convergence guidance for the quadrature are documented in O(3)-symmetrized models.

Requests for explicit TensorBlock gradients are rejected. When an input requires gradients, differentiating an averaged result through PyTorch autograd retains the source-model activations from all quadrature batches; batch_size does not bound their total size. Use torch.inference_mode() or torch.no_grad() when derivatives are not required.

Parameters:

• model (ModelInterface) – underlying ModelInterface. The wrap() method obtains this module from AtomisticModel.module.

• max_angular_momentum_target (int) – maximum angular momentum that can be transformed back to the input frame when an average or variance of an already-spherical output is requested. Cartesian outputs and character-only requests are not limited by this value.

• max_angular_momentum_input (int) – maximum angular momentum that can be rotated in already-spherical custom System data. The default of zero still allows Cartesian custom inputs. The ModelOutput declarations returned by a model’s requested_inputs() do not specify which angular momenta may occur in the corresponding TensorMaps, so this limit must be supplied before export for all required Wigner-D matrices to be serialized.

• max_angular_momentum_character (int | None) – maximum angular momentum included in character projections. None disables character-projection outputs; zero enables the scalar (o3_lambda = 0) contribution only.

• max_angular_momentum_grid (int | None) – quadrature integration degree. If None, use the larger of 2 * max_angular_momentum_target + 1 and 2 * max_angular_momentum_character when character projections are enabled. An explicit value must be non-negative and no larger than the highest available Lebedev order, 131; a value below 2 * max_angular_momentum_character is rejected.

• batch_size (int) – positive number of transformed systems evaluated in one call to model. The default is 32.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

static wrap(model: AtomisticModel, *, max_angular_momentum_target: int | None = None, max_angular_momentum_input: int | None = None, max_angular_momentum_character: int | None = None, max_angular_momentum_grid: int | None = None, batch_size: int = 32) → AtomisticModel

Wrap an exported model with O(3) averaging and diagnostics.

The returned model retains every output declared by model under its original name. Requesting such an output evaluates its O(3) average. Additional outputs named o3::variance::<name> provide the component-averaged equivariance variance. If max_angular_momentum_character is set, o3::character_projection::<name> outputs provide squared character projections through that angular momentum.

The original metadata, requested inputs, neighbor lists, and compatible capabilities are preserved.

Constructing a wrapper requires SciPy 1.15 or newer for its Lebedev quadrature. SciPy is not required to evaluate a wrapper that has already been saved.

Parameters:

• model (AtomisticModel) – the AtomisticModel to wrap

• max_angular_momentum_target (int | None) – maximum angular momentum accepted in already-spherical model outputs requested for averaging or variance. When None, it is guessed as the largest angular momentum of the standard quantities declared by model; non-standard outputs are skipped, and an explicit value is required if model declares outputs but none of them is a standard quantity.

• max_angular_momentum_input (int | None) – maximum angular momentum accepted in custom System data. When None, it is guessed the same way from the quantities in model.requested_inputs().

• max_angular_momentum_character (int | None) – maximum angular momentum in character projections, or None to disable them

• max_angular_momentum_grid (int | None) – quadrature integration degree, selected automatically when None

• batch_size (int) – number of transformed Systems evaluated in one model call

Return type:

AtomisticModel

requested_neighbor_lists() → List[NeighborListOptions]

Return the neighbor lists requested by the wrapped model.

Return type:

List[NeighborListOptions]

requested_inputs() → Dict[str, ModelOutput]

Return the custom System data requested by the wrapped model.

Return type:

Dict[str, ModelOutput]

forward(systems: List[System], outputs: Dict[str, ModelOutput], selected_atoms: Labels | None) → Dict[str, TensorMap]

Evaluate the requested O(3) averages and diagnostics.

Parameters:

• systems (List[System])

• outputs (Dict[str, ModelOutput])

• selected_atoms (Labels | None)

Return type:

Dict[str, TensorMap]