SegmentedTensorProduct#

class cuequivariance.SegmentedTensorProduct(
*,
operands: list[Operand] | None = None,
paths: list[Path] | None = None,
coefficient_subscripts: str = '',
)#

与 Irreps 无关且与数据流无关的分段张量积描述符

参数:

  • operands (list of operands) – 张量积的操作数。每个操作数对应下标和段列表。

  • paths (list of paths) – 每个路径包含系数和索引列表。索引是操作数段的索引。

  • coefficient_subscripts (str) – 系数的下标。

我们通常使用 from_subscripts 方法来创建描述符，然后逐个添加段和路径。

方法

classmethod from_subscripts(
subscripts: Subscripts,
) → SegmentedTensorProduct#

从下标字符串创建描述符。

示例

>>> d = cue.SegmentedTensorProduct.from_subscripts("uv,ui,vj+ij")
>>> i0 = d.add_segment(0, (2, 3))
>>> i1 = d.add_segment(1, (2, 5))
>>> i2 = d.add_segment(2, (3, 4))
>>> d.add_path(i0, i1, i2, c=np.ones((5, 4)))
0
>>> print(d)
uv,ui,vj+ij operands=[(2, 3)],[(2, 5)],[(3, 4)] paths=[op0[0]*op1[0]*op2[0]*c c.shape=(5, 4) c.nnz=20]
classmethod empty_segments(
num_segments: list[int],
) → SegmentedTensorProduct#

创建具有简单结构的描述符。

示例

>>> cue.SegmentedTensorProduct.empty_segments([2, 3, 4])
,, sizes=2,3,4 num_segments=2,3,4 num_paths=0
classmethod from_dict(
data: dict[str, Any],
) → SegmentedTensorProduct#

从字典创建描述符。

classmethod from_json(
data: str,
) → SegmentedTensorProduct#

从 JSON 字符串创建描述符。

classmethod from_bytes(
data: bytes,
) → SegmentedTensorProduct#

从压缩二进制数据创建描述符。

classmethod from_base64(
data: str,
) → SegmentedTensorProduct#

从 base64 字符串创建描述符。

property num_operands: int#

操作数数量。

property num_paths: int#

路径数量。

property subscripts: Subscripts#

张量积的下标。

property indices: ndarray#

路径的索引。

property coefficients_are_stackable: bool#

检查系数是否可堆叠。

property stacked_coefficients: ndarray#

单个数组中堆叠的路径系数。

to_text(
coefficient_formatter=<function SegmentedTensorProduct.<lambda>>,
) → str#

描述符的人类可读文本表示。

参数:

coefficient_formatter (callable, optional) – 用于格式化系数的函数。

示例

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1")
... ).d
>>> d = d.flatten_coefficient_modes()
>>> print(d.to_text())
uvw,u,v,w sizes=320,16,16,16 num_segments=5,4,4,4 num_paths=16 u=4 v=4 w=4
Flop cost: 0->1344 1->2368 2->2368 3->2368
Memory cost from 368 to 1216
Path indices: 0 0 0 0, 1 0 1 1, 1 0 2 2, 1 0 3 3, 2 1 0 1, 2 2 0 2, ...
Path coefficients: [0.17...]
to_dict(
extended: bool = False,
) → dict[str, Any]#

描述符的字典表示。

to_json(extended: bool = False) → str#

描述符的 JSON 表示。

to_bytes(extended: bool = False) → bytes#

描述符的压缩二进制表示。

to_base64(extended: bool = False) → str#

描述符的 Base64 表示。

示例

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1")
... ).d
>>> print(d.to_base64())
eJytkstuwjAQRX/F8r...lTF2zlX91/fHyvj2Z4=
get_dimensions_dict() → dict[str, set[int]]#

获取张量积的维度。

get_dims(m: str) → set[int]#

获取特定模式的维度。

示例

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+8x1"),
...     cue.Irreps("SO3", "3x0+3x1"),
...     cue.Irreps("SO3", "5x0+7x1")
... ).d
>>> d.get_dims("u")
{8, 4}
>>> d.get_dims("v")
{3}
>>> d.get_dims("w")
{5, 7}
get_path_dimensions_dict(
path: int | Path,
*,
returns_sets: bool = False,
) → dict[str, int | set[int]]#

获取特定路径的维度。

get_path_dim(
path: int | Path,
m: str,
*,
returns_set=False,
) → int | set[int]#

获取特定路径中特定模式的维度。

segment_slice(
operand: int,
path: int | Path,
) → slice#

获取给定路径选择的给定操作数中段的切片。

get_segment_shape(
operand: int,
path: int | Path,
) → tuple[int, ...]#

获取给定路径选择的给定操作数中段的形状。

all_segments_are_used() → bool#

检查张量积中是否使用了所有段。

compressed_path_segment(
operand: int,
) → ndarray#

指定操作数段的路径起始索引。

注意：此方法要求路径按指定操作数排序。

参数:

operand (int) – 要查找路径组的操作数的索引。

返回:

指定操作数中每个段的起始路径索引数组。

返回类型:

np.ndarray

示例

indices[:, operand], operands[operand].num_segments -> compressed_path_segment(operand)
[0, 0, 1, 1, 1], 2 -> [0, 2, 5]
[0, 0, 1, 1, 1], 3 -> [0, 2, 5, 5]
[0, 0, 2, 2, 2], 3 -> [0, 2, 2, 5]
[1, 1], 2 -> [0, 0, 2]
symmetries() → list[tuple[int, ...]]#

使张量积保持不变的置换列表。

coefficients_equal_one() → bool#

检查所有系数是否等于 1。

flop_cost(
operand: int,
algorithm: str = 'optimal',
) → int#

计算计算指定操作数所需的浮点运算次数。

参数:

  • operand (int) – 要计算浮点运算成本的操作数。

  • algorithm (str, optional) – 用于计算成本的算法。可以是 ‘optimal’ 或 ‘naive’。

返回:

计算指定操作数所需的浮点运算次数。

返回类型:

int

memory_cost(
algorithm: str = 'sequential',
) → int#

计算计算指定操作数所需的内存访问次数。

参数:

  • operand (int) – 要计算内存成本的操作数。

  • algorithm (str, optional) – 用于计算成本的算法。可以是 ‘sequential’ 或 ‘global’。

返回:

计算指定操作数所需的内存访问次数。

返回类型:

int

insert_path(
path_index,
*segments: int | tuple[int, ...] | dict[str, int] | None,
c: ndarray,
dims: dict[str, int] | None = None,
) → int#

在特定索引处插入路径。

add_path(
*segments: int | tuple[int, ...] | dict[str, int] | None,
c: ndarray,
dims: dict[str, int] | None = None,
) → int#

向描述符添加路径。

参数:

  • segments – 指定路径中收缩的操作数段。

  • c (np.ndarray) – 路径的系数。

  • dims (dict[str, int], optional) – 模式的范围。

返回:

添加路径的索引。

返回类型:

int

示例

>>> d = cue.SegmentedTensorProduct.from_subscripts("uv,ui,vj+ij")
>>> i1 = d.add_segment(1, (2, 3))
>>> i2 = d.add_segment(2, (2, 5))

我们可以使用 None 动态添加新段： >>> d.add_path(None, i1, i2, c=np.ones((3, 5))) 0

描述符现在在第一个操作数中有一个新段 (2, 2): >>> d uv,ui,vj+ij operands=[(2, 2)],[(2, 3)],[(2, 5)] paths=[op0[0]*op1[0]*op2[0]*c c.shape=(3, 5) c.nnz=15] >>> d.add_path(0, None, None, c=np.ones((10, 10))) 1 >>> d uv,ui,vj+ij sizes=4,26,30 num_segments=1,2,2 num_paths=2 i={10, 3} j={10, 5} u=2 v=2

当无法推断模式的维度时，我们可以提供它们： >>> d.add_path(None, None, None, c=np.ones((2, 2)), dims={“u”: 2, “v”: 2}) 2 >>> d uv,ui,vj+ij sizes=8,30,34 num_segments=2,3,3 num_paths=3 i={2, 10, 3} j={10, 2, 5} u=2 v=2

insert_segment(
operand: int,
sid: int,
segment: tuple[int, ...] | dict[str, int],
)#

在特定索引处插入一个段。

insert_segments(
operand: int,
sid: int,
segments: list[tuple[int, ...]] | Operand,
)#

在特定索引处插入多个段。

add_segment(
operand: int,
segment: tuple[int, ...] | dict[str, int],
) → int#

向描述符添加一个段。

add_segments(
operand: int,
segments: list[tuple[int, ...] | dict[str, int]],
)#

向描述符添加多个段。

canonicalize_subscripts() → SegmentedTensorProduct#

返回一个新的描述符，其中包含下标的规范化表示。

示例

>>> d = cue.SegmentedTensorProduct.from_subscripts("ab,ax,by+xy")
>>> d.canonicalize_subscripts()
uv,ui,vj+ij sizes=0,0,0 num_segments=0,0,0 num_paths=0

这对于识别等效的描述符很有用。

add_or_rename_modes(
subscripts: str,
*,
mapping: dict[str, str] | None = None,
) → SegmentedTensorProduct#

返回一个新的描述符，其中的模式根据新的下标重命名。

参数:

  • subscripts (str) – 包含模式新名称的新下标。新下标也可以是旧下标的超集。

  • mapping (dict of str to str, optional) – 旧模式和新模式之间的映射。

返回:

具有重命名模式的新描述符。

返回类型:

SegmentedTensorProduct

add_or_transpose_modes(
subscripts: str,
dims: dict[str, int] | None = None,
) → SegmentedTensorProduct#

返回一个新的描述符，其中的模式根据新的下标进行转置。

参数:

  • subscripts (str) – 一个新的下标，其中包含当前下标模式的排列。

  • dims (dict of str to int, optional) – 新模式的维度。

返回:

具有转置模式的新描述符。

返回类型:

SegmentedTensorProduct

append_modes_to_all_operands(
modes: str,
dims: dict[str, int],
) → SegmentedTensorProduct#

返回一个新的描述符，其中的模式被附加（在右侧）到所有操作数。

参数:

  • modes (str) – 要附加到所有操作数的新段模式。

  • dims (dict of str to int) – 新模式的维度。

返回:

具有附加模式的新描述符。

返回类型:

SegmentedTensorProduct

permute_operands(
perm: tuple[int, ...],
) → SegmentedTensorProduct#

排列描述符的操作数。

move_operand(
operand: int,
new_index: int,
) → SegmentedTensorProduct#

将一个操作数移动到新的索引位置。

move_operand_first(
operand: int,
) → SegmentedTensorProduct#

将一个操作数移动到第一个位置。

move_operand_last(
operand: int,
) → SegmentedTensorProduct#

将一个操作数移动到最后一个位置。

permute_segments(
operand: int,
perm: tuple[int, ...],
) → SegmentedTensorProduct#

排列一个操作数的段。

sort_paths(
operands_ordering: int | Sequence[int] | None = None,
) → SegmentedTensorProduct#

按路径索引的字典顺序对路径进行排序。

参数:

operands_ordering (int or sequence of int, optional) – 对路径进行排序的顺序。如果是一个整数，则按该操作数排序。如果是一个序列，则先按序列中的第一个操作数排序，如果相等，则按第二个操作数排序，依此类推。

返回:

排序后的描述符。

返回类型:

SegmentedTensorProduct

squeeze_modes(
modes: str | None = None,
) → SegmentedTensorProduct#

通过移除始终为 1 的维度来压缩描述符。

参数:

modes (str, optional) – 要压缩的模式。如果为 None，则压缩所有始终为 1 的模式。

返回:

压缩后的描述符。

返回类型:

SegmentedTensorProduct

split_mode(
mode: str,
size: int,
) → SegmentedTensorProduct#

将一个模式拆分为多个给定大小的模式。

参数:

  • mode (str) – 要拆分的模式。其所有维度都必须能被 size 整除。

  • size (int) – 新模式的大小。

返回:

新的描述符。

返回类型:

SegmentedTensorProduct

all_same_segment_shape() → bool#

检查所有段是否具有相同的形状。

normalize_paths_for_operand(
operand: int,
) → SegmentedTensorProduct#

规范化操作数的路径。

参数:

operand (int) – 要规范化的操作数的索引。

假设输入操作数具有单位方差，此方法计算所选操作数中每个段的方差，并将路径的系数统一地按段方差的平方根进行规范化。这对于确保输出具有单位方差很有用。

remove_zero_paths() → SegmentedTensorProduct#

移除系数为零的路径。

fuse_paths_with_same_indices() → SegmentedTensorProduct#

融合具有相同索引的路径。

consolidate_paths() → SegmentedTensorProduct#

通过合并重复项和移除零值来整合路径。

sort_indices_for_identical_operands(
operands: Sequence[int],
) → SegmentedTensorProduct#

通过对相同操作数的索引进行排序来减少路径数量。

symmetrize_operands(
operands: Sequence[int],
) → SegmentedTensorProduct#

对称化指定的通过排列索引的操作数。

remove_empty_segments() → SegmentedTensorProduct#

移除空段。

flatten_modes(
modes: Sequence[str],
*,
skip_zeros=True,
force=False,
) → SegmentedTensorProduct#

通过细分段和路径来移除指定的模式。

参数:

  • modes (Sequence of str) – 要移除的模式，它们必须在每个操作数中保留的模式之前。

  • skip_zeros (bool, optional) – 是否跳过系数为零的路径。默认为 True。

  • force (bool, optional) – 是否通过扁平化额外的必要模式来强制扁平化。默认为 False。

flatten_coefficient_modes(
*,
skip_zeros=True,
force=False,
) → SegmentedTensorProduct#

扁平化描述符的系数。创建新的段和路径以扁平化系数。

参数:

  • skip_zeros (bool, optional) – 是否跳过系数为零的路径。默认为 True。

  • force (bool, optional) – 是否通过扁平化额外的必要模式来强制扁平化。默认为 False。

consolidate_modes(
modes: str | None = None,
) → SegmentedTensorProduct#

通过合并模式来整合描述符。

参数:

modes (str, optional) – 要整合的模式。如果为 None，则整合所有重复的模式。

round_coefficients_to_rational(
max_denominator: int,
) → SegmentedTensorProduct#

将系数四舍五入到最接近的 p / q 数，并给定最大分母。

参数:

max_denominator (int) – 最大分母，q < max_denominator。

round_coefficients_to_sqrt_rational(
max_denominator: int,
) → SegmentedTensorProduct#

将系数四舍五入到最接近的 sqrt(p / q) 数，并给定最大分母。

参数:

max_denominator (int) – 最大分母，q < max_denominator。

modify_coefficients(
f: Callable[[ndarray], ndarray],
) → SegmentedTensorProduct#

修改描述符的系数。

参数:

f (callable) – 应用于系数的函数。