Indicator from independent variable sets distribution or zero

I am trying to figure out whether it is possible to have a distribution present when an indicator variable is =1, and not present when it is = 0.

e.g. in the following example, I want lp to only be Gaussian when v_i == 1, and to just be 0 otherwise.
(of course for this simple example I could just exclude obs where v_i == 0 prior to the model, but this is just a simplified example)

obs = np.random.normal(0,1,size=(20,1))
v_i = np.random.binomial(1,0.5, size=(20,1))
N_i = len(np.unique(v_i))
obs[v_i==0] = 0

model = pm.Model()

with model:
Ip_base = pm.Normal('Ip_base', mu=0, sd=1e-2)
Ip_mu = pm.Normal('Ip_mu', mu=0, sd=1)
Ip = Ip_mu[v_i] + Ip_base
nx = pm.Normal('nx', mu=Ip, sd = 1, observed=obs)

I think in various google searches a while ago, I found that it’s possible to do something like an if statement (I think it was theano switch) - not sure if that is the best way forwards though -


You can also stack the different RV together, and use the observed v_i to index them:

v = np.random.randint(2, size=100)
with pm.Model() as m:
    RV1 = pm.Normal('rv1', 0, 1)
    RV2 = pm.Flat('rv2') # Flat distribution returns 0 logp
    rvall_logp = tt.stack([RV1.logpt, RV2.logpt])
    observed = pm.Potential('obs', rvall_logp[v])

OK - sounds good! Thanks.