convert data from 2d to 3d numpy array

Question

If I have a 2d array like:

array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

And I want to end up with a 3d array like:

array([[[0, 4, 8],
        [1, 5, 9]],

       [[2, 6, 10],
        [3, 7, 11]]])

How should I reshape the array to get what I want?

Answer 1

Reshape and permute axes -

In [11]: a  # Input array
Out[11]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

In [12]: a.reshape(-1,2,2).transpose(1,2,0)
Out[12]: 
array([[[ 0,  4,  8],
        [ 1,  5,  9]],

       [[ 2,  6, 10],
        [ 3,  7, 11]]])

With np.moveaxis -

np.moveaxis(a.reshape(-1,2,2), 0,-1)

Generalizing it and assuming that you want the length along the first axis as half of no. of columns -

In [16]: m,n = a.shape

In [17]: a.reshape(m,-1,2).transpose(1,2,0)
Out[17]: 
array([[[ 0,  4,  8],
        [ 1,  5,  9]],

       [[ 2,  6, 10],
        [ 3,  7, 11]]])

If that length is supposed to be 2 -

In [15]: a.reshape(m,2,-1).transpose(1,2,0)
Out[15]: 
array([[[ 0,  4,  8],
        [ 1,  5,  9]],

       [[ 2,  6, 10],
        [ 3,  7, 11]]])