First crack at transforms

Author: zaxtax

pymc_experimental/tests/model/test_marginal_model.py

@@ -12,7 +12,7 @@
 from pymc.logprob.abstract import _logprob
 from pymc.util import UNSET
 from scipy.special import logsumexp
-from scipy.stats import norm
+from scipy.stats import halfnorm, norm
 
 from pymc_experimental.model.marginal_model import (
     FiniteDiscreteMarginalRV,
@@ -255,11 +255,12 @@ def test_marginalized_change_point_model_sampling(disaster_model):
 
 def test_recover_marginals_basic():
     with MarginalModel() as m:
+        sigma = pm.HalfNormal("sigma", transform=None)
         p = np.array([0.5, 0.2, 0.3])
         k = pm.Categorical("k", p=p)
         mu = np.array([-3.0, 0.0, 3.0])
         mu_ = pt.as_tensor_variable(mu)
-        y = pm.Normal("y", mu=mu_[k])
+        y = pm.Normal("y", mu=mu_[k], sigma=sigma)
 
     m.marginalize([k])
 
@@ -279,17 +280,17 @@ def test_recover_marginals_basic():
     assert idata.k.shape == idata.y.shape
     assert idata.lp_k.shape == idata.k.shape + (len(p),)
 
-    def true_logp(y):
+    def true_logp(y, sigma):
         y = y.repeat(len(p)).reshape(len(y), -1)
-        return np.log(p) + norm.logpdf(y, loc=mu)
+        sigma = sigma.repeat(len(p)).reshape(len(sigma), -1)
+        return np.log(p) + norm.logpdf(y, loc=mu, scale=sigma) + halfnorm.logpdf(sigma)
 
     np.testing.assert_almost_equal(
-        true_logp(idata.y.values.flatten()),
+        true_logp(idata.y.values.flatten(), idata.sigma.values.flatten()),
         idata.lp_k[0].values,
     )
 
 
-@pytest.mark.filterwarnings("error")
 def test_nested_recover_marginals():
     """Test that marginalization works when there are nested marginalized RVs"""