Reshape 4d numpy array to 2d array while preserving array locations

Question

I have a 4 dimensional numpy array of shape (N, N, Q, Q). So given a row and column index (i, j), mat[i,j] is a QxQ matrix. I want to reshape this array to shape (N*Q, N*Q) such that

array([[[[ 0,  1],
         [ 2,  3]],

        [[ 4,  5],
         [ 6,  7]]],


       [[[ 8,  9],
         [10, 11]],

        [[12, 13],
         [14, 15]]]])

goes to

array([[  0.,   1.,   4.,   5.],
       [  2.,   3.,   6.,   7.],
       [  8.,   9.,  12.,  13.],
       [ 10.,  11.,  14.,  15.]])

You can see that mat[0,0] goes to new_mat[0:2, 0:2]. Currently mat.reshape(N*Q, N*Q) takes mat[0,0] to new_mat[0:4, 0] (which is what I do not want). How can I use reshape or rollaxis or something similar to reshape this array? I eventually want to plot it with imshow, am currently stuck. I figure it's easy to do, I just haven't yet figured it out.

Answer 1

Nevermind, I figured it out. np.swapaxes(1, 2) was the missing piece I needed.

The answer is just to do mat.swapaxes(1, 2).reshape(N*Q, N*Q).

Feel foolish for posting without attempting to figure it out myself for too long, but I'll leave it up so others can benefit from it.