将形状为 (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" 感觉非常不满意数组，本质上只是 t.

的 y*z 个副本

当前解决方案示例：

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 和 [=24 之外的太多其他库=].

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

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)

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

Multiply and sum numpy arrays with shapes (x,y,z) and (x,)

python

arrays

optimization

numpy

linear-algebra