Reordering numpy arrays

Question

I want to resize my numpy array. The following code snippet works, but it's too slow.

for i in range(image_size): for j in range(image_size): for k in range(3): new_im[k, i, j] = im[i, j, k] 

After that I vectorize new_im:

 new_im_vec = new_im.reshape(image_size**2 * 3) 

However, I don't need new_im, and I only need to get to new_im_vec. Is there a better way to do this? image_size is around 256.

Answer 1

Note rollaxis , a function that moves axes around, allowing you to reorder the array in one command. If im has the form i, j, k

 rollaxis(im, 2) 

should return an array with the form k, i, j.

After that, you can smooth your array, ravel is a clear function for this purpose. Putting it all together, you have a nice one-liner:

 new_im_vec = ravel(rollaxis(im, 2)) 

Answer 2

 new_im = im.swapaxes(0,2).swapaxes(1,2) # First swap i and k, then i and j new_im_vec = new_im.flatten() # Vectorize 

This should be much faster, because swapaxes returns an array view, not a copy element.

And of course, if you want to skip new_im , you can do it on one line, and still only flatten does any copying.

 new_im_vec = im.swapaxes(0,2).swapaxes(1,2).flatten() 

Answer 3

With Einops:

 x = einops.rearrange(x, 'height width color -> color height width') 

Pros:

• You see how the axes were arranged in the input
• you see how the axes are ordered in the output
• you don’t need to think about the steps you need to take (and, for example, you don’t need to remember in which directions the axes rotate)