如何在numpy中将这个三重循环的二重数组矢量化？

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,CST输出是形状(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)

Answer 1

如果你使用numpy> 1.6,那就有很棒的np.einsum功能:

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

这相当于您的循环版本.一旦你在真正的问题中达到阵列的大小,我不确定这对效率是否公平(当我移动到那些尺寸时,我实际上在我的机器上遇到了段错误).我经常喜欢这些问题的另一个解决方案是使用cython重写方法.

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表示下一对-last等等).这样,您不必知道数组的原始形状; 只要他们的最后一个轴兼容,这将共同播出前三个:

>>> (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]]])