Numpy: stack array by the last dimension

Question

Suppose I have 3 numpy arrays a, b, c, of the same shape, say

a.shape == b.shape == c.shape == (7,9)

Now I'd like to create a 3-dimensional array of size (7,9,3), say x, such that

x[:,:,0] == a
x[:,:,1] == b
x[:,:,2] == c

What is the "pythonic" way of doing it (perhaps in one line)?

Thanks in advance!

Answer 1

There's a function that does exactly that: numpy.dstack ("d" for "depth"). For example:

In [10]: import numpy as np

In [11]: a = np.ones((7, 9))

In [12]: b = a * 2

In [13]: c = a * 3

In [15]: x = np.dstack((a, b, c))

In [16]: x.shape
Out[16]: (7, 9, 3)

In [17]: (x[:, :, 0] == a).all()
Out[17]: True

In [18]: (x[:, :, 1] == b).all()
Out[18]: True

In [19]: (x[:, :, 2] == c).all()
Out[19]: True

Answer 2

TL;DR:

Use numpy.stack (docs), which joins a sequence of arrays along a new axis of your choice.

---

Although @NPE answer is very good and cover many cases, there are some scenarios in which numpy.dstack isn't the right choice (I've just found that out while trying to use it). That's because numpy.dstack, according to the docs:

Stacks arrays in sequence depth wise (along third axis).

This is equivalent to concatenation along the third axis after 2-D arrays of shape (M,N) have been reshaped to (M,N,1) and 1-D arrays of shape (N,) have been reshaped to (1,N,1).

Let's walk through an example in which this function isn't desirable. Suppose you have a list with 512 numpy arrays of shape (3, 3, 3) and want to stack them in order to get a new array of shape (3, 3, 3, 512). In my case, those 512 arrays were filters of a 2D-convolutional layer. If you use numpy.dstack:

>>> len(arrays_list)
512
>>> arrays_list[0].shape
(3, 3, 3)
>>> numpy.dstack(arrays_list).shape
(3, 3, 1536)

That's because numpy.dstack always stacks the arrays along the third axis! Alternatively, you should use numpy.stack (docs), which joins a sequence of arrays along a new axis of your choice:

>>> numpy.stack(arrays_list, axis=-1).shape
(3, 3, 3, 512)

In my case, I passed -1 to the axis parameter because I wanted the arrays stacked along the last axis.