essentially I'm having to do a lot of Matrix-Vector products, where the 3x3 matrices are stored as the last two entries of a (n x n x 3 x 3) numpy array and the vectors are the last column in a corresponging (n x n x 3) array.

Currently I'm looping quite inelegantly over the first two indices and calculating each matrix-vector product separately:

length=1000
x=np.random.rand(length,length,3)
A=np.random.rand(length,length,3,3)
result=np.zeros((length,length,3))

for i in range(0,length):
   for j in range(0,length):
      result[i,j,:]=A[i,j,:,:].dot(x[i,j,:])

This however is quite slow. Is there a more efficient way to avoid this i,j for loop? I'm still trying to learn the numpy methods. ;)

Also, in my real-life case, the matrices are very sparse (>0.999 sparsity), so bonus points for anyone who can give me some hints on how to incorporate that.

Thank you very much!

CodePudding user response：

You could use