Python-Numpy：3D矩阵* 2D矢量快速计算

Question

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