Swapping the dimensions of a numpy array

Question

I would like to do the following:

for i in dimension1:
  for j in dimension2:
    for k in dimension3:
      for l in dimension4:
        B[k,l,i,j] = A[i,j,k,l]

without the use of loops. In the end both A and B contain the same information but indexed differently.

I must point out that the dimension 1,2,3 and 4 can be the same or different. So a numpy.reshape() seems difficult.

Answer 1

The canonical way of doing this in numpy would be to use np.transpose's optional permutation argument. In your case, to go from ijkl to klij, the permutation is (2, 3, 0, 1), e.g.:

In [16]: a = np.empty((2, 3, 4, 5))

In [17]: b = np.transpose(a, (2, 3, 0, 1))

In [18]: b.shape
Out[18]: (4, 5, 2, 3)

Answer 2

Please note: Jaime's answer is better. NumPy provides np.transpose precisely for this purpose.

---

Or use np.einsum; this is perhaps a perversion of its intended purpose, but the syntax is quite nice:

In [195]: A = np.random.random((2,4,3,5))

In [196]: B = np.einsum('klij->ijkl', A)

In [197]: A.shape
Out[197]: (2, 4, 3, 5)

In [198]: B.shape
Out[198]: (3, 5, 2, 4)

In [199]: import itertools as IT    
In [200]: all(B[k,l,i,j] == A[i,j,k,l] for i,j,k,l in IT.product(*map(range, A.shape)))
Out[200]: True

Answer 3

You could rollaxis twice:

>>> A = np.random.random((2,4,3,5))
>>> B = np.rollaxis(np.rollaxis(A, 2), 3, 1)
>>> A.shape
(2, 4, 3, 5)
>>> B.shape
(3, 5, 2, 4)
>>> from itertools import product
>>> all(B[k,l,i,j] == A[i,j,k,l] for i,j,k,l in product(*map(range, A.shape)))
True

or maybe swapaxes twice is easier to follow:

>>> A = np.random.random((2,4,3,5))
>>> C = A.swapaxes(0, 2).swapaxes(1,3)
>>> C.shape
(3, 5, 2, 4)
>>> all(C[k,l,i,j] == A[i,j,k,l] for i,j,k,l in product(*map(range, A.shape)))
True

Answer 4

One can also leverage numpy.moveaxis() for moving the required axes to desired locations. Here is an illustration, stealing the example from Jaime's answer:

In [160]: a = np.empty((2, 3, 4, 5))

# move the axes that are originally at positions [0, 1] to [2, 3]
In [161]: np.moveaxis(a, [0, 1], [2, 3]).shape 
Out[161]: (4, 5, 2, 3)

Answer 5

I would look at numpy.ndarray.shape and itertools.product:

import numpy, itertools
A = numpy.ones((10,10,10,10))
B = numpy.zeros((10,10,10,10))

for i, j, k, l in itertools.product(*map(xrange, A.shape)):
    B[k,l,i,j] = A[i,j,k,l]

By "without the use of loops" I'm assuming you mean "without the use of nested loops", of course. Unless there's some numpy built-in that does this, I think this is your best bet.

Answer 6

Jaime's answer is correct and np.transpose is the solution. but he did not explain how to calculate permutation argument. So, I want to explain how to calculate it correctly here:

If you have a 4 dimensional axis array like: (i, j, k, l), and you want to transpose them into (k, l, i, j), first you must number each axis pythonically, so in given example: "i" is axis=0, "j" is axis =1, "k" is axis=2, and "l" is axis=3. Then you should maintain original axis numbers and arrange those numbers into final desired axes arrangement, as in: k should be placed in axis=0 and k was represented by number 2 so 2 sits in position of axis=0 which gives => (2, , , )... and by maintaining the same condition for other axes the permutation argument of (2,3,0,1) will be obtained. However, It is easy to mistake final axis numbers instead of the original ones in calculating the permutation argument, like mistakenly considering: "i" should become axis=2, "j" should become axis=3, "k" should become axis=0, and "l" should become axis=1, which surprisingly gives (2,3,0,1) correctly again but this manner of calculation can give a wrong answer in other cases. For example, if you were to transpose (i, j, k, l) into (i, k, l, j), the first method calculates the permutation argument correctly as (0,2,3,1) but the second method calculates it mistakenly as (0,3,1,2), as it is shown in below:

import numpy as np
A = np.ones((500,3,512,512))
# Changing A dimensions into (500,512,512,3)
B = np.transpose(A, (0,2,3,1))
C = np.transpose(A, (0,3,1,2))
print(np.shape(B))  # permutation is correct
print(np.shape(C))  # permutation is wrong

which prints:

size of B is: (500, 512, 512, 3)
size of C is: (500, 512, 3, 512)