Add temporary fix for pretty representation of SymbolicDistributions

Author: ricardoV94

pymc/distributions/distribution.py

@@ -47,7 +47,7 @@
     resize_from_dims,
     resize_from_observed,
 )
-from pymc.printing import str_for_dist
+from pymc.printing import str_for_dist, str_for_symbolic_dist
 from pymc.util import UNSET
 from pymc.vartypes import string_types
 
@@ -483,15 +483,11 @@ def __new__(
             transform=transform,
             initval=initval,
         )
-
-        # TODO: Refactor this
         # add in pretty-printing support
-        rv_out.str_repr = lambda *args, **kwargs: name
-        rv_out._repr_latex_ = f"\\text{name}"
-        # rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
-        # rv_out._repr_latex_ = types.MethodType(
-        #     functools.partial(str_for_dist, formatting="latex"), rv_out
-        # )
+        rv_out.str_repr = types.MethodType(str_for_symbolic_dist, rv_out)
+        rv_out._repr_latex_ = types.MethodType(
+            functools.partial(str_for_symbolic_dist, formatting="latex"), rv_out
+        )
 
         return rv_out
 

pymc/printing.py

@@ -55,6 +55,18 @@ def str_for_dist(rv: TensorVariable, formatting: str = "plain", include_params:
             return rf"{print_name} ~ {dist_name}"
 
 
+def str_for_symbolic_dist(
+    rv: TensorVariable, formatting: str = "plain", include_params: bool = True
+) -> str:
+    """Make a human-readable string representation of a SymbolicDistribution in a model,
+    either LaTeX or plain, optionally with distribution parameter values included."""
+
+    if "latex" in formatting:
+        return rf"$\text{{{rv.name}}} \sim \text{{{rv.owner.op}}}$"
+    else:
+        return rf"{rv.name} ~ {rv.owner.op}"
+
+
 def str_for_model(model: Model, formatting: str = "plain", include_params: bool = True) -> str:
     """Make a human-readable string representation of Model, listing all random variables
     and their distributions, optionally including parameter values."""

pymc/tests/test_printing.py

@@ -9,6 +9,7 @@
     NegativeBinomial,
     Normal,
     Uniform,
+    ZeroInflatedPoisson,
 )
 from pymc.math import dot
 from pymc.model import Deterministic, Model, Potential
@@ -47,6 +48,9 @@ def setup_class(self):
             # )
             nb2 = NegativeBinomial("nb_with_p_n", p=Uniform("nbp"), n=10)
 
+            # Symbolic distribution
+            zip = ZeroInflatedPoisson("zip", 0.5, 5)
+
             # Expected value of outcome
             mu = Deterministic("mu", floatX(alpha + dot(X, b)))
 
@@ -76,7 +80,7 @@ def setup_class(self):
             # add a potential as well
             pot = Potential("pot", mu**2)
 
-        self.distributions = [alpha, sigma, mu, b, Z, nb2, Y_obs, pot]
+        self.distributions = [alpha, sigma, mu, b, Z, nb2, zip, Y_obs, pot]
         self.deterministics_or_potentials = [mu, pot]
         # tuples of (formatting, include_params
         self.formats = [("plain", True), ("plain", False), ("latex", True), ("latex", False)]
@@ -88,6 +92,7 @@ def setup_class(self):
                 r"beta ~ N(0, 10)",
                 r"Z ~ N(f(), f())",
                 r"nb_with_p_n ~ NB(10, nbp)",
+                r"zip ~ MarginalMixtureRV{inline=False}",
                 r"Y_obs ~ N(mu, sigma)",
                 r"pot ~ Potential(f(beta, alpha))",
             ],
@@ -98,6 +103,7 @@ def setup_class(self):
                 r"beta ~ N",
                 r"Z ~ N",
                 r"nb_with_p_n ~ NB",
+                r"zip ~ MarginalMixtureRV{inline=False}",
                 r"Y_obs ~ N",
                 r"pot ~ Potential",
             ],
@@ -108,6 +114,7 @@ def setup_class(self):
                 r"$\text{beta} \sim \operatorname{N}(0,~10)$",
                 r"$\text{Z} \sim \operatorname{N}(f(),~f())$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}(10,~\text{nbp})$",
+                r"$\text{zip} \sim \text{MarginalMixtureRV{inline=False}}$",
                 r"$\text{Y_obs} \sim \operatorname{N}(\text{mu},~\text{sigma})$",
                 r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",
             ],
@@ -118,6 +125,7 @@ def setup_class(self):
                 r"$\text{beta} \sim \operatorname{N}$",
                 r"$\text{Z} \sim \operatorname{N}$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}$",
+                r"$\text{zip} \sim \text{MarginalMixtureRV{inline=False}}$",
                 r"$\text{Y_obs} \sim \operatorname{N}$",
                 r"$\text{pot} \sim \operatorname{Potential}$",
             ],