Numpy reshape preserving some dimensions

Question

I want this

# assume each 2x2 as assigned number on right

[[[False  True]            # 1  
  [False False]]

 [[False False]            # 2
  [False  True]]

 [[ True False]            # 3
  [False False]]

 [[False  True]            # 4
  [False False]]]

to be reshape into

[[[[False  True False False]
   [False False False  True]
   [ True False False  True]
   [False False False False]]]]

as in

[[[[False  True | False False]                # 1 | 2
   [False False | False  True]                ___   ___
     ----------   -----------                  
   [ True False | False  True]                # 3 | 4
   [False False | False False]]]]

but I get when using arr.reshape(1,1,4,4)

[[[[False  True False False]                   # 1 flat
   [False False False  True]                   # 2 flat
   [ True False False False]                   # 3 flat
   [False  True False False]]]]                # 4 flat

Notice that each 2x2 is flattened. I want numpy to reshape such that 2x2 remains the same while the dimensions before that are adjusted. How do I do it?

---

EDIT: shape of arr is m,n,r and m could be odd

EDIT2:

Case 9x2x2 into 1x1x6x6

What I have:

[[[False False]
  [False  True]]

 [[False False]
  [False  True]]

 [[ True False]
  [False False]]

 [[False False]
  [False  True]]

 [[False False]
  [ True False]]

 [[False False]
  [ True False]]

 [[False  True]
  [False False]]

 [[False False]
  [ True False]]

 [[ True False]
  [False False]]]

Expected:

[[[[False False | False False |  True False]
   [False  True | False  True | False False]
    -----------   -----------   ----------
   [False False | False False | False False]
   [False  True |  True False |  True False]
     ----------   -----------   -----------
   [False  True | False False |  True False]
   [False False |  True False | False False]]]]

What I got:

[[[[False False False  True False False]          # 1 flattened; half of 2
   [False  True  True False False False]          # rem half of 2; flattened 3
   [False False False  True False  True]          # ...
   [False False False  True False False]
   [False False  True False False  True]
   [False False  True False False False]]]]

Answer 1

The important piece of the puzzle was that the first axis is a square number and we are splitting it by that square-root number, giving us a 4D array. If the first axis isn't a squared number, we would need another input argument telling us the number of blocks to be kept along the columns or rows in the final output. After splitting, swap axes 1 and 2 and finally reshape -

m = int(np.rint(np.sqrt(a.shape[0])))
out = a.reshape(m,m,2,2).swapaxes(1,2).reshape(m*2,-1)

Sample run -

1] Input :

In [69]: a
Out[69]: 
array([[[False, False],
        [False,  True]],

       [[False, False],
        [False,  True]],

       [[ True, False],
        [False, False]],

       [[False, False],
        [False,  True]],

       [[False, False],
        [ True, False]],

       [[False, False],
        [ True, False]],

       [[False,  True],
        [False, False]],

       [[False, False],
        [ True, False]],

       [[ True, False],
        [False, False]]], dtype=bool)

2] Output :

In [70]: m = int(np.sqrt(a.shape[0]))

In [71]: a.reshape(m,m,2,2).swapaxes(1,2).reshape(m*2,-1)
Out[71]: 
array([[False, False, False, False,  True, False],
       [False,  True, False,  True, False, False],
       [False, False, False, False, False, False],
       [False,  True,  True, False,  True, False],
       [False,  True, False, False,  True, False],
       [False, False,  True, False, False, False]], dtype=bool)

Answer 2

Use einops for such cases.

In: x

array([[[False, False],
        [False,  True]],

       [[False, False],
        [False,  True]],

       [[ True, False],
        [False, False]],

       [[False, False],
        [False,  True]],

       [[False, False],
        [ True, False]],

       [[False, False],
        [ True, False]],

       [[False,  True],
        [False, False]],

       [[False, False],
        [ True, False]],

       [[ True, False],
        [False, False]]])

In: einops.rearrange(x, '(h w) h2 w2 -> (h h2) (w w2)', h=3)

array([[False, False, False, False,  True, False],
       [False,  True, False,  True, False, False],
       [False, False, False, False, False, False],
       [False,  True,  True, False,  True, False],
       [False,  True, False, False,  True, False],
       [False, False,  True, False, False, False]])

See einops docs, basically it allows all kinds of rearrangements being written in explicit way.