What is the correct way to create an array of numpy transform matrices

Question

Given a list of rotation angles (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 corresponding to these corners by doing this:

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 will not take advantage of the numpy capacities for working with large arrays ... Therefore, if my array of angles is in millions, doing so will not be very fast.

Is there a better (faster) way to 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

Everything is much faster than your loop, so use what you like best :)

Answer 2

linear indexing, , :

# 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

zeros 1, :

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

, 1, , -

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

-

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