How do I multiply a matrix to every third dimensional vector in a numpy array?-CodePudding

For example, if I have an array

to_multiply = np.array([[[1,2,3], [2,3,4], [5,6,7]]])

And I want to multiply all the vectors in the 3rd dimension by a matrix. How would I do this without using for loops?

To do this with for loops, I did

for x in range(1):
    for y in range(2):
        to_multiply[y, x] @= matrix

CodePudding user response：

You could use einsum for that.

np.einsum('ijk,kl', to_multiply, m)

Returns and array of the same shape as to_multiply (assuming that m is squared), whose each [i,j] subarray is to_multiply[i,j] @ m.

To be more accurate, what it does is return a 3d-array, whose element [i,j,l] is the result of Σk to_multiply[i,j,k]*m[k,l]

So it computes

for i in range(len(to_multiply)):
    for j in range(to_multiply.shape[1]):
        for l in range(m.shape[1]):
            res[i,j,l]=0
            for k in range(to_multiply.shape[2]): # that must be same as m.shape[0]
                res[i,j,l]  = to_multiply[i,j,k]*m[k,l]

but, of course, that computation is not done in python, but in internal code of einsum, so it is fast.

CodePudding user response：

From the documentation of numpy.matmul,

If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly.

Hence

to_multiply = np.matmul(to_multiply, matrix)

is same as

to_multiply = np.einsum('ijk,kl', to_multiply, matrix)

which is again same as

to_multiply @ matrix

for this question.