Numpy array muliplication-CodePudding

I have two input arrays and one output array like this:

M=np.array([[1,2,3],[3,4,5],[6,7,8]])
u=np.array([[1,2,3],[4,5,7],[2,4,9]])
res=np.zeros((3,))

I want to do the following calculation:

for i in range(3):
      res[i]=np.matmul(np.matmul(u[0:,i].T,M),u[0:,i])
#res=array([ 231.,  594., 1957.])

Can I do it without doing for loop since for loop will take alot of time in larger size matrix

so the goal is to acheive quicker method

CodePudding user response：

The most intuitive way to do it is probably via np.einsum:

res = np.einsum('ki,kl,li->i', u, M, u)

CodePudding user response：

Sticking with matmul, you are treating the last dimension of u as a 'batch', the one you iterate over in the loop. If we reshape u to be 3d, with that batch dimension first:

In [430]: u1 = u.T[:,None,:]; u2 = u.T[:,:,None]

In [431]: u1@M@u2
Out[431]: 
array([[[ 231]],

       [[ 594]],

       [[1957]]])

This is (3,1,1), with that batch dimension first, and matrix product in the last 2 dimension. We can squeeze those out:

In [432]: np.squeeze(u1@M@u2)
Out[432]: array([ 231,  594, 1957])

The matmul is working with (n,1,3), (3,3), and (n,3,1) to produce (n,1,1), with sum-of-products on the share 3.