/bionemo-cuequivariance

name: cuequivariance
description: Use when defining cuEquivariance groups and irreps, building equivariant tensor products, or validating segmented polynomials with NumPy.
license: Apache-2.0
metadata:
  author: Mario Geiger <mgeiger@nvidia.com>
  tags: [equivariance, group-theory, tensor-products]

cuequivariance: Groups, Irreps, and Segmented Polynomials

Purpose

`cuequivariance` (imported as `cue`) provides two core abstractions:

1. **Group theory**: `Irrep` subclasses define irreducible representations of Lie groups (SO3, O3, SU2, or custom). `Irreps` manages collections with multiplicities. 2. **Segmented polynomials**: `SegmentedTensorProduct` describes tensor contractions over segments of varying shape, linked by `Path` objects carrying Clebsch-Gordan coefficients. `SegmentedPolynomial` wraps multiple STPs into a polynomial with named inputs/outputs. Two higher-level wrappers attach group representations:

• `EquivariantPolynomial` — dense operands with `IrrepsAndLayout` metadata
• `IrDictPolynomial` — operands already split by irrep, with per-group `Irreps` metadata for the `dict[Irrep, Array]` workflow

Inputs

Required inputs: requested operation, group, input irreps and multiplicities, output irrep types, layout, batch shapes, dtype, and weight ownership. Take explicit user requirements first, then existing project definitions; state defaults for unspecified values. Derive operand sizes and output block order from the descriptor. For custom groups, also establish the selection rules, generators, and Clebsch-Gordan coefficients. Optional inputs include test rotations and a framework backend.

Prerequisites

Use Python 3.10+ with `cuequivariance` and its NumPy, SciPy, opt_einsum, SymPy, and NetworkX dependencies. Reuse the project's interpreter and lockfile. If setup is needed outside an existing project, the bundled [CPU requirements](scripts/requirements.txt) pin the source revision used by these examples; install them with that interpreter's pip. Core NumPy execution requires no GPU stack or API credentials. Check the installed version and API availability before using `IrDictPolynomial` with an older package.

Instructions

1. Check the installed version and requested backend. Install `cuequivariance` for CPU NumPy work; JAX and PyTorch execution use the separate `cuequivariance-jax` and `cuequivariance-torch` skills. Inspect available APIs before using newer features such as `IrDictPolynomial` in an older environment. 2. Establish the group, input irreps, multiplicities, layout, and batch shapes before building a descriptor. Preserve the descriptor's actual output block order; a regrouped irrep summary does not reorder its output array. 3. Prefer a built-in descriptor. For a custom group, validate parsing, group relations, selection rules, and Clebsch-Gordan covariance on CPU first. Complex mathematical representations do not imply complex CUDA support. 4. Execute a small deterministic NumPy example and check shapes and equivariance. `IrDictPolynomial` carries metadata; execute its `.polynomial` with operands split according to `.input_irreps` and `.output_irreps`, including weights first when present. 5. Convert `mul_ir` and `ir_mul` by transposing the multiplicity and representation axes **within each irrep block**. A flat reshape alone changes the represented features.

Source paths below are relative to the `cuequivariance/` project inside a source checkout. When using this skill separately, inspect the installed package with `cue.__file__` or consult https://github.com/NVIDIA/cuEquivariance. The repository is not required for the examples.

Limitations

• Core descriptors and NumPy evaluation do not provide JAX or PyTorch execution kernels; use the corresponding frontend for framework execution and gradients.
• Custom-group mathematics must be checked independently. A valid complex representation does not establish support in a particular CUDA kernel.
• `IrDictPolynomial` describes split operands; it does not convert high-level dictionaries into the ordered arrays accepted by its underlying polynomial. Preserve the descriptor's metadata and layout when assembling those arrays.

Troubleshooting

| Symptom | Check and correction | | --- | --- | | Import fails or an API is absent | Check the interpreter, installed version, and module location. Use the project's compatible dependencies or the bundled requirements when setting up a standalone environment. | | Operand size or output shape is wrong | Compare each array with the descriptor's input/output dimensions; include the weights operand first when present. | | Values disagree after a layout conversion | Transpose multiplicity and representation axes within each irrep block before flattening; keep output data in the descriptor's block order. | | Custom-group equivariance fails | Check selection rules, generator relations, and CG covariance on a small NumPy example before changing execution backends. |

Examples

The examples below cover custom groups, descriptors, layouts, and NumPy execution. Some snippets extend variables from the preceding snippet; the final complete example can be run independently.

Defining a custom group

Subclass `cue.Irrep` (a frozen dataclass) and implement:

from __future__ import annotations
import dataclasses
import re
from typing import Iterator
import numpy as np
import cuequivariance as cue

@dataclasses.dataclass(frozen=True)
class Z2(cue.Irrep):
    odd: bool  # dataclass field -- required for correct __eq__ and __hash__

    # No __init__ needed -- @dataclass(frozen=True) generates it: Z2(odd=True)

    @classmethod
    def regexp_pattern(cls) -> re.Pattern:
        # Pattern whose first group is passed to from_string
        return re.compile(r"(odd|even)")

    @classmethod
    def from_string(cls, string: str) -> Z2:
        if string not in ("odd", "even"):
            raise ValueError(f"Invalid Z2 irrep: {string!r}")
        retur