简体   繁体   中英

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

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)

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]])

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)

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM