Intuition and idea behind reshaping 4D array to 2D array in NumPy

While implementing a Kronecker-product for pedagogical reasons (without using the obvious and readily available np.kron()), I obtained a 4 dimensional array as an intermediate result, which I've to reshape to get the final result.

But, I still can't wrap my head around reshaping these high dimensional arrays. I have this 4D array:

array([[[[ 0,  0],
         [ 0,  0]],

        [[ 5, 10],
         [15, 20]]],


       [[[ 6, 12],
         [18, 24]],

        [[ 7, 14],
         [21, 28]]]])

This is of shape (2, 2, 2, 2) and I'd like to reshape it to (4,4). One might think that this is obvious to do with

np.reshape(my4darr, (4,4))

But, the above reshape does not give me the expected result which is:

array([[ 0,  5,  0, 10],
       [ 6,  7, 12, 14],
       [ 0, 15,  0, 20],
       [18, 21, 24, 28]])

As you can see, all the elements in the expected result are present in the 4D array. I just can't get the hang of doing the reshape correctly as needed. In addition to the answer, some explanation of how to do the reshape for such high dimensional arrays would be really helpful. Thanks!

Solution 1:

General idea for nd to nd transformation

The idea with such nd to nd transformation is using just two things -

• Permute axes (with numpy.transpose or numpy.moveaxis or numpy.rollaxis if the needed permute order is a rolled one or numpy.swapaxes if just two axes need to be swapped) and

• Reshape.

Permute axes : To get the order such that the flattened version corresponds to the flattened version of output. So, if you somehow end up using it twice, look again because you shouldn't.

Reshape : To split the axes or bring the final output to the desired shape. Splitting axes is needed mostly at the start, when the input is of lower-dim and we are needed to split into blocks. Again, you shouldn't need this more than twice.

Hence, generally we would have three steps :

    [ Reshape ]      --->  [ Permute axes ]   --->  [ Reshape ]

 Create more axes             Bring axes             Merge axes
                          into correct order

Back-tracking method

The safest way to solve, given the input and output is through, what one could call as the back-tracking method, i.e. split the axes of the input (when going from smaller nd to bigger nd) or split the axes of the output (when going from bigger nd to smaller nd). The idea with the splitting is to bring the number of dims of the smaller nd one same as the bigger nd one. Then, study the strides of the output and match it up against the input to get the required permute order. Finally, a reshape (default way or C order) might be needed at the end, if the final one is a smaller nd one, to merge axes.

If both input and output are of same number of dims, then we would need to split both and break into blocks and study their strides against each other. In such cases, we should have the additional input parameter of block sizes, but that's probably off-topic.

Example

Let's use this specific case to demonstrate how to apply those strategies. In here, the input is 4D, while output is 2D. So, most probably, we won't need reshape to split. So, we need to start with permuting axes. Since, the final output is not 4D, but a 2D one, we would need a reshape at the end.

Now, the input here is :

In [270]: a
Out[270]: 
array([[[[ 0,  0],
         [ 0,  0]],

        [[ 5, 10],
         [15, 20]]],


       [[[ 6, 12],
         [18, 24]],

        [[ 7, 14],
         [21, 28]]]])

The expected output is :

In [271]: out
    Out[271]: 
    array([[ 0,  5,  0, 10],
           [ 6,  7, 12, 14],
           [ 0, 15,  0, 20],
           [18, 21, 24, 28]])

Also, this is a bigger nd to smaller nd transformation, so the back-tracking method would involve, splitting the output and studying its strides and matching up against the corresponding values in input :

                    axis = 3
                   ---      -->          
                                        
                    axis = 1                    
                   ------>           
axis=2|  axis=0|   [ 0,  5,  0, 10],        

               |   [ 6,  7, 12, 14],
               v  
      |            [ 0, 15,  0, 20],
      v
                   [18, 21, 24, 28]])

Hence, the permuted order needed is (2,0,3,1) :

In [275]: a.transpose((2, 0, 3, 1))
Out[275]: 
array([[[[ 0,  5],
         [ 0, 10]],

        [[ 6,  7],
         [12, 14]]],


       [[[ 0, 15],
         [ 0, 20]],

        [[18, 21],
         [24, 28]]]])

Then, simply reshape to the expected shape :

In [276]: a.transpose((2, 0, 3, 1)).reshape(4,4)
Out[276]: 
array([[ 0,  5,  0, 10],
       [ 6,  7, 12, 14],
       [ 0, 15,  0, 20],
       [18, 21, 24, 28]])

More examples

I dug up my history and found few Q&As based on nd to nd transformations. These could serve as other example cases, albeit with lesser explanation (mostly). As mentioned earlier, at most two reshapes and at most one swapaxes/transpose did the job everywhere. They are listed below :

• Python Reshape 3d array into 2d
• reshape an array using python/numpy
• Merging non-overlapping array blocks
• Conversion from a Numpy 3D array to a 2D array
• how to reshape an N length vector to a 3x(N/3) matrix in numpy using reshape
• Construct image from 4D list
• Reshaping/Combining several sub-matrices to one matrix in multi-dimensional space
• Interlace various small 2D matrices into a bigger one
• how to retrieve every section by 3X3?
• Reshaping 3D Numpy Array to a 2D array
• Iterate in submatrices through a bigger matrix
• Reorganizing a 2D numpy array into 3D
• Numpy change shape from (3, 512, 660, 4) to (3,2048,660,1)
• Numpy: rotate sub matrix m of M
• Split a 3D numpy array into 3D blocks
• Converting 3D matrix to cascaded 2D Matrices
• Rearranging numpy array
• Numpy: Reshape array along a specified axis
• How to construct 2d array from 2d arrays
• How to form a matrix from submatrices?
• Python: Reshape 3D image series to pixel series

Solution 2:

It seems like you're looking for a transpose followed by a reshape.

x.transpose((2, 0, 3, 1)).reshape(np.prod(x.shape[:2]), -1)

array([[ 0,  5,  0, 10],
       [ 6,  7, 12, 14],
       [ 0, 15,  0, 20],
       [18, 21, 24, 28]])

To help you understand why a transposition is needed, let's analyse your incorrectly shaped output (obtained by a single reshape call) to understand why it is incorrect.

A simple 2D reshaped version of this result (without any transposition) looks like this -

x.reshape(4, 4)

array([[ 0,  0,  0,  0],
       [ 5, 10, 15, 20],
       [ 6, 12, 18, 24],
       [ 7, 14, 21, 28]])

Now consider this output with respect to your expected output -

array([[ 0,  5,  0, 10],
       [ 6,  7, 12, 14],
       [ 0, 15,  0, 20],
       [18, 21, 24, 28]])

You'll notice that your actual result is obtained by a Z-like traversal of your incorrectly shaped output -

start
    | /|     /| /|
    |/ |    / |/ |
      /    /    / 
     /    /    /
    | /| /    | /|
    |/ |/     |/ |
                 end

This implies that you must move over the array in varying strides to get your actual result. In conclusion, a simple reshape is not enough. You must transpose the original array, in such a manner that these Z-like elements are made to be contiguous to each other, such that a subsequent reshape call gives you the output you need.

To understand how to transpose correctly, you should trace the elements along the input and figure out what axes you need to jump to get to each one in the output. The transposition follows accordingly. Divakar's answer does a sterling job of explaining this.