Add nested model test

Author: zaxtax

pymc_experimental/tests/model/test_marginal_model.py

@@ -291,6 +291,62 @@ 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