How to rearrange an array by subarray elegantly in numpy?

Question

Let's say I have a 3-D array:

[[[0,1,2],
  [0,1,2],
  [0,1,2]],

 [[3,4,5],
  [3,4,5],
  [3,4,5]]]

And I want to rearrange this by the columns:

[[0,1,2,3,4,5],
 [0,1,2,3,4,5],
 [0,1,2,3,4,5]]

What would be an elegant python numpy code for doing this for essentially a 3-D np.array of arbitrary shape and depth? Could there be a fast method that bypasses for loop? All the approaches I made were terribly adhoc and brute they were basically too slow and useless...

Thanks!!

I guess one could say I'm trying to ravel this 3D array by the column index... — StatsNoob, Jan 27 '17 at 23:21

score 4 · Answer 1 · answered Jun 27 '19 at 18:11

Using einops:

einops.rearrange(a, 'x y z -> y (x z) ')

And I would recommend to give meaningful names to axes (instead of x y z) depending on the context (e.g. time, height, etc.). This will make it easy to understand what the code does

In : einops.rearrange(a, 'x y z -> y (x z) ')
Out:
array([[0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5]])

score 3 · Accepted Answer · answered Jan 27 '17 at 23:21

Swap axes and reshape -

a.swapaxes(0,1).reshape(a.shape[1],-1)

Sample run -

In [115]: a
Out[115]: 
array([[[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]],

       [[3, 4, 5],
        [3, 4, 5],
        [3, 4, 5]]])

In [116]: a.swapaxes(0,1).reshape(a.shape[1],-1)
Out[116]: 
array([[0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5]])

How to rearrange an array by subarray elegantly in numpy?

2 Answers2