Easy way to fold a multidimensional NumPy array

Question

I have a NumPy 2D array:

a = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])

where each column of this array (i.e.: [1,3, 5, 7], and [2, 4, 6, 8]) need to be transformed to a matrix of given size M1xM2 (using order='F' when reshaping). In this case, M1 = M2 = 2. So, my desired output would be:

b = np.array([[[1, 5], [3, 7]], [[2, 6], [4, 8]]]).

I can easily achieve this by iterating over the columns. However, the number of columns can be any and the initial 2D array can be of up to 8 dimensions. How can I easily extend this solution to be used for more dimensions?

I suspect this is a common procedure and there is a built-in function to solve it, but haven't been able to find it.

What do you mean "up to 8 dimensions"? Can you give examples of these higher dimensional arrays? — Nils Werner, Dec 11 '20 at 13:21
@Daniela - I also would like to try to solve this in higher dimensions (perhaps in another question), but I'm not completely shure what an example would look like. — Michael Szczesny, Dec 11 '20 at 13:24

score 3 · Accepted Answer · answered Dec 11 '20 at 13:19

3

It is as simple as reshape, then transpose:

a.reshape(2, 2, -1).T
# array([[[1, 5],
#         [3, 7]],
# 
#        [[2, 6],
#         [4, 8]]])

answered Dec 11 '20 at 13:19

Nils Werner

34,832
7
76
98

I had to add order='F' in my case, but this answer worked correctly. – 6659081 Dec 14 '20 at 16:20

score 2 · Answer 2 · answered Dec 11 '20 at 13:12

2

You can use reshape, swapaxes, reshape. @divakar posted a detailed explanation.

a.T.reshape(2,2,2,1).swapaxes(1,2).reshape(2,2,-1)

Out:

array([[[1, 5],
        [3, 7]],

       [[2, 6],
        [4, 8]]])

answered Dec 11 '20 at 13:12

Michael Szczesny

4,911
5
15
32

You can simplify it to `a.reshape(2, 2, -1).transpose(2, 1, 0)` – Nils Werner Dec 11 '20 at 13:15
Yes that's true, but for me it's hard to find that solution. I memorized the pattern in my answer. Can you post your solution as an answer? – Michael Szczesny Dec 11 '20 at 13:18

Easy way to fold a multidimensional NumPy array

2 Answers2