cuequivariance

cuEquivariance 核心

使用 cuequivariance（import cuequivariance as cue）描述群表示和张量收缩。如果内置描述符能表达所需操作，优先使用；只有自定义收缩时才构造 SegmentedTensorProduct。核心多项式可以在 CPU 上用 NumPy 求值。它们不会创建可训练的框架参数，也不会选择 GPU 内核。

选择工作流

• 对于内置群和稠密描述符，请使用下方示例。
• 对于新群，请阅读自定义不可约表示。
• 对于自定义收缩、CG 路径或多项式检查，请阅读分段多项式。
• 对于以不可约表示为键的字典，请阅读 ir_dict 工作流。 其中包含 CPU 示例，以及与 JAX 执行的边界说明。

对于解释说明或要求提供可运行代码的请求，请改编最接近的示例；回答问题不以安装软件包为前提。请阅读所选工作流的参考资料。用户要求运行代码时，或集成或正确性问题需要运行时检查时，再执行代码。若已安装的 API 与预期不符，或数值检查失败，请检查上游源码。

检查执行环境

这些示例已通过 cuEquivariance 0.12.0 检查。执行时，请同时检查所选解释器以及该工作流所需的 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 的签名及其版本的后端要求。 官方文档和 0.12.0 源码介绍了这些独立的软件包。

群、重数与布局

群	构造方式	维度和含义
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；非负整数或半整数自旋

cue.Irreps("SO3", "16x0 + 4x1 + 2x2") 的维度为 38。迭代时会得到 (multiplicity, irrep) 块。将 Irreps 与整数相乘会缩放重数；单个不可约表示上的 ir1 * ir2 则给出选择规则。对于 SO(3)，允许的输出角动量范围为从 abs(l1-l2) 到 l1+l2。对于 O(3)，输出奇偶性还必须等于 p1*p2。

布局定义在每个不可约表示块内部，而非整个向量范围内：

布局	展平前的块形状
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 不会改变元素顺序。 同样，对不可约表示元数据排序或简化，也不会重排已有张量。求值时请保留描述符的块顺序。

可运行的稠密张量积

当用户要求输出所有不可约表示时，请从 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()

# A multiplicity summary for display; keep e.outputs[0] for out's actual layout.
summary = e.outputs[0].irreps.sort().irreps.simplify()
assert summary == cue.Irreps("SO3", "3x0 + 5x1 + 3x2 + 1x3")

选择描述符

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

描述符	用途和操作数顺序
linear(irreps_in, irreps_out)	混合匹配不可约表示的重数；权重，输入 → 输出
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)	不带权重的配对通道；输入需要具有相同数量的不可约表示副本，且块边界兼容
spherical_harmonics(cue.O3(1, -1), [0, 1, 2, 3])	向量 → 带宇称的调和函数；若需要单位球面值，请单独归一化输入向量
symmetric_contraction(irreps_in, irreps_out, (1, 2, 3))	MACE 风格的多项式相关；检查描述符形状和重数约束

通道式乘积的第三个参数用于筛选不可约表示类型；它不会设置输出重数。通道式乘积会保留通道对（u, v），而逐元素乘积使用匹配的通道（u）。当稠密的通道式描述符应合并重复的输出不可约表示时，请使用 simplify_irreps3=True。

EquivariantPolynomial 在 inputs 和 outputs 中存储稠密表示元数据，并在 polynomial 中存储其收缩结果。调用它会返回一个 NumPy 输出列表，即使只有一个输出也是如此。对应的 _ir_dict 变体会返回元数据以及已按不可约表示拆分的多项式；请参阅链接的工作流。

对于自定义实现，请检查输出维度和有限值，然后验证其在非平凡群变换下的协变性。输入和预期输出必须使用相同的群、基、块顺序和布局。