Rearrange elements in numpy array to plot a 3d array in 2d, preserving grouping

Hello I have a 3d numpy array of shape (nc, nx, ny). I would like to plot it (with matpoltlib) using x and y as the axis, but in such a way that all c values are close to each other and arranged in a rectangle around a central value (you can assume nc is a power of 2).

Thanks to @mad-physicist’s answer I think I managed to come up with a general solution, but it is not very efficient.

My code looks like this:

nc1 , nc2 = 4, 4
(nx, ny, nc) = (2,3,nc1*nc2)
data = np.reshape(np.arange(nx* ny* nc),(nc, nx, ny))

res = np.zeros((nx*nc1,ny*nc2))

for i in range(nx) :
    for j in range(ny) :
        tmp = data[:,i,j].reshape(nc1,nc2)
        res[i*nc1:(i+1)*nc1,j*nc2:(j+1)*nc2] = tmp

I am trying to find a way to avoid the loops using numpy functions. Any suggestion on how to do this?

To clarify, here is a picture of what I have in mind.

Answer

Numpy assumes C order. So when you reshape, if your array is not C contiguous, it generally makes a copy. It also means that you can rearrange the memory layout of an array by say transposing and then reshaping it.

You want to take an array of shape (nc1 * nc2, nx, ny) into one of shape (nx * nc1, ny * nc2). One approach is to split the array into 4D first. Adding an extra dimension will not change the memory layout:

data = data.reshape(nc1, nc2, nx, ny)

Now you can move the dimensions around to get the memory layout you need. This will change your strides without copying memory, since the array is still C contiguous in the last step:

data = data.transpose(2, 0, 3, 1) # nx, nc1, ny, nc2

Notice that we moved nc1 after nx and nc2 after ny. In C order, we want the groups to be contiguous. Later dimensions are close together than earlier ones. So raveling [[1, 2], [3, 4]] produces [1, 2, 3, 4], while raveling [[1, 3], [2, 4]] produces [1, 3, 2, 4].

Now you can get the final array. Since the data is no longer C contiguous after rearranging the dimensions, the following reshape will make the copy that actually rearranges the data:

data = data.reshape(nx * nc1, ny * nc2)

You can write this in one line:

result = data.reshape(nc1, nc2, nx, ny).transpose(2, 0, 3, 1).reshape(nx * nc1, ny * nc2)

In principle, this is as efficient as you can get: the first reshape changes the shape and strides; the transpose swaps around the order of the shape and strides. Only a single copy operation happens during the last reshape.