numpy create Mx2 matrix based on Mx1 matrix

Viewed 77

I have a np.array with 800 values. Each value is either 0 or 1.

If the value is 0, I want to replace it with mu_0, which is a 1x2 array; else, I want to replace it with mu_1, which is also a 1x2 array.

I tried using np.where(y == 0, mu_0, mu_1), but python would only broadcast the value of mu to match y, not the other way around. In particular, the error I get is

ValueError: operands could not be broadcast together with shapes (800,) (2,) (2,) 

I tried expanding y into (800, 2), by padding y_pad = np.c_[y, np.zeros(800)], but I am unsure how to condition on the first value of each row.

If I use np.where(y_pad[:, 0] == 0, ...), the array gets sliced back into (800,) again.

2 Answers

You can indeed expand y into (800, 2) with, for example, np.repeat so that 0s are 0s, 1s are 1s for each row. Then we can use np.where:

# first casting to (800, 1) with `newaxis` then repetition
y_rep = np.repeat(y[:, np.newaxis], repeats=2, axis=1)
result = np.where(y_rep == 0, mu_0, mu_1)

sample run:

mu_0 = np.array([ 9, 17])
mu_1 = np.array([-3, -5])

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

then

>>> result

array([[ 9, 17],
       [ 9, 17],
       [-3, -5],
       [-3, -5],
       [ 9, 17],
       [-3, -5],
       [-3, -5],
       [-3, -5]])

where the condition became:

>>> y_rep == 0

array([[ True,  True],
       [ True,  True],
       [False, False],
       [False, False],
       [ True,  True],
       [False, False],
       [False, False],
       [False, False]])

You can use indexing:

## Dummy data
# The 1x2 array
mu_0 = np.array([[1,1]])
mu_1 = np.array([[2,2]])
# The boolean array
y    = np.array([0,1,0,1,0])


## Get the result:
res = np.vstack((mu_0,mu_1))[y,:]  

And we obtain the following array:

array([[1, 1],
       [2, 2],
       [1, 1],
       [2, 2],
       [1, 1]]) 
Related