Use numpy to generate data with condition

Question

I want to generate data to be modeled. Based on the following logic:

Input: 2 numpy array named z and d. z : array, 0/1 d : array, 0/1

Return: y: array. norm random numbers.

If z == 0 and d==0, y ~ norm(1,1),

if z == 0 and d == 1, y~ norm(0,1),

if z == 1 and d == 0, y ~ norm(1,1),

if z == 1 and d == 1, y ~ norm(2,1).

I want to do it in a super faster,clear and pythonic way.

It seems basic math and np.where is faster. In this case, I only have 3 conditions (you can see clearly from basic mathmatics part). If I have 10 or more condition, type them in if-else is sometimes confusing, what is the best way to do it?

What I have tried:

# generate data
n = 2000
z = np.random.binomial(1,0.5,n)
d = np.random.binomial(1,0.5,n)

dict case-when

def myfun(x):
    return {(0,1):np.random.normal(0,1),\
            (0,0):np.random.normal(1,1),\
            (1,0):np.random.normal(1,1),\
            (1,1):np.random.normal(2,1)}[x]
%%timeit
y = [myfun(i) for i in zip(z,d)]

Out:

16.2 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

simple if-else

%%timeit
y = np.random.normal([0 if i == 0 & j ==1 else 2 if i == 1 & j == 1 else 1 for i in z for j in d],1)

Out:

2.31 s ± 12.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

basic mathmatics

%%timeit
h0 = np.random.normal(0,1,n)
h1 = np.random.normal(1,1,n)
h2 = np.random.normal(2,1,n)
y = (1-z)*d*h0 + (1-d)*h1 + z*d*h2

Out:

140 µs ± 135 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

np.where

%%timeit
h0 = np.random.normal(0,1,n)
h1 = np.random.normal(1,1,n)
h2 = np.random.normal(2,1,n)
y = np.where((d== 0),h1,0) + np.where((z ==1) & (d== 1),h2,0) + np.where((z ==0) & (d== 1),h0,0)

Out:

156 µs ± 598 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)