# Python - Numpy : 3D matrix * 2D vector fast calculation - Python

TAGS :
Viewed: 3 - Published at: a few seconds ago

#### [ Python - Numpy : 3D matrix * 2D vector fast calculation ]

Hello everyone,

Here is what I want to do. I've got two arrays :

• rotation_matrices that contains 50 two dimensional rotation matrices. Each rotation matrix is shaped as (2,2). Thus, rotation_matrices is shaped as (2,2,50).
• vectors that contains 50 two dimensional vectors. Thus, it is shaped as (2,50).

I want (if it exists) a one-line numpy operation that gives me the (2,50) array that contains the rotated vectors, let's call it rotated_vectors. What I mean is that the k-ith element of rotated_vectors contains the product of the k-ith rotation matrix with the k-ith vector.

For the moment, I have come up with the loop that follows :

``````for ind,elt in enumerate(np.arange(nb_of_vectors)):
rotated_vector[ind] = np.dot( rotation_matrices[:,:,ind], vectors[:,ind] )
``````

I think there is room for improvement. If you have any suggestion, you are welcome.

Thank you for your time.

Jagaral

# Answer 1

For such reductions that require alignment along one or more axes, one can use `np.einsum` -

``````rotated_vector = np.einsum('ijk,jk->ki',rotation_matrices,vectors)
``````

Please note that the output would be of shape `(N,2)`, where `N` is the number of vectors. If instead, you were looking to have an output of shape `(2,N)` and which would have required the original code to be : `rotated_vector[:,ind] = np.dot(...)` instead, just edit the output string notation to `ik` instead of `ki`.

Runtime test -

``````In [24]: def org_app(rotation_matrices,vectors):
...:     nb_of_vectors = vectors.shape[1]
...:     r = np.zeros((nb_of_vectors,2))
...:     for ind,elt in enumerate(np.arange(nb_of_vectors)):
...:         r[ind] = np.dot( rotation_matrices[:,:,ind], vectors[:,ind] )
...:     return r
...:

In [25]: # Input arrays
...: rotation_matrices = np.random.rand(2,2,50)
...: vectors = np.random.rand(2,50)
...:

In [26]: out1 = org_app(rotation_matrices,vectors)

In [27]: out2 = np.einsum('ijk,jk->ki',rotation_matrices,vectors)

In [28]: np.allclose(out1,out2) # Verify results
Out[28]: True

In [29]: %timeit org_app(rotation_matrices,vectors)
10000 loops, best of 3: 196 µs per loop

In [30]: %timeit np.einsum('ijk,jk->ki',rotation_matrices,vectors)
100000 loops, best of 3: 5.12 µs per loop
``````

That proves again why iterating in NumPy is basically just bad!

# Answer 2

Your axes are in an unusual order. First you'll want to put the matrix axes last:

``````rotation_matrices = np.rollaxis(rotation_matrices, -1)  # shape (50, 2, 2)
vectors = np.rollaxis(vectors, -1)                      # shape (50, 2)
``````

Which would allow you to make your existing loop more readable:

``````for ind in np.arange(nb_of_vectors):
rotated_vector[ind] = np.dot(rotation_matrices[ind], vectors[ind])
``````

But instead, you can use the matrix multiply operator (or `np.matmul` in python < 3.5)

``````rotated_vectors = (a @ vectors[...,None])[...,0]
# rotated_vectors = np.matmul(a, vectors[...,None])[...,0]
``````

The `[...,None]` converts the array of vectors (shape `(n,)` into an array of column matrices (shape `(n, 1)`), and the trailing `[...,0]` converts the column matrices back into vectors

# Answer 3

This is an explicit formula version

``````result = np.array([vectors[0,:]*rotation_matrices[0,0,:] +
vectors[1,:]*rotation_matrices[0,1,:],
vectors[0,:]*rotation_matrices[1,0,:] +
vectors[1,:]*rotation_matrices[1,1,:]]).transpose()
``````

much (14x) faster than your original code but slower (2.6x) than `einsum` on my machine