Assume that I have 9 arrays (A, B, C, .. J) of size N. I want to create a new array of N 3x3 matrices such that e.g.
matrices[i] = [[A[i], B[i], C[i]],
[D[i], E[i], F[i]],
[G[i], H[i], J[i]]]
A simple solution is to add each entry to the array matrices
in a for-loop as:
for i in range(len(matrices)):
matrices[i] = [[A[i], B[i], C[i]],
[D[i], E[i], F[i]],
[G[i], H[i], J[i]]]
Anybody got some tips on how this can be done in a faster, vectorized way avoiding the for-loop? If there exists some smart indexing operations or something.
One approach would be to stack those in columns with np.column_stack
and reshape with np.reshape
-
np.column_stack((A,B,C,D,E,F,G,H,J)).reshape(-1,3,3)
Concatenating with np.concatenate
is known to be much faster, so using it with 2D transpose
and reshaping -
np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(9,-1).T.reshape(-1,3,3)
Another with np.concatenate
, 3D transpose
and reshaping -
np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(3,3,-1).transpose(2,0,1)
Runtime tests -
In [59]: # Setup input arrays
...: N = 1000
...: A = np.random.randint(0,9,(N,))
...: B = np.random.randint(0,9,(N,))
...: C = np.random.randint(0,9,(N,))
...: D = np.random.randint(0,9,(N,))
...: E = np.random.randint(0,9,(N,))
...: F = np.random.randint(0,9,(N,))
...: G = np.random.randint(0,9,(N,))
...: H = np.random.randint(0,9,(N,))
...: J = np.random.randint(0,9,(N,))
...:
In [60]: %timeit np.column_stack((A,B,C,D,E,F,G,H,J)).reshape(-1,3,3)
10000 loops, best of 3: 84.4 µs per loop
In [61]: %timeit np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(9,-1).T.reshape(-1,3,3)
100000 loops, best of 3: 15.8 µs per loop
In [62]: %timeit np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(3,3,-1).transpose(2,0,1)
100000 loops, best of 3: 14.8 µs per loop