重复维度的 Einsum 公式

Question

我有这段代码：

other = np.random.rand((m,n,o))
prev = np.random.rand((m,n,o,m,n,o))
mu = np.zeros((m,n,o,m,n,o))
for c in range(m):
   for i in range(n):
      for j in range(o):
         mu[c,i,j,c,i,j] = other[c,i,j]*prev[c,i,j,c,i,j]

我想使用 einsum 符号来简化它（可能通过跳过 python 中的 for 循环来节省时间）。然而，经过几次尝试，我最终不确定如何解决这个问题。我目前的尝试是：

np.einsum('cijklm,cij->cijklm', prev, other)

它没有实现与“for-loop”代码段相同的结果。

Answer 1

对于形状 (2,3,4)，我得到：

In [52]: mu.shape
Out[52]: (2, 3, 4, 2, 3, 4)

这个 einsum 表达式抱怨输出中重复了维度：

In [53]: np.einsum('cijcij,cij->cijcij', prev, other).shape
Traceback (most recent call last):
  File "<ipython-input-53-92862a0865a2>", line 1, in <module>
    np.einsum('cijcij,cij->cijcij', prev, other).shape
  File "<__array_function__ internals>", line 180, in einsum
  File "/usr/local/lib/python3.8/dist-packages/numpy/core/einsumfunc.py", line 1359, in einsum
    return c_einsum(*operands, **kwargs)
ValueError: einstein sum subscripts string includes output subscript 'c' multiple times

没有重复：

In [55]: x=np.einsum('cijcij,cij->cij', prev, other)
In [56]: x.shape
Out[56]: (2, 3, 4)

非零值匹配：

In [57]: np.allclose(mu[np.nonzero(mu)].ravel(), x.ravel())
Out[57]: True

或者从 mu 中提取对角线：

In [59]: I,J,K = np.ix_(np.arange(2),np.arange(3),np.arange(4))
In [60]: mu[I,J,K,I,J,K].shape
Out[60]: (2, 3, 4)
In [61]: np.allclose(mu[I,J,K,I,J,K],x)
Out[61]: True

您的 einsum 满足相同的 'diagonals' 测试：

In [68]: y=np.einsum('cijklm,cij->cijklm', prev, other)
In [69]: y.shape
Out[69]: (2, 3, 4, 2, 3, 4)
In [70]: np.allclose(y[I,J,K,I,J,K],x)
Out[70]: True

因此 mu 值也出现在 y 中，但分布方式不同。但是数组太大了，不容易查看和比较。

OK，每个y[i,j,k]都一样，等于x。在 mu 中，大多数这些值都是 0，只有选定的对角线是非零的。

虽然 einsum 可以生成相同的非零值，但它不能以与循环相同的 3d 对角线方式分布它们。

更改您的 mu 计算以生成 3d 数组：

In [76]: nu = np.zeros((m,n,o))
    ...: for c in range(m):
    ...:    for i in range(n):
    ...:       for j in range(o):
    ...:          nu[c,i,j] = other[c,i,j]*prev[c,i,j,c,i,j]
    ...: 
In [77]: np.allclose(nu,x)
Out[77]: True

编辑

我们可以将 einsum 结果分配给对角线：

In [134]: out = np.zeros((2,3,4,2,3,4))
In [135]: out[I,J,K,I,J,K] = x
In [136]: np.allclose(out, mu)
Out[136]: True

从概念上讲，它可能比 as_strided 解决方案更简单。并且可能一样快。 as_strided 在制作 view 时不如 reshape 那种 view.

快

In [143]: %%timeit
     ...: out = np.zeros((m, n, o, m, n, o))
     ...: mu_view = np.lib.stride_tricks.as_strided(out,
     ...:                      shape=(m, n, o),
     ...:                      strides=[sum(mu.strides[i::3]) for i in range(3)
     ...: ]
     ...:                      )
     ...: np.einsum('cijcij,cij->cij', prev, other, out=mu_view)
     ...: 
     ...: 
31.6 µs ± 69.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [144]: %%timeit
     ...: out = np.zeros((2,3,4,2,3,4))
     ...: out[I,J,K,I,J,K] =np.einsum('cijcij,cij->cij', prev, other)
     ...: 
     ...: 
18.5 µs ± 178 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

或在时间循环中包括I,J,K世代

In [146]: %%timeit
     ...: I,J,K = np.ix_(np.arange(2),np.arange(3),np.arange(4))
     ...: out = np.zeros((2,3,4,2,3,4))
     ...: out[I,J,K,I,J,K] =np.einsum('cijcij,cij->cij', prev, other)
40.4 µs ± 1.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

或直接创建 IJK：

In [151]: %%timeit
     ...: I,J,K = np.arange(2)[:,None,None],np.arange(3)[:,None],np.arange(4)
     ...: out = np.zeros((2,3,4,2,3,4))
     ...: out[I,J,K,I,J,K] =np.einsum('cijcij,cij->cij', prev, other)
25.1 µs ± 38.3 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 2

单独使用 np.einsum() 无法获得此结果，但您可以尝试这样做：

import numpy as np
from numpy.lib.stride_tricks import as_strided

m, n, o = 2, 3, 5
np.random.seed(0)
other = np.random.rand(m, n, o)
prev = np.random.rand(m, n, o, m, n, o)
mu = np.zeros((m, n, o, m, n, o))

mu_view = as_strided(mu,
                     shape=(m, n, o),
                     strides=[sum(mu.strides[i::3]) for i in range(3)]
                     )
np.einsum('cijcij,cij->cij', prev, other, out=mu_view)

数组 mu 应该与问题中使用嵌套循环的代码生成的数组相同。

一些解释。无论 numpy 数组的形状如何，其元素在内部都存储在连续的内存块中。数组结构的一部分是步幅，它指定当数组元素的一个索引递增 1 时需要跳转多少字节。因此，在二维数组中 arr，arr.stride[0] 是将元素 arr[i, j] 与 arr[i+1, j] 分开的字节数，arr.stride[1] 是将 arr[i, j] 与 a[i, j+1] 分开的字节数。使用步幅信息，numpy 可以根据其索引在数组中找到给定元素。参见例如 post 了解更多详情。

numpy.lib.stride_tricks.as_strided 是一个函数，它以 custom-made 步幅创建给定数组的视图。通过指定步幅，可以更改哪个数组元素对应于哪个索引。在上面的代码中，这用于创建 mu_view，它是 mu 和属性的视图，元素 mu_view[c, i, j] 是元素 mu[c, i, j, c, i, j]。这是通过根据 mu 的步幅指定 mu_view 的步幅来完成的。例如mu_view[c, i, j]和mu_view[c+1, i, j]之间的距离设置为mu[c, i, j, c, i, j]和mu[c+1, i, j, c+1, i, j]之间的距离，即mu.strides[0] + mu.strides[3].

重复维度的 Einsum 公式

Einsum formula for repeating dimensions

python

numpy

numpy-einsum

编辑