{"id":17524,"plugin_id":"plugins_6a76572d8f8081918362aa7ff90947fb","kind":"skill","collection_source":null,"comparison_source":null,"observed_at":"2026-09-30T23:14:15.012Z","digest":"217aaeb248be96220cfdcc4bc7c533bd8088393ae80e0760b51fb08728a717b3","against":null,"payload":{"name":"cuequivariance","description":"Define custom groups (Irrep subclasses), build segmented tensor products with CG coefficients, create equivariant polynomials and IrDictPolynomials, and use built-in descriptors (linear, tensor products, spherical harmonics). Use when working with cuequivariance group theory, irreps, or segmented polynomials.","included_files":[],"skill_md_contents":"---\nname: cuequivariance\ndescription: Define custom groups (Irrep subclasses), build segmented tensor products with CG coefficients, create equivariant polynomials and IrDictPolynomials, and use built-in descriptors (linear, tensor products, spherical harmonics). Use when working with cuequivariance group theory, irreps, or segmented polynomials.\n---\n\n# cuequivariance: Groups, Irreps, and Segmented Polynomials\n\n## Overview\n\n`cuequivariance` (imported as `cue`) provides two core abstractions:\n\n1. **Group theory**: `Irrep` subclasses define irreducible representations of Lie groups (SO3, O3, SU2, or custom). `Irreps` manages collections with multiplicities.\n2. **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:\n   - `EquivariantPolynomial` — dense operands with `IrrepsAndLayout` metadata\n   - `IrDictPolynomial` — operands already split by irrep, with per-group `Irreps` metadata for the `dict[Irrep, Array]` workflow\n\n## Defining a custom group\n\nSubclass `cue.Irrep` (a frozen dataclass) and implement:\n\n```python\nfrom __future__ import annotations\nimport dataclasses\nimport re\nfrom typing import Iterator\nimport numpy as np\nimport cuequivariance as cue\n\n@dataclasses.dataclass(frozen=True)\nclass Z2(cue.Irrep):\n    odd: bool  # dataclass field -- required for correct __eq__ and __hash__\n\n    # No __init__ needed -- @dataclass(frozen=True) generates it: Z2(odd=True)\n\n    @classmethod\n    def regexp_pattern(cls) -> re.Pattern:\n        # Pattern whose first group is passed to from_string\n        return re.compile(r\"(odd|even)\")\n\n    @classmethod\n    def from_string(cls, string: str) -> Z2:\n        return cls(odd=string == \"odd\")\n\n    def __repr__(rep: Z2) -> str:\n        return \"odd\" if rep.odd else \"even\"\n\n    def __mul__(rep1: Z2, rep2: Z2) -> Iterator[Z2]:\n        # Selection rule: which irreps appear in the tensor product rep1 x rep2\n        return [Z2(odd=rep1.odd ^ rep2.odd)]\n\n    @classmethod\n    def clebsch_gordan(cls, rep1: Z2, rep2: Z2, rep3: Z2) -> np.ndarray:\n        # Shape: (num_paths, rep1.dim, rep2.dim, rep3.dim)\n        if rep3 in rep1 * rep2:\n            return np.array([[[[1]]]])\n        else:\n            return np.zeros((0, 1, 1, 1))\n\n    @property\n    def dim(rep: Z2) -> int:\n        return 1\n\n    def __lt__(rep1: Z2, rep2: Z2) -> bool:\n        # Ordering for sorting; dimension is compared first by the base class\n        return rep1.odd < rep2.odd\n\n    @classmethod\n    def iterator(cls) -> Iterator[Z2]:\n        # Must yield trivial irrep first\n        for odd in [False, True]:\n            yield Z2(odd=odd)\n\n    def discrete_generators(rep: Z2) -> np.ndarray:\n        # Shape: (num_generators, dim, dim)\n        if rep.odd:\n            return -np.ones((1, 1, 1))\n        else:\n            return np.ones((1, 1, 1))\n\n    def continuous_generators(rep: Z2) -> np.ndarray:\n        # Shape: (lie_dim, dim, dim) -- Z2 is discrete, so lie_dim=0\n        return np.zeros((0, rep.dim, rep.dim))\n\n    def algebra(self) -> np.ndarray:\n        # Shape: (lie_dim, lie_dim, lie_dim) -- structure constants [X_i, X_j] = A_ijk X_k\n        return np.zeros((0, 0, 0))\n\n\n# Usage:\nirreps = cue.Irreps(Z2, \"3x odd + 2x even\")  # dim=5\n```\n\n### Required methods summary\n\n| Method | Returns | Purpose |\n|--------|---------|---------|\n| `regexp_pattern()` | `re.Pattern` | Parse string like `\"1\"`, `\"0e\"`, `\"odd\"` |\n| `from_string(s)` | `Irrep` | Construct from matched string |\n| `__repr__` | `str` | Canonical string form |\n| `__mul__(a, b)` | `Iterator[Irrep]` | Selection rule for tensor product |\n| `clebsch_gordan(a, b, c)` | `ndarray (n, d1, d2, d3)` | CG coefficients |\n| `dim` (property) | `int` | Dimension of representation |\n| `__lt__(a, b)` | `bool` | Ordering (dimension first, then custom) |\n| `iterator()` | `Iterator[Irrep]` | All irreps, trivial first |\n| `continuous_generators()` | `ndarray (lie_dim, dim, dim)` | Lie algebra generators |\n| `discrete_generators()` | `ndarray (n, dim, dim)` | Finite symmetry generators |\n| `algebra()` | `ndarray (lie_dim, lie_dim, lie_dim)` | Structure constants |\n\n### Built-in groups\n\n- **`cue.SO3(l)`**: 3D rotations. `l` is a non-negative integer. `dim = 2l+1`. String: `\"0\"`, `\"1\"`, `\"2\"`.\n- **`cue.O3(l, p)`**: 3D rotations + parity. `p=+1` (even) or `p=-1` (odd). String: `\"0e\"`, `\"1o\"`, `\"2e\"`.\n- **`cue.SU2(j)`**: Spin group. `j` is a non-negative half-integer. String: `\"0\"`, `\"1/2\"`, `\"1\"`.\n\n## Irreps and layout\n\n```python\nirreps = cue.Irreps(\"SO3\", \"16x0 + 4x1 + 2x2\")  # 16 scalars, 4 vectors, 2 rank-2\nirreps.dim   # 16*1 + 4*3 + 2*5 = 38\n\nfor mul, ir in irreps:\n    print(mul, ir, ir.dim)  # 16 0 1, then 4 1 3, then 2 2 5\n```\n\n`IrrepsLayout` controls memory order within each `(mul, ir)` block:\n\n- `cue.ir_mul`: data ordered as `(ir.dim, mul)` — **used by all descriptors and ir_dict internally**\n- `cue.mul_ir`: data ordered as `(mul, ir.dim)` — **used by nnx `dict[Irrep, Array]` and PyTorch**\n\n`IrrepsAndLayout` combines irreps with a layout into a `Rep`:\n\n```python\nrep = cue.IrrepsAndLayout(cue.Irreps(\"SO3\", \"4x0 + 2x1\"), cue.ir_mul)\nrep.dim  # 10\n```\n\n## Building a SegmentedTensorProduct from scratch\n\nThe subscripts string uses Einstein notation. Operands are comma-separated, coefficient modes follow `+`.\n\n```python\n# Matrix-vector multiply: y_i = sum_j M_ij * x_j\nd = cue.SegmentedTensorProduct.from_subscripts(\"ij,j,i\")\nd.add_segment(0, (3, 4))  # operand 0: matrix segment of shape (3, 4)\nd.add_segment(1, (4,))     # operand 1: vector of size 4\nd.add_segment(2, (3,))     # operand 2: output vector of size 3\nd.add_path(0, 0, 0, c=1.0) # link segments 0,0,0 with coefficient=1.0\n\npoly = cue.SegmentedPolynomial.eval_last_operand(d)  # last operand becomes output\n[y] = poly(M_flat, x)  # numpy evaluation\n```\n\n### Multi-segment STP (how descriptors work internally)\n\nDescriptors build STPs with multiple segments per operand. Each segment corresponds to an irrep block:\n\n```python\n# Linear equivariant map: output[iv] = sum_u weight[uv] * input[iu]\nd = cue.SegmentedTensorProduct.from_subscripts(\"uv,iu,iv\")\n\n# Segment for l=1: ir_dim=3, mul_in=2, mul_out=5\ns_in_0 = d.add_segment(1, (3, 2))    # input block\ns_out_0 = d.add_segment(2, (3, 5))   # output block\nd.add_path((2, 5), s_in_0, s_out_0, c=1.0)\n\n# Segment for l=0: ir_dim=1, mul_in=4, mul_out=3\ns_in_1 = d.add_segment(1, (1, 4))\ns_out_1 = d.add_segment(2, (1, 3))\nd.add_path((4, 3), s_in_1, s_out_1, c=1.0)\n```\n\n### Weights operand\n\nFor weighted tensor products (subscript starting with `uvw` or `uv`), the first operand is always weights. The weight segment shape is `(mul_1, mul_2, ...)` matching the multiplicity modes. The weights operand gets `new_scalars()` irreps since weights are invariant.\n\n### CG coefficients as path coefficients\n\n```python\nd = cue.SegmentedTensorProduct.from_subscripts(\"uvw,iu,jv,kw+ijk\")\n# For each pair of input irreps and each output irrep in the selection rule:\nfor cg in cue.clebsch_gordan(ir1, ir2, ir3):\n    # cg has shape (ir1.dim, ir2.dim, ir3.dim)\n    d.add_path((mul1, mul2, mul3), seg_in1, seg_in2, seg_out, c=cg)\n```\n\n## Descriptors\n\nAll descriptors come in two variants:\n\n- **Original** — returns `EquivariantPolynomial` with dense operands\n- **`_ir_dict`** — returns `IrDictPolynomial` with operands already split by irrep\n\n### EquivariantPolynomial descriptors\n\n```python\n# Fully connected tensor product (all input-output irrep combinations)\ne = cue.descriptors.fully_connected_tensor_product(\n    16 * cue.Irreps(\"SO3\", \"0 + 1 + 2\"),\n    16 * cue.Irreps(\"SO3\", \"0 + 1 + 2\"),\n    16 * cue.Irreps(\"SO3\", \"0 + 1 + 2\"),\n)\n\n# Channelwise tensor product (same-channel only, sparse)\ne = cue.descriptors.channelwise_tensor_product(\n    64 * cue.Irreps(\"SO3\", \"0 + 1\"), cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"), simplify_irreps3=True,\n)\n\n# Full (weightless) tensor product\ne = cue.descriptors.full_tensor_product(\n    cue.Irreps(\"SO3\", \"2x0 + 1x1\"), cue.Irreps(\"SO3\", \"0 + 1\"),\n)\n\n# Elementwise tensor product (paired channels)\ne = cue.descriptors.elementwise_tensor_product(\n    cue.Irreps(\"SO3\", \"4x0 + 4x1\"), cue.Irreps(\"SO3\", \"4x0 + 4x1\"),\n)\n\n# Linear equivariant map (weight x input)\ne = cue.descriptors.linear(\n    cue.Irreps(\"SO3\", \"4x0 + 2x1\"),\n    cue.Irreps(\"SO3\", \"3x0 + 5x1\"),\n)\n\n# Spherical harmonics\ne = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2, 3])\n\n# Symmetric contraction (MACE-style)\ne = cue.descriptors.symmetric_contraction(\n    64 * cue.Irreps(\"SO3\", \"0 + 1 + 2\"),\n    64 * cue.Irreps(\"SO3\", \"0 + 1\"),\n    (1, 2, 3),\n)\n```\n\n### IrDictPolynomial descriptors\n\nEach `_ir_dict` variant returns an `IrDictPolynomial` whose polynomial is already split by irrep. The `input_irreps` and `output_irreps` tuples describe the operand groups.\n\n```python\n# Channelwise tensor product\ndesc = cue.descriptors.channelwise_tensor_product_ir_dict(\n    64 * cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n)\n# desc.polynomial       — SegmentedPolynomial, already split by irrep\n# desc.input_irreps     — (weight_irreps, irreps1, irreps2)\n# desc.output_irreps    — (irreps_out,)\n\n# Fully connected tensor product\ndesc = cue.descriptors.fully_connected_tensor_product_ir_dict(irreps1, irreps2, irreps3)\n\n# Full (weightless) tensor product\ndesc = cue.descriptors.full_tensor_product_ir_dict(irreps1, irreps2)\n\n# Elementwise tensor product\ndesc = cue.descriptors.elementwise_tensor_product_ir_dict(irreps1, irreps2)\n\n# Linear\ndesc = cue.descriptors.linear_ir_dict(irreps_in, irreps_out)\n\n# Spherical harmonics\ndesc = cue.descriptors.spherical_harmonics_ir_dict(cue.O3(1, -1), [0, 1, 2, 3])\n\n# Symmetric contraction\ndesc = cue.descriptors.symmetric_contraction_ir_dict(irreps_in, irreps_out, (1, 2, 3))\n```\n\n### IrDictPolynomial\n\n`IrDictPolynomial` pairs a `SegmentedPolynomial` (already split by irrep) with the `Irreps` that describe each operand group.\n\n```python\ndesc = cue.descriptors.channelwise_tensor_product_ir_dict(\n    32 * cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n)\n\ndesc.polynomial       # SegmentedPolynomial — each operand is one (mul, ir) block\ndesc.input_irreps     # (weight_irreps, irreps1, irreps2)\ndesc.output_irreps    # (irreps_out,)\n\n# Scale coefficients\nscaled_poly = desc.polynomial * 0.5\n\n# Access individual operand info\nfor i, op in enumerate(desc.polynomial.inputs):\n    print(f\"Input {i}: size={op.size}, num_segments={op.num_segments}\")\n```\n\nContract: for each `(mul, ir)` block in `input_irreps` / `output_irreps`, the corresponding polynomial operand has size `mul * ir.dim`.\n\n### split_polynomial_by_irreps\n\nThe low-level function underlying `_ir_dict` descriptors. Splits one polynomial operand at irrep boundaries:\n\n```python\npoly = e.polynomial  # from an EquivariantPolynomial\npoly = cue.split_polynomial_by_irreps(poly, 2, irreps_sh)   # split input 2\npoly = cue.split_polynomial_by_irreps(poly, 1, irreps_in)   # split input 1\npoly = cue.split_polynomial_by_irreps(poly, -1, irreps_out) # split output\n```\n\n### EquivariantPolynomial key methods\n\n```python\ne.inputs     # tuple of Rep (group representations for each input)\ne.outputs    # tuple of Rep\ne.polynomial # the underlying SegmentedPolynomial\n\n# Numpy evaluation\n[out] = e(weights, input1, input2)\n\n# Preparing for uniform_1d execution (see cuequivariance_jax SKILL.md)\ne_ready = e.squeeze_modes().flatten_coefficient_modes()\n\n# Split an operand into per-irrep pieces (for ir_dict interface)\ne_split = e.split_operand_by_irrep(1).split_operand_by_irrep(-1)\n\n# Scale all coefficients\ne_scaled = e * 0.5\n\n# Fuse compatible STPs\ne_fused = e.fuse_stps()\n```\n\n### normalize_paths_for_operand\n\nCalled internally by descriptors. Normalizes path coefficients so that a random input produces unit-variance output for the specified operand. Critical for numerical stability.\n\n## SegmentedPolynomial structure\n\n```python\npoly = e.polynomial\npoly.num_inputs    # number of input operands\npoly.num_outputs   # number of output operands\npoly.inputs        # tuple of SegmentedOperand\npoly.outputs       # tuple of SegmentedOperand\npoly.operations    # tuple of (Operation, SegmentedTensorProduct)\n\n# Each operation maps buffers to STP operands\nfor op, stp in poly.operations:\n    print(op.buffers)  # e.g., (0, 1, 2) means inputs[0], inputs[1] -> outputs[0]\n    print(stp.subscripts)\n```\n\n### SegmentedOperand\n\n```python\noperand = poly.inputs[0]\noperand.num_segments     # how many segments\noperand.segments         # tuple of shape tuples, e.g., ((3, 4), (1, 2))\noperand.size             # total flattened size (sum of products of segment shapes)\noperand.ndim             # number of dimensions per segment\noperand.all_same_segment_shape()  # True if all segments have identical shape\noperand.segment_shape    # the common shape (only if all_same_segment_shape)\n```\n\n## Custom equivariant polynomial from scratch\n\n```python\nimport numpy as np\nimport cuequivariance as cue\n\n# Build a fully-connected SO3(1)xSO3(1)->SO3(0) tensor product manually\ncg = cue.clebsch_gordan(cue.SO3(1), cue.SO3(1), cue.SO3(0))  # shape (1, 3, 3, 1)\n\nd = cue.SegmentedTensorProduct.from_subscripts(\"uvw,iu,jv,kw+ijk\")\nd.add_segment(1, (3, 4))   # input1: 4x SO3(1), shape=(ir_dim, mul)\nd.add_segment(2, (3, 4))   # input2: 4x SO3(1)\nd.add_segment(3, (1, 16))  # output: 16x SO3(0) (4*4 fully connected)\n\nfor c in cg:\n    d.add_path((4, 4, 16), 0, 0, 0, c=c)\n\nd = d.normalize_paths_for_operand(-1)\n\npoly = cue.SegmentedPolynomial.eval_last_operand(d)\nep = cue.EquivariantPolynomial(\n    [\n        cue.IrrepsAndLayout(cue.Irreps(\"SO3\", \"4x1\").new_scalars(d.operands[0].size), cue.ir_mul),\n        cue.IrrepsAndLayout(cue.Irreps(\"SO3\", \"4x1\"), cue.ir_mul),\n        cue.IrrepsAndLayout(cue.Irreps(\"SO3\", \"4x1\"), cue.ir_mul),\n    ],\n    [cue.IrrepsAndLayout(cue.Irreps(\"SO3\", \"16x0\"), cue.ir_mul)],\n    poly,\n)\n\n# Numpy evaluation\nw = np.random.randn(ep.inputs[0].dim)\nx = np.random.randn(ep.inputs[1].dim)\ny = np.random.randn(ep.inputs[2].dim)\n[out] = ep(w, x, y)\n```\n\n## Key file locations\n\n| Component | Path |\n|-----------|------|\n| `Irrep` base class | `cuequivariance/group_theory/representations/irrep.py` |\n| `Rep` base class | `cuequivariance/group_theory/representations/rep.py` |\n| `SO3` | `cuequivariance/group_theory/representations/irrep_so3.py` |\n| `O3` | `cuequivariance/group_theory/representations/irrep_o3.py` |\n| `SU2` | `cuequivariance/group_theory/representations/irrep_su2.py` |\n| `Irreps` | `cuequivariance/group_theory/irreps_array/irreps.py` |\n| `IrrepsLayout` | `cuequivariance/group_theory/irreps_array/irreps_layout.py` |\n| `IrrepsAndLayout` | `cuequivariance/group_theory/irreps_array/irreps_and_layout.py` |\n| `SegmentedTensorProduct` | `cuequivariance/segmented_polynomials/segmented_tensor_product.py` |\n| `SegmentedPolynomial` | `cuequivariance/segmented_polynomials/segmented_polynomial.py` |\n| `EquivariantPolynomial` | `cuequivariance/group_theory/equivariant_polynomial.py` |\n| `IrDictPolynomial` | `cuequivariance/group_theory/ir_dict_polynomial.py` |\n| Descriptors | `cuequivariance/group_theory/descriptors/` |\n| Tensor product descriptors | `cuequivariance/group_theory/descriptors/irreps_tp.py` |\n| `spherical_harmonics` | `cuequivariance/group_theory/descriptors/spherical_harmonics_.py` |\n| `symmetric_contraction` | `cuequivariance/group_theory/descriptors/symmetric_contractions.py` |\n"},"changes":[],"summary":"First saved snapshot. No earlier version is available for comparison.","summary_kind":"deterministic","summary_metadata":{}}