How to "zip" several N-D arrays in Numpy?

Question

The conditions are following:

1) we have a list of N-D arrays and this list is of unknown length M

2) dimensions each arrays are equal, but unknown

3) each array should be splitted along 0-th dimension and resulting elements should be grouped along 1-st dimension of length M and then stacked back along 0-th dimension of the same length it was

4) resulting rank should be N+1 and the lenght of 1-st dimension should be M

Above is the same as zip, but in the world of N-D arrays.

Currently I do the following way:

xs = [list of numpy arrays]
grs = []
for i in range(len(xs[0])):
   gr = [x[i] for x in xs] 
   gr = np.stack(gr)
   grs.append(gr)
grs = np.stack(grs)

Can I write shorter with bulk operations?

UPDATE

Here is what I want

import numpy as np

sz = 2
sh = (30, 10, 10, 3)

xs = []
for i in range(sz):
    xs.append(np.zeros(sh, dtype=np.int))

value = 0

for i in range(sz):
    for index, _ in np.ndenumerate(xs[i]):
        xs[i][index] = value
        value += 1

grs = []
for i in range(len(xs[0])):
   gr = [x[i] for x in xs]
   gr = np.stack(gr)
   grs.append(gr)
grs = np.stack(grs)

print(np.shape(grs))

This code apparantly works correctly, producing arrays of shape (30, 2, 10, 10, 3). Is it possible to avoid loop?

Answer 1

Seems you need to transpose the array with respect to its 1st and 2nd dimension; You can use swapaxes for this:

np.asarray(xs).swapaxes(1,0)

Example:

xs = [np.array([[1,2],[3,4]]), np.array([[5,6],[7,8]])]
grs = []
for i in range(len(xs[0])):
    gr = [x[i] for x in xs] 
    gr = np.stack(gr)
    grs.append(gr)
grs = np.stack(grs)

grs
#array([[[1, 2],
#        [5, 6]],

#       [[3, 4],
#        [7, 8]]])

np.asarray(xs).swapaxes(1,0)
#array([[[1, 2],
#        [5, 6]],

#       [[3, 4],
#        [7, 8]]])

Answer 2

np.stack takes an axis parameter; looking at the shape of grs, I guessed that np.stack(xs, 1) does the same thing.

In [490]: x
Out[490]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])
In [491]: x.shape
Out[491]: (2, 3, 4)
In [494]: xs = [x, x+10, x+100]
In [495]: grs = []
     ...: for i in range(len(xs[0])):
     ...:    gr = [x[i] for x in xs] 
     ...:    gr = np.stack(gr)
     ...:    grs.append(gr)
     ...: grs = np.stack(grs)
     ...: 
In [496]: grs
Out[496]: 
array([[[[  0,   1,   2,   3],
         [  4,   5,   6,   7],
         [  8,   9,  10,  11]],

        [[ 10,  11,  12,  13],
         [ 14,  15,  16,  17],
         [ 18,  19,  20,  21]],

        [[100, 101, 102, 103],
         [104, 105, 106, 107],
         ...
         [116, 117, 118, 119],
         [120, 121, 122, 123]]]])
In [497]: grs.shape
Out[497]: (2, 3, 3, 4)

testing np.stack:

In [499]: np.allclose(np.stack(xs, 1),grs)
Out[499]: True