Enable nested SymbolicRandomVariables

Author: ricardoV94

pymc/distributions/censored.py

@@ -86,7 +86,9 @@ class Censored(Distribution):
 
     @classmethod
     def dist(cls, dist, lower, upper, **kwargs):
-        if not isinstance(dist, TensorVariable) or not isinstance(dist.owner.op, RandomVariable):
+        if not isinstance(dist, TensorVariable) or not isinstance(
+            dist.owner.op, (RandomVariable, SymbolicRandomVariable)
+        ):
             raise ValueError(
                 f"Censoring dist must be a distribution created via the `.dist()` API, got {type(dist)}"
             )

pymc/distributions/mixture.py

@@ -191,7 +191,7 @@ def dist(cls, w, comp_dists, **kwargs):
             # TODO: Allow these to not be a RandomVariable as long as we can call `ndim_supp` on them
             #  and resize them
             if not isinstance(dist, TensorVariable) or not isinstance(
-                dist.owner.op, RandomVariable
+                dist.owner.op, (RandomVariable, SymbolicRandomVariable)
             ):
                 raise ValueError(
                     f"Component dist must be a distribution created via the `.dist()` API, got {type(dist)}"

pymc/distributions/multivariate.py

@@ -1180,7 +1180,7 @@ def dist(cls, n, eta, sd_dist, **kwargs):
         if not (
             isinstance(sd_dist, Variable)
             and sd_dist.owner is not None
-            and isinstance(sd_dist.owner.op, RandomVariable)
+            and isinstance(sd_dist.owner.op, (RandomVariable, SymbolicRandomVariable))
             and sd_dist.owner.op.ndim_supp < 2
         ):
             raise TypeError("sd_dist must be a scalar or vector distribution variable")

pymc/distributions/timeseries.py

@@ -155,7 +155,7 @@ def dist(cls, init_dist, innovation_dist, steps=None, **kwargs) -> at.TensorVari
         if not (
             isinstance(init_dist, at.TensorVariable)
             and init_dist.owner is not None
-            and isinstance(init_dist.owner.op, RandomVariable)
+            and isinstance(init_dist.owner.op, (RandomVariable, SymbolicRandomVariable))
             # TODO: Lift univariate constraint on init_dist
             and init_dist.owner.op.ndim_supp == 0
         ):
@@ -165,7 +165,7 @@ def dist(cls, init_dist, innovation_dist, steps=None, **kwargs) -> at.TensorVari
         if not (
             isinstance(innovation_dist, at.TensorVariable)
             and init_dist.owner is not None
-            and isinstance(init_dist.owner.op, RandomVariable)
+            and isinstance(init_dist.owner.op, (RandomVariable, SymbolicRandomVariable))
             # TODO: Lift univariate constraint on inovvation_dist
             and init_dist.owner.op.ndim_supp == 0
         ):
@@ -434,7 +434,7 @@ def dist(
 
         if init_dist is not None:
             if not isinstance(init_dist, TensorVariable) or not isinstance(
-                init_dist.owner.op, RandomVariable
+                init_dist.owner.op, (RandomVariable, SymbolicRandomVariable)
             ):
                 raise ValueError(
                     f"Init dist must be a distribution created via the `.dist()` API, "

pymc/printing.py

@@ -151,6 +151,9 @@ def str_for_potential_or_deterministic(
 
 
 def _str_for_input_var(var: Variable, formatting: str) -> str:
+    # Avoid circular import
+    from pymc.distributions.distribution import SymbolicRandomVariable
+
     def _is_potential_or_determinstic(var: Variable) -> bool:
         try:
             return var.str_repr.__func__.func is str_for_potential_or_deterministic
@@ -160,7 +163,9 @@ def _is_potential_or_determinstic(var: Variable) -> bool:
 
     if isinstance(var, TensorConstant):
         return _str_for_constant(var, formatting)
-    elif isinstance(var.owner.op, RandomVariable) or _is_potential_or_determinstic(var):
+    elif isinstance(
+        var.owner.op, (RandomVariable, SymbolicRandomVariable)
+    ) or _is_potential_or_determinstic(var):
         # show the names for RandomVariables, Deterministics, and Potentials, rather
         # than the full expression
         return _str_for_input_rv(var, formatting)
@@ -194,15 +199,18 @@ def _str_for_constant(var: TensorConstant, formatting: str) -> str:
 
 
 def _str_for_expression(var: Variable, formatting: str) -> str:
+    # Avoid circular import
+    from pymc.distributions.distribution import SymbolicRandomVariable
+
     # construct a string like f(a1, ..., aN) listing all random variables a as arguments
     def _expand(x):
-        if x.owner and (not isinstance(x.owner.op, RandomVariable)):
+        if x.owner and (not isinstance(x.owner.op, (RandomVariable, SymbolicRandomVariable))):
             return reversed(x.owner.inputs)
 
     parents = [
         x
         for x in walk(nodes=var.owner.inputs, expand=_expand)
-        if x.owner and isinstance(x.owner.op, RandomVariable)
+        if x.owner and isinstance(x.owner.op, (RandomVariable, SymbolicRandomVariable))
     ]
     names = [x.name for x in parents]
 

pymc/tests/distributions/test_mixture.py

@@ -599,13 +599,17 @@ def test_list_mvnormals_predictive_sampling_shape(self):
         assert prior["mu0"].shape == (n_samples, D)
         assert prior["chol_cov_0"].shape == (n_samples, D * (D + 1) // 2)
 
-    @pytest.mark.xfail(reason="Nested mixtures not refactored yet")
     def test_nested_mixture(self):
         if aesara.config.floatX == "float32":
             rtol = 1e-4
         else:
             rtol = 1e-7
         nbr = 4
+
+        norm_x = generate_normal_mixture_data(
+            np.r_[0.75, 0.25], np.r_[0.0, 5.0], np.r_[1.0, 1.0], size=1000
+        )
+
         with Model() as model:
             # mixtures components
             g_comp = Normal.dist(
@@ -622,7 +626,7 @@ def test_nested_mixture(self):
             l_mix = Mixture.dist(w=l_w, comp_dists=l_comp)
             # mixture of mixtures
             mix_w = Dirichlet("mix_w", a=floatX(np.ones(2)), transform=None, shape=(2,))
-            mix = Mixture("mix", w=mix_w, comp_dists=[g_mix, l_mix], observed=np.exp(self.norm_x))
+            mix = Mixture("mix", w=mix_w, comp_dists=[g_mix, l_mix], observed=np.exp(norm_x))
 
         test_point = model.initial_point()
 
@@ -658,21 +662,23 @@ def mixmixlogp(value, point):
             )
             return priorlogp, mixmixlogpg
 
-        value = np.exp(self.norm_x)[:, None]
+        value = np.exp(norm_x)[:, None]
         priorlogp, mixmixlogpg = mixmixlogp(value, test_point)
 
         # check logp of mixture
-        assert_allclose(mixmixlogpg, mix.logp_elemwise(test_point), rtol=rtol)
+        mix_logp_fn = model.compile_logp(vars=[mix], sum=False)
+        assert_allclose(mixmixlogpg, mix_logp_fn(test_point)[0], rtol=rtol)
 
         # check model logp
-        assert_allclose(priorlogp + mixmixlogpg.sum(), model.logp(test_point), rtol=rtol)
+        model_logp_fn = model.compile_logp()
+        assert_allclose(priorlogp + mixmixlogpg.sum(), model_logp_fn(test_point), rtol=rtol)
 
         # check input and check logp again
         test_point["g_w"] = np.asarray([0.1, 0.1, 0.2, 0.6])
         test_point["mu_g"] = np.exp(np.random.randn(nbr))
         priorlogp, mixmixlogpg = mixmixlogp(value, test_point)
-        assert_allclose(mixmixlogpg, mix.logp_elemwise(test_point), rtol=rtol)
-        assert_allclose(priorlogp + mixmixlogpg.sum(), model.logp(test_point), rtol=rtol)
+        assert_allclose(mixmixlogpg, mix_logp_fn(test_point)[0], rtol=rtol)
+        assert_allclose(priorlogp + mixmixlogpg.sum(), model_logp_fn(test_point), rtol=rtol)
 
     def test_iterable_single_component_warning(self):
         with warnings.catch_warnings():

pymc/tests/test_printing.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from pymc import Bernoulli, Censored, Mixture
 from pymc.aesaraf import floatX
 from pymc.distributions import (
     DirichletMultinomial,
@@ -48,9 +49,14 @@ def setup_class(self):
             # )
             nb2 = NegativeBinomial("nb_with_p_n", p=Uniform("nbp"), n=10)
 
-            # Symbolic distribution
+            # SymbolicRV
             zip = ZeroInflatedPoisson("zip", 0.5, 5)
 
+            # Nested SymbolicRV
+            comp_1 = ZeroInflatedPoisson.dist(0.5, 5)
+            comp_2 = Censored.dist(Bernoulli.dist(0.5), -1, 1)
+            nested_mix = Mixture("nested_mix", [0.5, 0.5], [comp_1, comp_2])
+
             # Expected value of outcome
             mu = Deterministic("mu", floatX(alpha + dot(X, b)))
 
@@ -80,7 +86,7 @@ def setup_class(self):
             # add a potential as well
             pot = Potential("pot", mu**2)
 
-        self.distributions = [alpha, sigma, mu, b, Z, nb2, zip, Y_obs, pot]
+        self.distributions = [alpha, sigma, mu, b, Z, nb2, zip, nested_mix, Y_obs, pot]
         self.deterministics_or_potentials = [mu, pot]
         # tuples of (formatting, include_params
         self.formats = [("plain", True), ("plain", False), ("latex", True), ("latex", False)]
@@ -93,6 +99,11 @@ def setup_class(self):
                 r"Z ~ N(f(), f())",
                 r"nb_with_p_n ~ NB(10, nbp)",
                 r"zip ~ MarginalMixture(f(), DiracDelta(0), Pois(5))",
+                (
+                    r"nested_mix ~ MarginalMixture(<constant>, "
+                    r"MarginalMixture(f(), DiracDelta(0), Pois(5)), "
+                    r"Censored(Bern(0.5), -1, 1))"
+                ),
                 r"Y_obs ~ N(mu, sigma)",
                 r"pot ~ Potential(f(beta, alpha))",
             ],
@@ -104,6 +115,7 @@ def setup_class(self):
                 r"Z ~ N",
                 r"nb_with_p_n ~ NB",
                 r"zip ~ MarginalMixture",
+                r"nested_mix ~ MarginalMixture",
                 r"Y_obs ~ N",
                 r"pot ~ Potential",
             ],
@@ -115,6 +127,11 @@ def setup_class(self):
                 r"$\text{Z} \sim \operatorname{N}(f(),~f())$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}(10,~\text{nbp})$",
                 r"$\text{zip} \sim \operatorname{MarginalMixture}(f(),~\text{\$\operatorname{DiracDelta}(0)\$},~\text{\$\operatorname{Pois}(5)\$})$",
+                (
+                    r"$\text{nested_mix} \sim \operatorname{MarginalMixture}(\text{<constant>},"
+                    r"~\text{\$\operatorname{MarginalMixture}(f(),~\text{\\$\operatorname{DiracDelta}(0)\\$},~\text{\\$\operatorname{Pois}(5)\\$})\$},"
+                    r"~\text{\$\operatorname{Censored}(\text{\\$\operatorname{Bern}(0.5)\\$},~-1,~1)\$})$"
+                ),
                 r"$\text{Y_obs} \sim \operatorname{N}(\text{mu},~\text{sigma})$",
                 r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",
             ],
@@ -126,6 +143,7 @@ def setup_class(self):
                 r"$\text{Z} \sim \operatorname{N}$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}$",
                 r"$\text{zip} \sim \operatorname{MarginalMixture}$",
+                r"$\text{nested_mix} \sim \operatorname{MarginalMixture}$",
                 r"$\text{Y_obs} \sim \operatorname{N}$",
                 r"$\text{pot} \sim \operatorname{Potential}$",
             ],