Split a 3D numpy array into smaller 3D arrays

Question

I have a 3D np.array

arr = np.array([ 
                [ [0, 205, 25], [210, 150, 30], [0, 0, 0], [1, 2, 3], [4, 5, 6], [7, 8, 9] ],
                [ [0, 255, 0], [255, 40, 0], [0, 0, 200], [7, 8, 9], [10, 11, 12], [120, 51, 58] ],
                [ [0, 0, 30], [0, 40, 0], [200, 100, 20], [12, 13, 14], [15, 16, 17], [13, 78, 84], ],
                [ [0, 205, 25], [210, 150, 30], [0, 0, 0], [1, 2, 3], [4, 5, 6], [7, 8, 9] ],
                [ [0, 255, 0], [255, 40, 0], [0, 0, 200], [7, 8, 9], [10, 11, 12], [120, 51, 58] ],
                [ [0, 0, 30], [0, 40, 0], [200, 100, 20], [12, 13, 14], [15, 16, 17], [13, 78, 84], ],
                [ [0, 205, 25], [210, 150, 30], [0, 0, 0], [1, 2, 3], [4, 5, 6], [7, 8, 9] ],
                [ [0, 255, 0], [255, 40, 0], [0, 0, 200], [7, 8, 9], [10, 11, 12], [120, 51, 58] ],
                [ [0, 0, 30], [0, 40, 0], [200, 100, 20], [12, 13, 14], [15, 16, 17], [13, 78, 84], ],
              ])

And I need to split it to 3x2x3 3D arrays

[ [0, 205, 25], [210, 150, 30],    [0, 0, 0], [1, 2, 3],             [4, 5, 6], [7, 8, 9] ],
[ [0, 255, 0],  [255, 40, 0],      [0, 0, 200], [7, 8, 9],           [10, 11, 12], [120, 51, 58] ],
[ [0, 0, 30],   [0, 40, 0],        [200, 100, 20], [12, 13, 14],     [15, 16, 17], [13, 78, 84], ],

[ [0, 205, 25], [210, 150, 30],    [0, 0, 0], [1, 2, 3],             [4, 5, 6], [7, 8, 9] ],
[ [0, 255, 0],  [255, 40, 0],      [0, 0, 200], [7, 8, 9],           [10, 11, 12], [120, 51, 58] ],
[ [0, 0, 30],   [0, 40, 0],        [200, 100, 20], [12, 13, 14],     [15, 16, 17], [13, 78, 84], ],

[ [0, 205, 25], [210, 150, 30],    [0, 0, 0], [1, 2, 3],             [4, 5, 6], [7, 8, 9] ],
[ [0, 255, 0],  [255, 40, 0],      [0, 0, 200], [7, 8, 9],           [10, 11, 12], [120, 51, 58] ],
[ [0, 0, 30],   [0, 40, 0],        [200, 100, 20], [12, 13, 14],     [15, 16, 17], [13, 78, 84], ],

to get a 4D array with these 3D blocks I've selected by spaces. Zero element must be

[ 
    [[0, 205, 25], [210, 150, 30]],
    [[0, 255, 0], [255, 40, 0]],
    [[0, 0, 30], [0, 40, 0]] 
]

and so on.

I've read this question but still don't undersatand how to do this (Why we need to reshape, transpose and reshape again and what a magical numbers in transpose() ). I could try to write my own function but I want to know how to do it native way.

Answer 1

You can reshape and transpose it

arr.reshape(3, 3, 3, 2, 3).transpose(2, 0, 1, 3, 4)
# array([[[[[  0, 205,  25],
#           [210, 150,  30]],
# 
#          [[  0, 255,   0],
#           [255,  40,   0]],
# 
#          [[  0,   0,  30],
#           [  0,  40,   0]]],
# 
# 
#         [[[  0, 205,  25],
#           [210, 150,  30]],
# 
#          [[  0, 255,   0],
#           [255,  40,   0]],
# 
#          [[  0,   0,  30],
#           [  0,  40,   0]]],
# 
# 
# ...