Как "сжать" несколько массивов ND в Numpy?

Question

Условия следующие:

1) у нас есть список массивов ND, и этот список имеет неизвестную длину M

2) размеры каждого массива равны, но неизвестны

3) каждый массив должен быть разбит по 0-му измерению, а полученные элементы должны быть сгруппированы по 1-му измерению длины M а затем сложены обратно вдоль 0-го измерения той же длины, что было

4) итоговый ранг должен быть N+1 и длина 1-го измерения должна быть M

Выше так же, как zip, но в мире ND массивов.

В настоящее время я делаю следующим образом:

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)

Можно ли писать короче с массовыми операциями?

ОБНОВИТЬ

Вот что я хочу

импортировать numpy как 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))

Этот код явно работает правильно, создавая массивы формы (30, 2, 10, 10, 3), Можно ли избежать петли?

Answer 1

Кажется, вам нужно транспонировать массив относительно его 1-го и 2-го измерения; Ты можешь использовать swapaxes за это:

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

Пример:

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 принимает параметр оси; глядя на форму grsЯ догадался что np.stack(xs, 1) делает то же самое.

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)

тестирование np.stack:

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