search

pymc-devs / pymc-experimental

https://pymc-experimental.readthedocs.io
Other
77 stars 49 forks source link

Add test for MarginalModel where variable depends on two marginalized variables #286

Closed ricardoV94 closed 9 months ago

ricardoV94 commented 9 months ago

I am positively surprised this is working! Would probably be worth adding as a test.

We already have a test for when multiple RVs depend on a marginalized RV or when there are nested marginalized RVs. This will also create a nested marginalized RV even though x and y are independent a priori

import pymc as pm
from pymc_experimental.model.marginal_model import MarginalModel
import numpy as np

with MarginalModel() as m:
    x = pm.Bernoulli("x", 0.1)
    y = pm.Bernoulli("y", 0.3)
    z = pm.DiracDelta("z", c=x+y)

m.marginalize([x, y])
logp = m.compile_logp()

np.testing.assert_allclose(np.exp(logp({"z": 0})), 0.9 * 0.7)
np.testing.assert_allclose(np.exp(logp({"z": 1})), 0.9 * 0.3 + 0.1 * 0.7)
np.testing.assert_allclose(np.exp(logp({"z": 2})), 0.1 * 0.3)