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!!
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]])
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]])
