将形状为 (x,y,z) 和 (x,) 的 numpy 数组相乘和求和

Question

所以我有一个 3D 数据集 (x,y,z)，我想用一组权重对一个轴 (x) 求和，w = w(x)。 我求和的开始和结束索引对于每个 (y,z) 都是不同的，我通过屏蔽 3D 数组解决了这个问题。 对于我没有总结的两个变量，权重是恒定的。 关于实现和数学的两个答案都受到赞赏（是否有一个聪明的 linalg。这样做的方法？）。

我有一个形状为 (x,y,z) 的 3D 掩码数组 (A) 和一个形状为 (x,) 的一维数组 (t)。 有没有一种好方法可以将 A 中的每个 (y,z) 元素与 t 中的相应数字相乘，而无需将 t 扩展为 3D 数组？ 我目前的解决方案是使用 np.tensordot 制作一个与 A 形状相同的 3D 数组，该数组包含所有 t 值，但是花费运行时间构建“new_t”数组感觉非常不令人满意，这基本上只是 y*z t 的副本。

当前解决方案示例：

a1 = np.array([[1,2,3,4],
               [5,6,7,8],
               [9,10,11,12]])

a2 = np.array([[0,1,2,3],
               [4,5,6,7],
               [8,9,10,11]])

#note: A is a masked array, mask is a 3D array of bools
A = np.ma.masked_array([a1,a2],mask)
t = np.array([10,11])

new_t = np.tensordot(t, np.ones(A[0].shape), axes = 0)
return np.sum(A*new_t, axis=0)

本质上，我想以尽可能短的运行时间为所有 i,j 执行 t*A[:,i,j]，最好不使用除 numpy 和 scipy 之外的太多其他库。

另一种产生所需输出的方法（同样，运行时间太长）：

B = [[t*A[:,i,j] for j in range(A.shape[2])] for i in range(A.shape[1])]
return np.sum(B,axis=2)

Answer 1

灵感来自@phipsgabler 评论

arr1 = np.tensordot(A.T,t,axes=1).T
arr1
array([[ 10,  31,  52,  73],
       [ 94, 115, 136, 157],
       [178, 199, 220, 241]])

Answer 2

谢谢你的好答案！ 使用像@alyhosny 建议的 tensordot 工作，但使用零替换掩码值

A = np.ma.MaskedArray.filled(A,0)

在用 einsum 求和之前（感谢@phipsgabler）给了一半的运行时间。 最终代码：

A = np.ma.MaskedArray(A,mask)
A = np.ma.MaskedArray.filled(A,0)
return np.einsum('ijk,i->jk',A,t)