Question

我需要执行一种特殊类型的张量收缩。我想要这样的东西：

A_ {bg} = Sum_ {a，a'，a''}（B_ {a} C_ {a'b} D_ {a''g}）

其中所有指数的值都为0,1且总和超过a，a'和a''适用于a + a'+ a''= 1或a + a'+ a''的所有情况= 2.所以它就像爱因斯坦求和公约的反面：我只想在三个指数中的一个与其他指数不同时求和。

此外，我想要一些灵活性与未被求和的索引的数量：在示例中，得到的张量有2个索引，并且总和超过3个张量的元素的产品，一个具有一个索引，另外两个有两个指数。这些指数的数量会有所不同，所以一般来说我希望能写出这样的东西：

A _ {...} = Sum_ {a，a'，a''}（B_ {a ...} C_ {a ...} D_ {a''...}}

我想指出索引的数量不是固定的，但它是受控的：我可以知道并指定每个张量中每个步长有多少个索引。

我试过np.einsum()，但显然我被迫在标准的爱因斯坦惯例中总结重复的指数，我不知道如何实现我在这里暴露的条件。

我不能用各种来写所有内容，因为正如我所说，所涉及的张量索引的数量并不固定。

有人有想法吗？

来自评论：

我会用这样的编程语言写下我所写的内容：

tensa = np.zeros((2,2))
for be in range(2):
    for ga in range(2):
        for al in range(2):
            for alp in range(2):
                for alpp in range(res(al,alp),prod(al,alp)):
                    tensa[be,ga] += tensb[al] * tensc[alp,be] * tensd[alpp,ga]

其中res和prod是确保al+alp+alpp = 1 or 2的两个函数。这个问题是我需要指定所有涉及的索引，而且我不能在所有格子的一般计算中这样做。

Answer 1

首先，让我们在Python循环中编写您的示例，以获得比较基线。如果我理解正确，这就是你想要做的事情：

b, g = 4, 5
B = np.random.rand(2)
C = np.random.rand(2, b)
D = np.random.rand(2, g)

out = np.zeros((b, g))
for j in (0, 1):
    for k in (0, 1):
        for l in (0, 1):
            if j + k + l in (1, 2):
                out += B[j] * C[k, :, None] * D[l, None, :]

当我运行它时，我得到了这个输出：

>>> out
array([[ 1.27679643,  2.26125361,  1.32775173,  1.5517918 ,  0.47083151],
       [ 0.84302586,  1.57516142,  1.1335904 ,  1.14702252,  0.34226837],
       [ 0.70592576,  1.34187278,  1.02080112,  0.99458563,  0.29535054],
       [ 1.66907981,  3.07143067,  2.09677013,  2.20062463,  0.65961165]])

您无法直接使用np.einsum获取此信息，但您可以运行两次，并将结果视为这两者的差异：

>>> np.einsum('i,jk,lm->km', B, C, D) - np.einsum('i,ik,im->km', B, C, D)
array([[ 1.27679643,  2.26125361,  1.32775173,  1.5517918 ,  0.47083151],
       [ 0.84302586,  1.57516142,  1.1335904 ,  1.14702252,  0.34226837],
       [ 0.70592576,  1.34187278,  1.02080112,  0.99458563,  0.29535054],
       [ 1.66907981,  3.07143067,  2.09677013,  2.20062463,  0.65961165]])

对np.einsum的第一次调用是添加所有内容，无论索引加起来是什么。第二个只是将所有三个指数相同的那些加起来。显然你的结果是两者的区别。

理想情况下，您现在可以继续写下这样的内容：

>>>(np.einsum('i...,j...,k...->...', B, C, D) -
... np.einsum('i...,i...,i...->...', B, C, D))

并获得您的结果，无论您的C和D数组的尺寸如何。如果您尝试第一个，您将收到以下错误消息：

ValueError: operands could not be broadcast together with remapped shapes
[original->remapped]: (2)->(2,newaxis,newaxis) (2,4)->(4,newaxis,2,newaxis)
                      (2,5)->(5,newaxis,newaxis,2)

问题在于，由于您未指定要对张量的b和g尺寸进行操作，因此会尝试将它们一起广播，因为它们不同，所以失败。您可以通过添加大小为1的额外维度来实现它：

>>> (np.einsum('i...,j...,k...->...', B, C, D[:, None]) -
...  np.einsum('i...,i...,i...->...', B, C, D[:, None]))
array([[ 1.27679643,  2.26125361,  1.32775173,  1.5517918 ,  0.47083151],
       [ 0.84302586,  1.57516142,  1.1335904 ,  1.14702252,  0.34226837],
       [ 0.70592576,  1.34187278,  1.02080112,  0.99458563,  0.29535054],
       [ 1.66907981,  3.07143067,  2.09677013,  2.20062463,  0.65961165]])

如果您希望B的所有轴都放在C的所有轴之前，并且这些轴放在D的所有轴之前，则以下似乎可行，至少就创建正确形状的输出而言，尽管你可能想要仔细检查结果是否真的如你所愿：

>>> B = np.random.rand(2, 3)
>>> C = np.random.rand(2, 4, 5)
>>> D = np.random.rand(2, 6)
>>> C_idx = (slice(None),) + (None,) * (B.ndim - 1)
>>> D_idx = C_idx + (None,) * (C.ndim - 1)
>>> (np.einsum('i...,j...,k...->...', B, C[C_idx], D[D_idx]) -
...  np.einsum('i...,i...,i...->...', B, C[C_idx], D[D_idx])).shape
(3L, 4L, 5L, 6L)

编辑从评论中，如果不是每个张量的第一个轴都必须减少，那么它是前两个，那么上面的内容可以写成：

>>> B = np.random.rand(2, 2, 3)
>>> C = np.random.rand(2, 2, 4, 5)
>>> D = np.random.rand(2, 2, 6)
>>> C_idx = (slice(None),) * 2 + (None,) * (B.ndim - 2)
>>> D_idx = C_idx + (None,) * (C.ndim - 2)
>>> (np.einsum('ij...,kl...,mn...->...', B, C[C_idx], D[D_idx]) -
...  np.einsum('ij...,ij...,ij...->...', B, C[C_idx], D[D_idx])).shape
(3L, 4L, 5L, 6L)

更一般地说，如果减少d个索引，C_idx和D_idx会是这样的：

>>> C_idx = (slice(None),) * d + (None,) * (B.ndim - d)
>>> D_idx = C_idx + (None,) * (C.ndim - d)

并且对np.einsum的调用需要在索引中包含d个字母，在第一个调用中是唯一的，在第二个调用中重复。

编辑2 那么C_idx和D_idx究竟发生了什么？以最后一个示例为例，B，C和D包含形状(2, 2, 3)，(2, 2, 4, 5)和(2, 2, 6)。 C_idx由两个空切片组成，加上None维数B减去C[C_idx]减去2，因此当我们采用(2, 2, 1, 4, 5)时，结果具有形状{ {1}}。同样，D_idx为C_idx加上None与C的维数减2相同，因此D[D_idx]的结果具有(2, 2, 1, 1, 1, 6)的形状}。这三个数组不是一起编写的，但是np.einsum增加了大小为1的附加维度，即上面错误的“重映射”形状，因此得到的数组结果有额外的尾随，并且形状为如下：

(2, 2, 3, 1, 1, 1)
(2, 2, 1, 4, 5, 1)
(2, 2, 1, 1, 1, 6)

前两个轴缩小，因此从输出中消失，在其他情况下，应用广播，其中复制大小为1的维度以匹配较大的维度，因此输出为(3, 4, 5, 6) as我们想要。

@hpaulj提出了一种使用“Levi-Civita like”张量的方法，理论上应该更快，请看我对原始问题的评论。这是一些用于比较的代码：

b, g = 5000, 2000
B = np.random.rand(2)
C = np.random.rand(2, b)
D = np.random.rand(2, g)

def calc1(b, c, d):
    return (np.einsum('i,jm,kn->mn', b, c, d) -
            np.einsum('i,im,in->mn', b, c, d))

def calc2(b, c, d):
    return np.einsum('ijk,i,jm,kn->mn', calc2.e, b, c, d)
calc2.e = np.ones((2,2,2))
calc2.e[0, 0, 0] = 0
calc2.e[1, 1, 1] = 0

但在运行时：

%timeit calc1(B, C, D)
1 loops, best of 3: 361 ms per loop

%timeit calc2(B, C, D)
1 loops, best of 3: 643 ms per loop

np.allclose(calc1(B, C, D), calc2(B, C, D))
Out[48]: True

一个令人惊讶的结果，我无法解释......

Python中的特殊张量收缩

来自评论：

1 个答案: