沿一维的 1D 数组和 3D 数组的有效乘积 - NumPy

我有两个 numpy 数组：

• 一个名为 t 的一维数组，形状为 (70L,)，其元素名为 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

您可以使用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