cuequivariance

用途

cuEquivariance 用于描述群表示（group representations）与张量收缩（tensor contractions）。导入方式为 import cuequivariance as cue。当内置描述符（descriptor）能表达所需操作时，应优先使用内置描述符；仅在需要自定义收缩时，才构造 SegmentedTensorProduct。核心多项式可在 CPU 上用 NumPy 求值，它们不创建可训练框架参数，也不选择 GPU 内核。

主要能力

群、重数与布局

支持的群构造方式：

SO(3)：cue.SO3(l)，维度 2*l + 1，表示旋转，l 为非负整数。
O(3)：cue.O3(l, p)，维度 2*l + 1，还包含宇称 p=+1 或 -1，例如 0e、1o。
SU(2)：cue.SU2(j)，维度 2*j + 1，j 为非负整数或半整数自旋。

cue.Irreps("SO3", "16x0 + 4x1 + 2x2") 的维度为 38。迭代它会产出 (multiplicity, irrep) 块。将 Irreps 乘以整数会缩放重数；对单个 irrep 执行 ir1 * ir2 会给出选择定则。对 SO(3)，允许的输出角动量从 abs(l1-l2) 到 l1+l2。对 O(3)，输出宇称还必须等于 p1*p2。

布局定义在每个 irrep 块内部，而不是整个向量上：

cue.ir_mul：块形状为 (..., ir.dim, mul)
cue.mul_ir：块形状为 (..., mul, ir.dim)

可从描述符读取 e.inputs[i].layout 与 e.outputs[i].layout。此处的张量积示例使用 ir_mul；许多 PyTorch 层默认使用 mul_ir 并接受显式布局选项。JAX 字典功能常使用 mul_ir，而多项式执行可能要求 ir_mul。转换实际数据时需逐块转置后再展平；仅更改元数据或调用 reshape 不会改变元素顺序。同样，对 irreps 元数据排序或简化不会重排已有张量。求值时需保留描述符的块顺序。

可运行的稠密张量积

当用户要求所有输出 irreps 时，从 full_tensor_product 开始。它有两个输入且无权重。示例：

import numpy as np
import cuequivariance as cue

irreps1 = cue.Irreps("SO3", "2x0 + 1x1 + 1x2")
irreps2 = cue.Irreps("SO3", "1x0 + 1x1")
e = cue.descriptors.full_tensor_product(irreps1, irreps2)
rng = np.random.default_rng(0)
x1 = rng.normal(size=(4, e.inputs[0].dim))
x2 = rng.normal(size=(4, e.inputs[1].dim))
[out] = e(x1, x2)
assert out.shape == (4, 40)
assert np.isfinite(out).all()

# 用于展示的重数摘要；out 的实际布局请保留 e.outputs[0]。
summary = e.outputs[0].irreps.sort().irreps.simplify()
assert summary == cue.Irreps("SO3", "3x0 + 5x1 + 3x2 + 1x3")

选择描述符

以下名称均位于 cue.descriptors 下。应检查返回的输入/输出元数据来分配数组（包括权重）；不要猜测权重数量。

linear(irreps_in, irreps_out)：混合匹配 irreps 的重数；权重，输入 → 输出。
fully_connected_tensor_product(a, b, out)：选择精确的输出重数；权重，输入1，输入2 → 输出。
full_tensor_product(a, b)：所有允许的耦合，无权重；输入1，输入2 → 输出。
channelwise_tensor_product(a, b, irreps3_filter=...)：加权通道对；每条路径的输出重数为 mul1*mul2。
elementwise_tensor_product(a, b)：无权重配对通道；输入需要相同数量的 irrep 副本且块边界兼容。
spherical_harmonics(cue.O3(1, -1), [0, 1, 2, 3])：向量 → 带宇称的谐波；若需要单位球面值，需单独归一化输入向量。
symmetric_contraction(irreps_in, irreps_out, (1, 2, 3))：MACE 风格多项式关联；检查描述符形状与重数约束。

channelwise 积的第三个参数用于过滤 irrep 类型，不设置输出重数。channelwise 保留通道对 (u, v)，不同于 elementwise 的匹配通道 (u)。当稠密 channelwise 描述符应合并重复输出 irreps 时，使用 simplify_irreps3=True。

EquivariantPolynomial 在 inputs 与 outputs 中存储稠密表示元数据，其收缩存储在 polynomial 中。调用它会返回 NumPy 输出列表，即使只有一个输出也是如此。对应的 _ir_dict 变体返回元数据以及已按 irrep 拆分的多项式。

使用方式

选择工作流：

内置群与稠密描述符：使用上述示例。
新群：阅读 custom irreps。
自定义收缩、CG 路径或多项式检查：阅读 segmented polynomials。
以 irrep 为键的字典：阅读 ir_dict 工作流，其中包含 CPU 示例以及与 JAX 执行的边界。

对于解释或展示可运行代码的请求，应改编最接近的示例；安装包不是回答的前提。阅读所选工作流的参考文档。当用户请求运行，或集成/正确性问题需要运行时检查时，执行代码。当已安装 API 不同或数值检查失败时，检查上游源码。

检查执行环境：

from importlib.metadata import PackageNotFoundError, version
import sys

print("Python:", sys.executable)
try:
    installed = version("cuequivariance")
except PackageNotFoundError:
    print("cuequivariance is not installed in this interpreter")
else:
    import cuequivariance as cue

    print("cuEquivariance:", installed)
    print("IrDictPolynomial available:", hasattr(cue, "IrDictPolynomial"))

复用兼容的项目环境及其依赖约束。若执行需要缺失的包，使用隔离的示例执行。保留现有虚拟环境以及系统/用户 Python 安装；不要删除环境或绕过包管理器保护来修复设置。当安装不可用或被禁止时，提供可运行代码并说明哪些检查未执行。核心示例需要 NumPy，但不需要 CUDA 或 JAX/PyTorch。若 API 不可用，在满足任务时使用受支持的稠密描述符；不要发明方法或强制升级环境。

对于框架任务，检查已安装的 cuequivariance_jax 或 cuequivariance_torch 签名及其版本的 backend 要求。官方文档与 0.12.0 源码覆盖这些独立包。

适用场景

构建与调试 cuEquivariance irreps、自定义 Irrep 子类。
Clebsch-Gordan 张量积。
等变或分段多项式。
以 irrep 为键的字典工作流。
自定义收缩与 CG 路径。

限制

核心多项式不创建可训练框架参数，也不选择 GPU 内核。
核心示例需要 NumPy，但不需要 CUDA 或 JAX/PyTorch。
布局转换需逐块转置，仅改元数据或 reshape 无效。
排序或简化 irreps 元数据不会重排已有张量。
不要猜测权重数量，应检查描述符元数据。
自定义实现需检查输出维度与有限值，并在非平凡群变换下验证协变性；输入与期望输出需使用相同的群、基、块顺序与布局。
资料未提供收费、价格、平台支持范围等具体信息。