How to manipulate numpy arrays in this way?

Question

I have a (28,28) array a. And I want to obtain a (28,28,3) array b s.t. b[i][j][0] = b[i][j][1] = b[i][j][2] = a[i][j].

Is there any numpy shortcut to do this without tedious for loops?

Answer 1

>>> import numpy as np
>>> a = np.zeros((28,28))
>>> b = np.dstack((a,a,a))
>>> a.shape
(28, 28)
>>> b.shape
(28, 28, 3)

Example:

>>> a = np.array([[1,2],[3,4]])
>>> b = np.dstack((a,a,a))
>>> a
array([[1, 2],
       [3, 4]])
>>> b
array([[[1, 1, 1],
        [2, 2, 2]],

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

Answer 2

You can append an axis to a and then repeat it n=3 times:

>>> a = np.random.randn(28, 28)
>>> b = np.repeat(np.atleast_3d(a), repeats=3, axis=-1)
>>> b.shape 
(28, 28, 3)

And as required,

>>> np.all(a == b[...,0])
True
>>> np.all(b[...,0] == b[...,1])
True
>>> np.all(b[...,1] == b[...,2])
True

Answer 3

You can allocate b and use broadcasting:

b = np.empty(a.shape + (3,), a.dtype)
b[...] = a[..., None]

If you only need read access, then a very efficient way is creating a strided view:

c = np.lib.stride_tricks.as_strided(a, a.shape + (3,), a.strides + (0,))

This shares its data with a meaning that when you write to c, a will also change. Moreover along the last axis the stride is zero meaning that for example c[1, 1, 0] and c[1, 1, 2] are the same memory, change one and the other will also change. If that's not desired make a copy:

b = c.copy()