python - Multiple matrix multiplication

2
votes

In numpy, I have an array of N 3x3 matrices. This would be an example of how I'm storing them (I'm abstracting away the contents):

N = 10
matrices = np.ones((N, 3, 3))

I also have an array of 3-vectors, this would be an example:

vectors = np.ones((N, 3))

I can't seem to figure out how to multiply those via numpy, so as to achieve something like this:

result_vectors = []
for matrix, vector in zip(matrices, vectors):
    result_vectors.append(matrix @ vector)

with the result_vector's shape (upon casting to array) being (N, 3). However, a list implementation is out of the question due to speed.

I've tried np.dot with various transpositions, but the end result didn't get the shape right.

1
pythonperformancenumpymatrixmatrix-multiplication

1 Answers

2
votes

Use np.einsum -

np.einsum('ijk,ik->ij',matrices,vectors)

Steps :

1) Keep the first axes aligned.

2) Sum-reduce the last axes from the input arrays against each other.

3) Let the remainining axes(second axis from matrices) be element-wise multiplied.