Sum distributions based on another distribution

Hello,

I am a beginner in probabilistic concepts, and I have a question.

I have 3 random variables:

a = pm.Categorical('a', [0.8, 0.1, 0.1])
b = pm.Categorical('b', [0.5, 0.2, 0.2, 0.1])
c = pm.Categorical('c', [0.9, 0.0, 0.0, 0.1])

and I want to sum them.

The simple solution is just to do

d = a + b + c

but the tricky part is that I want to do it based on another random variable

d = pm.Categorical('d', [0, 0.1, 0.85, 0.05])

meaning that if 100 iterations of sampling are done:

  • in 10 of them, only the variable “a” will be summed,

  • in 85 of the iterations only the variables “a” and “b” will be summed and

  • in 5 iterations all the variables “a”, “b” and “c” will be summed.

Thanks in advance

1 Like

You can use pm.math.switch() to build conditional logic. So maybe something along these lines?

d = pm.Categorical('d', [0.1, 0.85, 0.05])
b_logic = pm.math.switch(pm.math.eq(d, 1), 0, b)
c_logic = pm.math.switch(pm.math.eq(d, 2), 0, c)
sum = pm.Deterministic('sum', a + b_logic + c_logic)

Not pretty, but hopefully it helps.

3 Likes