如何正确堆叠 numpy arrays？

Question

I am having trouble understanding how data is being stacked in a numpy array and why I cannot match the last data that I added to an array with the last generated data. Here is a MWE:

import numpy as np
np.random.seed(1)

# build storage
container = []

# gen data
x = np.random.random((13, 1, 64, 768))

# add to container
container.append(x)

# gen data
x2 = np.random.random((13, 1, 64, 768))

# add to container
container.append(x2)

# convert to np array
container = np.asarray(container)

# reshape to [13, 2, 64, 768]
container = container.reshape(13, 2, 64, 768)

# check that the last generated data matches the last appended data
assert np.all(x2.flatten() == container[:, -1, :, :].flatten()), 'not a match'

Answer 1

Instead of stacking manually with appending to lists and then reshaping you could use the vstack or the concatenate function of numpy.

# gen data
x1 = np.random.random((13, 1, 64, 768))
x2 = np.random.random((13, 1, 64, 768))

container = np.vstack((x1,x2))
assert np.all(x2.flatten()) == np.all(container[:, -1, :, :].flatten()), 'not a match'

To answer your question: your code does work, just make sure to put np.all() at both sides of the comparison. It's always a good idea to make your input much smaller (say (2,1,2,2)) so you can see what actually happens.

Answer 2

In [152]: alist = []
In [154]: alist.append(np.random.random((2,1,3)))
In [155]: alist.append(np.random.random((2,1,3)))
In [156]: alist
Out[156]: 
[array([[[0.85221826, 0.56088315, 0.06232853]],
 
        [[0.0966469 , 0.89513922, 0.44814579]]]),
 array([[[0.86207845, 0.88895573, 0.62069196]],
 
        [[0.11475614, 0.29473531, 0.11179268]]])]

Using np.array to join the list elements produces a 4d array - it has joined them on a new leading dimension:

In [157]: arr = np.array(alist)
In [158]: arr.shape
Out[158]: (2, 2, 1, 3)
In [159]: arr[-1,]        # same as alist[-1]
Out[159]: 
array([[[0.86207845, 0.88895573, 0.62069196]],

       [[0.11475614, 0.29473531, 0.11179268]]])

If we concatenate on one of the dimensions:

In [160]: arr = np.concatenate(alist, axis=1)
In [161]: arr
Out[161]: 
array([[[0.85221826, 0.56088315, 0.06232853],
        [0.86207845, 0.88895573, 0.62069196]],

       [[0.0966469 , 0.89513922, 0.44814579],
        [0.11475614, 0.29473531, 0.11179268]]])
In [162]: arr.shape
Out[162]: (2, 2, 3)     # note the shape - that 2nd 2 is the join axis
In [163]: arr[:,-1]
Out[163]: 
array([[0.86207845, 0.88895573, 0.62069196],
       [0.11475614, 0.29473531, 0.11179268]])

[163] has the same numbers as [159], but a (2,3) shape.

reshape keeps the values, but may 'shuffle' them:

In [164]: np.array(alist).reshape(2,2,3)
Out[164]: 
array([[[0.85221826, 0.56088315, 0.06232853],
        [0.0966469 , 0.89513922, 0.44814579]],

       [[0.86207845, 0.88895573, 0.62069196],
        [0.11475614, 0.29473531, 0.11179268]]])

We have transpose the leading 2 axes before reshape to match [161]

In [165]: np.array(alist).transpose(1,0,2,3)
Out[165]: 
array([[[[0.85221826, 0.56088315, 0.06232853]],

        [[0.86207845, 0.88895573, 0.62069196]]],


       [[[0.0966469 , 0.89513922, 0.44814579]],

        [[0.11475614, 0.29473531, 0.11179268]]]])
In [166]: np.array(alist).transpose(1,0,2,3).reshape(2,2,3)
Out[166]: 
array([[[0.85221826, 0.56088315, 0.06232853],
        [0.86207845, 0.88895573, 0.62069196]],

       [[0.0966469 , 0.89513922, 0.44814579],
        [0.11475614, 0.29473531, 0.11179268]]])