沿着一个维度的1D阵列和3D阵列的高效产品 - NumPy

时间:2016-03-22 09:16:46

标签: python arrays performance numpy multiplication

我有两个numpy数组:

  • 名为t of shape(70L,)的一维数组,其元素名为let s say ti
  • 一个名为I的3D阵列,具有形状(70L,1024L,1024L),每个元素称为Ii。因此,Ii具有尺寸(1024L,1024L)

我想在第一维上制作两个数组的乘积,即:

tI = t1*I1,t2*I2,...,tN*IN

再次获得一个新的维度数组(70L,1024L,1024L),然后沿第一维获取总和,以获得维度数组(1024L,1024L):

tsum = t1*I1 + t2*I2 + ... +tN*IN

目前我对以下内容感到满意:

tI = np.asarray([t[i]*I[i,:,:] for i in range(t.shape[0])])
tsum = np.sum(tI,axis=0)

但是我的阵列的尺寸越来越大,这将会有点慢。 我想知道是否存在numpy或scipy函数,是否针对该特定任务进行了更优化?

先感谢任何链接或信息。

Greg

2 个答案:

答案 0 :(得分:3)

您可以使用np.tensordot -

np.tensordot(t,I, axes=([0],[0]))

您还可以使用np.einsum -

np.einsum('i,ijk->jk',t,I)

运行时测试和输出验证 -

In [21]: def original_app(t,I):
    ...:     tI = np.asarray([t[i]*I[i,:,:] for i in range(t.shape[0])])
    ...:     tsum = np.sum(tI,axis=0)
    ...:     return tsum
    ...: 

In [22]: # Inputs with random elements
    ...: t = np.random.rand(70,)
    ...: I = np.random.rand(70,1024,1024)
    ...: 

In [23]: np.allclose(original_app(t,I),np.tensordot(t,I, axes=([0],[0])))
Out[23]: True

In [24]: np.allclose(original_app(t,I),np.einsum('i,ijk->jk',t,I))
Out[24]: True

In [25]: %timeit np.tensordot(t,I, axes=([0],[0]))
1 loops, best of 3: 110 ms per loop

In [26]: %timeit np.einsum('i,ijk->jk',t,I)
1 loops, best of 3: 201 ms per loop

答案 1 :(得分:2)

Divakar提供最好(最有效)的答案。为了完整性'另一种方法是使用Numpy's broadcasting capabilities

(t[:,np.newaxis,np.newaxis]*I).sum(axis=0)

通过向t添加两个轴,可以进行广播,并且可以使用常规的Numpy操作,对于某些操作可能更具可读性。