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

Question

So I have a 3D data-set (x,y,z), and i want to sum over one of the axes (x) with a set of weights, w = w(x). The start and end index i am summing over is different for every (y,z), I have solved this by masking the 3D-array. The weights are constant with regard to the two variables i am not summing over. Both answers regarding implementation and mathematics are appreciated (is there a smart linalg. way of doing this?).

I have a 3D masked array (A) of shape (x,y,z) and a 1D array (t) of shape (x,). Is there a good way to multiply every (y,z) element in A with the corresponding number in t without expanding t to a 3D array? My current solution is using np.tensordot to make a 3D array of the same shape as A, that holds all the t-values, but it feels very unsatisfactory to spend runtime building the "new_t" array, which is essensially just y*z copies of t.

Example of current solution:

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)

In essence i want to perform t*A[:,i,j] for all i,j with the shortest possible runtime, preferably without using too many other libraries than numpy and scipy.

Another way of producing desired output (again, with far too high run time):

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

inspired by @phipsgabler comment

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

Thanks for good answers! Using tensordot like @alyhosny proposed worked, but replacing masked values with zeros using

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

before summing with einsum (thanks @phipsgabler) gave half the run time. Final code:

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