將形狀為 (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)