如何在numpy中对2D数组进行三重循环向量化？

Question

我能否在这个计算中消除所有的Python循环：

result[i,j,k] = (x[i] * y[j] * z[k]).sum()

其中x[i]、y[j]和z[k]是长度为N的向量，x、y和z的第一个维度分别具有长度A、B、C，输出的形状为(A,B,C)，每个元素是三重积的总和（逐个元素相乘）。

我可以将代码从3个循环减少到1个循环（下方代码），但我卡在尝试消除最后一个循环上。

如果必要，我可以通过轻微填充来使A=B=C。

# Example with 3 loops, 2 loops, 1 loop (testing omitted)

N = 100 # more like 100k in real problem
A =   2 # more like 20 in real problem
B =   3 # more like 20 in real problem
C =   4 # more like 20 in real problem

import numpy
x = numpy.random.rand(A, N)
y = numpy.random.rand(B, N)
z = numpy.random.rand(C, N)

# outputs of each variant
result_slow = numpy.empty((A,B,C))
result_vec_C = numpy.empty((A,B,C))
result_vec_CB = numpy.empty((A,B,C))

# 3 nested loops
for i in range(A):
    for j in range(B):
        for k in range(C):
            result_slow[i,j,k] = (x[i] * y[j] * z[k]).sum()

# vectorize loop over C (2 nested loops)
for i in range(A):
    for j in range(B):
        result_vec_C[i,j,:] = (x[i] * y[j] * z).sum(axis=1)

# vectorize one C and B (one loop)
for i in range(A):
    result_vec_CB[i,:,:] = numpy.dot(x[i] * y, z.transpose())

numpy.testing.assert_almost_equal(result_slow, result_vec_C)
numpy.testing.assert_almost_equal(result_slow, result_vec_CB)

- Joseph Hastings

---

很遗憾，这不是一个作业问题。如果有关于“如何向量化”的通用主题的课程/教科书，我会非常兴奋！ - Joseph Hastings

Answer 1

如果您使用的是numpy > 1.6版本，那么有一个非常棒的np.einsum函数：

np.einsum('im,jm,km->ijk',x,y,z)

这与您的循环版本相当。我不确定一旦您在实际问题中使用数组的大小，效率会如何（当我转移到那些大小时，我的机器实际上会出现segfault）。我通常喜欢针对这些问题重新编写方法并使用Cython。

- JoshAdel

Answer 2

在你的情况下，使用einsum很有意义; 但是您可以轻松地手动完成此操作。诀窍是使数组相互广播。这意味着重新整形它们，以便每个数组沿其自己的轴独立变化。然后将它们相乘，让numpy处理广播; 然后沿着最后（最右边）的轴进行求和。

>>> x = numpy.arange(2 * 4).reshape(2, 4)
>>> y = numpy.arange(3 * 4).reshape(3, 4)
>>> z = numpy.arange(4 * 4).reshape(4, 4)
>>> (x.reshape(2, 1, 1, 4) * 
...  y.reshape(1, 3, 1, 4) *
...  z.reshape(1, 1, 4, 4)).sum(axis=3)
array([[[  36,   92,  148,  204],
        [  92,  244,  396,  548],
        [ 148,  396,  644,  892]],

       [[  92,  244,  396,  548],
        [ 244,  748, 1252, 1756],
        [ 396, 1252, 2108, 2964]]])

您可以使用切片符号、newaxis（它等于None，因此以下内容同样适用于None）和 sum 可以接受负的轴值（-1 表示最后一个，-2 表示次后一个，以此类推），使这个过程更加通用化。 这样做的好处是，您不必知道数组的原始形状；只要它们的最后一维兼容，这将广播前三个数组：

>>> (x[:, numpy.newaxis, numpy.newaxis, :] *
...  y[numpy.newaxis, :, numpy.newaxis, :] *
...  z[numpy.newaxis, numpy.newaxis, :, :]).sum(axis=-1)
array([[[  36,   92,  148,  204],
        [  92,  244,  396,  548],
        [ 148,  396,  644,  892]],

       [[  92,  244,  396,  548],
        [ 244,  748, 1252, 1756],
        [ 396, 1252, 2108, 2964]]])

- senderle