Refactoring printing module with aeppl/printing.py with additions

pymc-devs · Jan 9, 2023 · e9d7210 · e9d7210
1 parent 76003e8
commit e9d7210
Show file tree

Hide file tree

Showing 5 changed files with 565 additions and 229 deletions.
diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py
@@ -1249,9 +1249,13 @@ def _zero_inflated_mixture(*, name, nonzero_p, nonzero_dist, **kwargs):
         nonzero_dist,
     ]
     if name is not None:
-        return Mixture(name, weights, comp_dists, **kwargs)
+        out_rv = Mixture(name, weights, comp_dists, **kwargs)
     else:
-        return Mixture.dist(weights, comp_dists, **kwargs)
+        out_rv = Mixture.dist(weights, comp_dists, **kwargs)
+
+    out_rv.is_zero_inflated = True
+
+    return out_rv
 
 
 class ZeroInflatedPoisson:

diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
@@ -14,7 +14,6 @@
 import contextvars
 import functools
 import sys
-import types
 import warnings
 
 from abc import ABCMeta
@@ -57,7 +56,6 @@
 )
 from pymc.logprob.rewriting import logprob_rewrites_db
 from pymc.model import BlockModelAccess
-from pymc.printing import str_for_dist
 from pymc.pytensorf import collect_default_updates, convert_observed_data
 from pymc.util import UNSET, _add_future_warning_tag
 from pymc.vartypes import string_types
@@ -317,12 +315,6 @@ def __new__(
             initval=initval,
         )
 
-        # add in pretty-printing support
-        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.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
         rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
         rv_out.random = _make_nice_attr_error("rv.random()", "pm.draw(rv)")

diff --git a/pymc/distributions/mixture.py b/pymc/distributions/mixture.py
@@ -162,6 +162,7 @@ class Mixture(Distribution):
     """
 
     rv_type = MarginalMixtureRV
+    is_zero_inflated = False
 
     @classmethod
     def dist(cls, w, comp_dists, **kwargs):

diff --git a/pymc/model.py b/pymc/model.py
@@ -12,7 +12,7 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-import functools
+import itertools
 import threading
 import types
 import warnings
@@ -589,10 +589,7 @@ def __init__(
 
         from pymc.printing import str_for_model
 
-        self.str_repr = types.MethodType(str_for_model, self)
-        self._repr_latex_ = types.MethodType(
-            functools.partial(str_for_model, formatting="latex"), self
-        )
+        self._repr_latex_ = types.MethodType(str_for_model, self)
 
     @property
     def model(self):
@@ -2015,17 +2012,17 @@ def Deterministic(name, var, model=None, dims=None):
     model.deterministics.append(var)
     model.add_named_variable(var, dims)
 
-    from pymc.printing import str_for_potential_or_deterministic
+    # from pymc.printing import str_for_potential_or_deterministic
 
-    var.str_repr = types.MethodType(
-        functools.partial(str_for_potential_or_deterministic, dist_name="Deterministic"), var
-    )
-    var._repr_latex_ = types.MethodType(
-        functools.partial(
-            str_for_potential_or_deterministic, dist_name="Deterministic", formatting="latex"
-        ),
-        var,
-    )
+    # var.str_repr = types.MethodType(
+    #     functools.partial(str_for_potential_or_deterministic, dist_name="Deterministic"), var
+    # )
+    # var._repr_latex_ = types.MethodType(
+    #     functools.partial(
+    #         str_for_potential_or_deterministic, dist_name="Deterministic", formatting="latex"
+    #     ),
+    #     var,
+    # )
 
     return var
 
@@ -2047,16 +2044,16 @@ def Potential(name, var, model=None):
     model.potentials.append(var)
     model.add_named_variable(var)
 
-    from pymc.printing import str_for_potential_or_deterministic
-
-    var.str_repr = types.MethodType(
-        functools.partial(str_for_potential_or_deterministic, dist_name="Potential"), var
-    )
-    var._repr_latex_ = types.MethodType(
-        functools.partial(
-            str_for_potential_or_deterministic, dist_name="Potential", formatting="latex"
-        ),
-        var,
-    )
+    # from pymc.printing import str_for_potential_or_deterministic
+
+    # var.str_repr = types.MethodType(
+    #     functools.partial(str_for_potential_or_deterministic, dist_name="Potential"), var
+    # )
+    # var._repr_latex_ = types.MethodType(
+    #     functools.partial(
+    #         str_for_potential_or_deterministic, dist_name="Potential", formatting="latex"
+    #     ),
+    #     var,
+    # )
 
     return var