collapsing all dimensions of numpy array except the first two

Question

I have a variable dimension numpy array, for example it could have the following shapes

(64, 64)
(64, 64, 2, 5)
(64, 64, 40)
(64, 64, 10, 20, 4)

What I want to do is that if the number of dimensions is greater than 3, I want to collapse/stack everything else into the third dimension while preserving order. So, in my above example the shapes after the operation should be:

(64, 64)
(64, 64, 10)
(64, 64, 40)
(64, 64, 800)

Also, the order needs to be preserved. For example, the array of the shape (64, 64, 2, 5) should be stacked as

(64, 64, 2)
(64, 64, 2)
(64, 64, 2)
(64, 64, 2)
(64, 64, 2)

i.e. the 3D slices one after the other. Also, after the operation I would like to reshape it back to the original shape without any permutation i.e. preserve the original order.

One way I could do is multiply all the dimension values from 3 to the last dimension i.e.

shape = array.shape
if len(shape) > 3:
    final_dim = 1
    for i in range(2, len(shape)):
        final_dim *= shape[i]

and then reshape the array. Something like:

array.reshape(64, 64, final_dim)

However, I was first of all not sure if the order is preserved as I want and whether there is a better pythonic way to achieve this?

Answer 1

EDIT: As pointed out in the other answers it is even easier to just provide -1 as the third dimension for reshape. Numpy automatically determines the correct shape then.

I am not sure what the problem here is. You can just use np.reshape and it preserves order. See the following code:

import numpy as np

A = np.random.rand(20,20,2,2,18,5)
print A.shape

new_dim = np.prod(A.shape[2:])
print new_dim
B = np.reshape(A, (A.shape[0], A.shape[1], np.prod(A.shape[2:])))
print B.shape

C = B.reshape((20,20,2,2,18,5))
print np.array_equal(A,C)

The output is:

(20L, 20L, 2L, 2L, 18L, 5L)
360
(20L, 20L, 360L)
True

This accomplishes exactly what you asked for.

Answer 2

reshape accept automatic re-dimension :

a=rand(20,20,8,6,4)
s=a.shape[:2]
if a.ndim>2 : s = s+ (-1,)
b=a.reshape(s)