Reshape 4D numpy array into 3D

Question

I have a numpy array with the following dimensions - (256, 128, 4, 200) - basically the first two can form an image, the third is channels and the fourth is frames ("time instances"). How can I reshape the array so the frames are "stacked" one after the other, in other words the array would have a shape of (256, 128*200, 4)? It is important that the concatenating is frame-wise, so the order of the values in a frame is preserved.

Essentially, what is needed is to optimize:

data_new = data[:, :, :, 0]
for i in range(1, data.shape[3]):
    data_new = np.concatenate((data_new, data[:, :, :, i]), axis=1)

Answer 1

Permute axes with np.transpose and reshape -

m,n = data.shape[::2]
data_new = data.transpose(0,3,1,2).reshape(m,-1,n)

Or roll-axis and reshape -

data_new = np.rollaxis(data,3,1).reshape(m,-1,n)

Runtime test -

In [40]: data = np.random.randint(0,9,(256,128,4,200))

In [46]: %%timeit
    ...: data_new = data[:, :, :, 0]
    ...: for i in range(1, data.shape[3]):
    ...:     data_new = np.concatenate((data_new, data[:, :, :, i]), axis=1)
    ...: 
1 loop, best of 3: 3.56 s per loop

In [49]: m,n = data.shape[::2]

In [50]: %timeit data.transpose(0,3,1,2).reshape(m,-1,n)
10 loops, best of 3: 47.1 ms per loop

In [51]: %timeit np.rollaxis(data,3,1).reshape(m,-1,n)
10 loops, best of 3: 46.8 ms per loop

Thus, 76x+ speedup is the vectorization profit.