Rearranging axes in numpy?

Question

I have an ndarray such as

>>> arr = np.random.rand(10, 20, 30, 40)
>>> arr.shape
(10, 20, 30, 40)

whose axes I would like to swap around into some arbitrary order such as

>>> rearranged_arr = np.swapaxes(np.swapaxes(arr, 1,3), 0,1)
>>> rearranged_arr.shape
(40, 10, 30, 20)

Is there a function which achieves this without having to chain together a bunch of np.swapaxes?

Answer 1

There are two options: np.moveaxis and np.transpose.

• np.moveaxis(a, sources, destinations) docs

  This function can be used to rearrange specific dimensions of an array. For example, to move the 4th dimension to be the 1st and the 2nd dimension to be the last:

  >>> rearranged_arr = np.moveaxis(arr, [3, 1], [0, 3])
  >>> rearranged_arr.shape
  (40, 10, 30, 20)
  

  This can be particularly useful if you have many dimensions and only want to rearrange a small number of them. e.g.

  >>> another_arr = np.random.rand(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
  >>> np.moveaxis(another_arr, [8, 9], [0, 1]).shape
  (8, 9, 0, 1, 2, 3, 4, 5, 6, 7)
  

• np.transpose(a, axes=None) docs

  This function can be used to rearrange all dimensions of an array at once. For example, to solve your particular case:

  >>> rearranged_arr = np.transpose(arr, axes=[3, 0, 2, 1])
  >>> rearranged_arr.shape
  (40, 10, 30, 20)
  

  or equivalently

  >>> rearranged_arr = arr.transpose(3, 0, 2, 1)
  >>> rearranged_arr.shape
  (40, 10, 30, 20)
  

Answer 2

In [126]: arr = np.random.rand(10, 20, 30, 40)                                                               
In [127]: arr.transpose(3,0,2,1).shape                                                                       
Out[127]: (40, 10, 30, 20)