How to expand a 2d matrix to 3d by replacing the scalar with a vector in numpy-CodePudding

I have a 2d array like this:

a = np.array([[1, 0, 1],
              [0, 0 ,0]])

with shape (2,3). Indexing at some point would give a scalar, for example a[0,0] = 1 How can I turn this into a 3d array by replacing the scalar with a vector of length x filled with the initial scalar? For example say x = 5, I would want a[0,0] = [1,1,1,1,1] , a[0,1] = [0,0,0,0,0], ... and the shape of a to be (2,3,x)

CodePudding user response：

With np.einsum

import numpy as np
a = np.array([[1, 0, 1],
              [0, 0 ,0]])
x = 5

np.einsum('ij,k', a, np.ones(x, dtype=a.dtype))

Output

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

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

CodePudding user response：

Just use a.repeat() and then reshape:

x = 5
new_a = a.repeat(x).reshape(*a.shape, -1)

Output:

>>> new_a
array([[[1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1]],

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

>>> new_a[0, 0]
array([1, 1, 1, 1, 1])

>>> new_a[0, 1]
array([0, 0, 0, 0, 0])