e3nn

Warnings

• breaking e3nn is under active development. Breaking changes are introduced in minor version increments (e.g., from 0.x.x to 0.y.x). Always check the CHANGELOG when updating.
  Fix: Consult the official CHANGELOG (https://github.com/e3nn/e3nn/blob/main/CHANGELOG.md) before upgrading to new minor versions to adapt your code for breaking API changes.
• breaking Normalization constants for `o3.TensorProduct` and `o3.Linear` were changed in version 0.4.0. Models with inhomogeneous multiplicities may be affected.
  Fix: If upgrading from pre-0.4.0, existing models using these operations might produce different results. Retrain or re-evaluate models if using custom normalization or assuming previous constant values.
• gotcha The output of an e3nn model (or any equivariant operation) must always have equal or higher symmetry than its input. Designing models that violate this principle can lead to unexpected behavior or errors.
  Fix: Carefully design your network architecture, especially the irreducible representations (Irreps) for intermediate and output layers, to ensure symmetry constraints are respected. Refer to the 'Simple Tasks and Symmetry' tutorial for examples.
• deprecated Older tutorials and examples (especially those for e3nn versions < 0.2) may be outdated and incompatible with current API versions. Always refer to the latest User Guide.
  Fix: Prioritize the official e3nn documentation (e3nn.org) and the e3nn GitHub repository for up-to-date examples and best practices. Avoid relying on older, unmaintained tutorial repositories.
• breaking Python 3.6 support was dropped in earlier versions. The current library requires Python 3.8 or higher.
  Fix: Ensure your development environment uses Python 3.8 or a newer compatible version.

Install

• pip install --upgrade e3nn Install e3nn (requires PyTorch pre-installed)

Imports

• o3

  from e3nn import o3

  The 'o3' submodule contains core O(3) irreducible representations and operations.
• Irreps

  from e3nn.o3 import Irreps

  While 'from e3nn import o3' is common, specific classes like Irreps are found within the 'o3' submodule.

Quickstart

This example demonstrates how to define irreducible representations (Irreps), apply a linear transformation, and compute a tensor product using e3nn's o3 submodule. It showcases basic operations for building equivariant neural networks.

import torch
from e3nn import o3

# Create a random array made of scalar (0e) and a vector (1o) representations
# '0e' denotes a scalar with even parity, '1o' denotes a vector with odd parity.
irreps_in = o3.Irreps("0e + 1o")
x = irreps_in.randn(-1)

# Define output representations and apply a linear layer
irreps_out = o3.Irreps("2x0e + 2x1o") # Two scalars and two vectors
linear = o3.Linear(irreps_in=irreps_in, irreps_out=irreps_out)
y = linear(x)

# Compute a tensor product of the input with itself
tp = o3.FullTensorProduct(irreps_in1=irreps_in, irreps_in2=irreps_in)
z = tp(x, x)

print(f"Input (x) shape: {x.shape}, irreps: {irreps_in}")
print(f"Linear output (y) shape: {y.shape}, irreps: {irreps_out}")
print(f"Tensor product output (z) shape: {z.shape}, irreps: {tp.irreps_out}")

# For performance, modules can optionally be compiled with torch.compile
# tp_pt2 = torch.compile(tp, fullgraph=True)
# z_pt2 = tp_pt2(x, x)
# torch.testing.assert_close(z, z_pt2)