Reshaping array of matrices in Python

Question

I have a Numpy array X of n 2x2 matrices, arranged so that X.shape = (2,2,n), that is, to get the first matrix I call X[:,:,0]. I would like to reshape X into an array Y such that I can get the first matrix by calling Y[0] etc., but performing X.reshape(n,2,2) messes up the matrices. How can I get it to preserve the matrices while reshaping the array?

I am essentially trying to do this:

import numpy as np
Y = np.zeros([n,2,2])
for i in range(n):
    Y[i] = X[:,:,i]

but without using the for loop. How can I do this with reshape or a similar function?

(To get an example array X, try X = np.concatenate([np.identity(2)[:,:,None]] * n, axis=2) for some n.)

Answer 1

numpy.moveaxis can be used to take a view of an array with one axis moved to a different position in the shape:

numpy.moveaxis(X, 2, 0)

numpy.moveaxis(a, source, destination) takes a view of array a where the axis originally at position source ends up at position destination, so numpy.moveaxis(X, 2, 0) makes the original axis 2 the new axis 0 in the view.

There's also numpy.transpose, which can be used to perform arbitrary rearrangements of an array's axes in one go if you pass it the optional second argument, and numpy.rollaxis, an older version of moveaxis with a more confusing calling convention.

Answer 2

Use swapaxis:

Y = X.swapaxes(0,2)