Python Numpy stack 2d arrays in vector

Question

So, I would like to stack couple 2d arrays to vector so it would look like this:

[[[0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]]

 [[0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]]

 [[0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]]]

I can make smth like this:

import numpy as np
a = np.zeros((5, 5), dtype=int)
b = np.zeros((5, 5), dtype=int)
c = np.stack((a, b), 0)
print(c)

To get this:

[[[0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]]

 [[0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]
  [0 0 0 0 0]]]

But I cant figure out how to add third 2d array to such vector or how to create such vector of 2d arrays iteratively in a loop. Append, stack, concat just dont keep the needed shape

So, any suggestions? Thank you!

Conclusion: Thanks to Tom and Mozway we've got two answers

Tom's:

data_x_train = x_train[np.where((y_train==0) | (y_train==1))

Mozway's:

out = np.empty((0,5,5))

while condition:
    # get new array
    a = XXX
    out = np.r_[out, a[None]]
out

Answer 1

Do you mean something like:

np.tile(a, (3, 1, 1))

array([[[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]]])

Edit: Do you mean something like:

test = np.tile(a, (3000, 1, 1))
filtered_subset = test[[1, 10, 100], :, :]

array([[[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]]])

Answer 2

Assuming the following arrays:

a = np.ones((5, 5), dtype=int)
b = np.ones((5, 5), dtype=int)*2
c = np.ones((5, 5), dtype=int)*3

You can stack all at once using:

np.stack((a, b, c), 0)

If you really need to add the arrays iteratively, you can use np.r_ :

out = a[None]

for i in (b,c):
    out = np.r_[out, i[None]]

output:

array([[[1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1]],

       [[2, 2, 2, 2, 2],
        [2, 2, 2, 2, 2],
        [2, 2, 2, 2, 2],
        [2, 2, 2, 2, 2],
        [2, 2, 2, 2, 2]],

       [[3, 3, 3, 3, 3],
        [3, 3, 3, 3, 3],
        [3, 3, 3, 3, 3],
        [3, 3, 3, 3, 3],
        [3, 3, 3, 3, 3]]])

edit: if you do not know the arrays in advance

out = np.empty((0,5,5))

while condition:
    # get new array
    a = XXX
    out = np.r_[out, a[None]]
out