Add nested model test

Author: zaxtax

pymc_experimental/tests/model/test_marginal_model.py

@@ -4,15 +4,13 @@
 import numpy as np
 import pandas as pd
 import pymc as pm
-import pytensor
 import pytensor.tensor as pt
 import pytest
 from arviz import InferenceData, dict_to_dataset
 from pymc import ImputationWarning, inputvars
 from pymc.distributions import transforms
 from pymc.logprob.abstract import _logprob
 from pymc.util import UNSET
-from pytensor.graph import vectorize_graph
 from scipy.special import logsumexp
 from scipy.stats import norm
 
@@ -291,6 +289,58 @@ def true_logp(y):
     )
 
 
+@pytest.mark.filterwarnings("error")
+def test_nested_recover_marginals():
+    """Test that marginalization works when there are nested marginalized RVs"""
+
+    with MarginalModel() as m:
+        idx = pm.Bernoulli("idx", p=0.75)
+        sub_idx = pm.Bernoulli("sub_idx", p=pt.switch(pt.eq(idx, 0), 0.15, 0.95))
+        sub_dep = pm.Normal("y", mu=idx + sub_idx, sigma=1.0)
+
+    m.marginalize([idx, sub_idx])
+
+    rng = np.random.default_rng(211)
+
+    with m:
+        prior = pm.sample_prior_predictive(
+            samples=20,
+            random_seed=rng,
+            return_inferencedata=False,
+        )
+        idata = InferenceData(posterior=dict_to_dataset(prior))
+
+    idata = m.recover_marginals(idata, include_samples=True)
+    assert "idx" in idata
+    assert "lp_idx" in idata
+    assert idata.idx.shape == idata.y.shape
+    assert idata.lp_idx.shape == idata.idx.shape + (2,)
+    assert "sub_idx" in idata
+    assert "lp_sub_idx" in idata
+    assert idata.sub_idx.shape == idata.y.shape
+    assert idata.lp_sub_idx.shape == idata.sub_idx.shape + (2,)
+
+    def true_idx_logp(y):
+        idx_0 = np.log(0.85 * 0.25 * norm.pdf(y, loc=0) + 0.15 * 0.25 * norm.pdf(y, loc=1))
+        idx_1 = np.log(0.05 * 0.75 * norm.pdf(y, loc=1) + 0.95 * 0.75 * norm.pdf(y, loc=2))
+        return np.stack([idx_0, idx_1]).T
+
+    np.testing.assert_almost_equal(
+        true_idx_logp(idata.y.values.flatten()),
+        idata.lp_idx[0].values,
+    )
+
+    def true_sub_idx_logp(y):
+        sub_idx_0 = np.log(0.85 * 0.25 * norm.pdf(y, loc=0) + 0.05 * 0.75 * norm.pdf(y, loc=1))
+        sub_idx_1 = np.log(0.15 * 0.25 * norm.pdf(y, loc=1) + 0.95 * 0.75 * norm.pdf(y, loc=2))
+        return np.stack([sub_idx_0, sub_idx_1]).T
+
+    np.testing.assert_almost_equal(
+        true_sub_idx_logp(idata.y.values.flatten()),
+        idata.lp_sub_idx[0].values,
+    )
+
+
 @pytest.mark.filterwarnings("error")
 def test_not_supported_marginalized():
     """Marginalized graphs with non-Elemwise Operations are not supported as they