What is the proper way to create a numpy array of transformation matrices

Question

Given a list of rotation angles (lets say about the X axis):

import numpy as np
x_axis_rotations = np.radians([0,10,32,44,165])

I can create an array of matrices matching these angles by doing so:

matrices = []
for angle in x_axis_rotations:
    matrices.append(np.asarray([[1 , 0 , 0],[0, np.cos(angle), -np.sin(angle)], [0, np.sin(angle), np.cos(angle)]]))
matrices = np.array(matrices)

This will work but it doesn't take advantage of numpy's strengths for dealing with large arrays... So if my array of angles is in the millions, doing it this way won't be very fast.

Is there a better (faster) way to do create an array of transform matrices from an array of inputs?

Answer 1

Here's a direct and simple approach:

c = np.cos(x_axis_rotations)
s = np.sin(x_axis_rotations)
matrices = np.zeros((len(x_axis_rotations), 3, 3))
matrices[:, 0, 0] =  1
matrices[:, 1, 1] =  c
matrices[:, 1, 2] = -s
matrices[:, 2, 1] =  s
matrices[:, 2, 2] =  c

---

timings, for the curious:

In [30]: angles = 2 * np.pi * np.random.rand(1000)

In [31]: timeit OP(angles)
100 loops, best of 3: 5.46 ms per loop

In [32]: timeit askewchan(angles)
10000 loops, best of 3: 39.6 µs per loop

In [33]: timeit divakar(angles)
10000 loops, best of 3: 93.8 µs per loop

In [34]: timeit divakar_oneline(angles)
10000 loops, best of 3: 56.1 µs per loop

In [35]: timeit divakar_combine(angles)
10000 loops, best of 3: 43.9 µs per loop

All are much faster than your loop, so use whichever you like the most :)

Answer 2

You can use linear indexing to help out, like so -

# Get cosine and sine values in one-go
cosv = np.cos(x_axis_rotations)
sinv = np.sin(x_axis_rotations)

# Get size parameter
N = x_axis_rotations.size

# Initialize output array
out = np.zeros((N,3,3))

# Set the first element in each 3D slice as 1
out[:,0,0] = 1

# Calculate the first of positions where cosine valued elements are to be put
idx1 = 4 + 9*np.arange(N)[:,None]

# One by one put those 4 values in 2x2 blocks across all 3D slices
out.ravel()[idx1] = cosv
out.ravel()[idx1+1] = -sinv

out.ravel()[idx1+3] = sinv
out.ravel()[idx1+4] = cosv

Alternatively, you can set the elements in one-go after you have initialized the output array with zeros and set the first element in each slice as 1, like so -

out.reshape(N,-1)[:,[4,5,7,8]] = np.column_stack((cosv,-sinv,sinv,cosv))

Between the above mentioned two approaches, two more middleground approaches could evolve, again put right after initializing with zeros and setting the first element in each 3D slice as 1, like so -

out.reshape(N,-1)[:,[4,8]] = cosv[:,None]
out.reshape(N,-1)[:,[5,7]] = np.column_stack((-sinv[:,None],sinv[:,None]))

The last one would be -

out.reshape(N,-1)[:,[4,8]] = cosv[:,None]
out.reshape(N,-1)[:,5] = -sinv
out.reshape(N,-1)[:,7] = sinv