How to Transpose each element in a 3D np array

Question

Given a 3D array a, I want to call np.tranpose on each of the element in its first index.
For example, given the array:

array([[[1, 1, 1, 1],
        [1, 1, 1, 1],
        [1, 1, 1, 1]],

       [[2, 2, 2, 2],
        [2, 2, 2, 2],
        [2, 2, 2, 2]],

       [[3, 3, 3, 3],
        [3, 3, 3, 3],
        [3, 3, 3, 3]])

I want:

array([[[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],

       [[2, 2, 2],
        [2, 2, 2],
        [2, 2, 2],
        [2, 2, 2]],

       [[3, 3, 3],
        [3, 3, 3],
        [3, 3, 3],
        [3, 3, 3]]])

Essentially I want to transpose each element inside the array. I tried to reshape it but I can't find a good way of doing it. Looping through it and calling transpose on each would be too slow. Any advice?

Answer 1

You can use the built-in numpy transpose method and directly specify the axes to transpose

>>> a = np.array([[[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]],
                  [[2, 2, 2, 2], [2, 2, 2, 2], [2, 2, 2, 2]],
                  [[3, 3, 3, 3], [3, 3, 3, 3], [3, 3, 3, 3]]])

>>> print(a.transpose((0, 2, 1)))

[[[1 1 1]
  [1 1 1]
  [1 1 1]
  [1 1 1]]

 [[2 2 2]
  [2 2 2]
  [2 2 2]
  [2 2 2]]

 [[3 3 3]
  [3 3 3]
  [3 3 3]
  [3 3 3]]]