add str/repr formatting options and change defaults accordingly

+ new options "latex_with_params" and "plain_with_params" replace the current default behavior of including input parameters (going back to default behaviour of the 3.9.3 release)
+ __latex__ and _repr_latex default to "latex_with_params"
+ __str__ and _str_repr default to "plain"
+ new formatting kwarg for model_to_graphviz enables switching between "plain" (default) and "plain_with_params"

Author: michaelosthege

pymc3/distributions/distribution.py

@@ -164,34 +164,51 @@ def _distr_name_for_repr(self):
         return self.__class__.__name__
 
     def _str_repr(self, name=None, dist=None, formatting="plain"):
-        """Generate string representation for this distribution, optionally
+        """
+        Generate string representation for this distribution, optionally
         including LaTeX markup (formatting='latex').
+
+        Parameters
+        ----------
+        name : str
+            name of the distribution
+        dist : Distribution
+            the distribution object
+        formatting : str
+            one of { "latex", "plain", "latex_with_params", "plain_with_params" }
         """
         if dist is None:
             dist = self
         if name is None:
             name = "[unnamed]"
+        supported_formattings = {"latex", "plain", "latex_with_params", "plain_with_params"}
+        if not formatting in supported_formattings:
+            raise ValueError(f"Unsupported formatting ''. Choose one of {supported_formattings}.")
 
         param_names = self._distr_parameters_for_repr()
         param_values = [
             get_repr_for_variable(getattr(dist, x), formatting=formatting) for x in param_names
         ]
 
-        if formatting == "latex":
+        if "latex" in formatting:
             param_string = ",~".join(
                 [fr"\mathit{{{name}}}={value}" for name, value in zip(param_names, param_values)]
             )
-            return r"$\text{{{var_name}}} \sim \text{{{distr_name}}}({params})$".format(
-                var_name=name, distr_name=dist._distr_name_for_repr(), params=param_string
+            if formatting == "latex_with_params":
+                return r"$\text{{{var_name}}} \sim \text{{{distr_name}}}({params})$".format(
+                    var_name=name, distr_name=dist._distr_name_for_repr(), params=param_string
+                )
+            return r"$\text{{{var_name}}} \sim \text{{{distr_name}}}$".format(
+                var_name=name, distr_name=dist._distr_name_for_repr()
             )
         else:
-            # 'plain' is default option
+            # one of the plain formattings
             param_string = ", ".join(
                 [f"{name}={value}" for name, value in zip(param_names, param_values)]
             )
-            return "{var_name} ~ {distr_name}({params})".format(
-                var_name=name, distr_name=dist._distr_name_for_repr(), params=param_string
-            )
+            if formatting == "plain_with_params":
+                return f"{name} ~ {dist._distr_name_for_repr()}({param_string})"
+            return f"{name} ~ {dist._distr_name_for_repr()}"
 
     def __str__(self, **kwargs):
         try:
@@ -201,7 +218,7 @@ def __str__(self, **kwargs):
 
     def _repr_latex_(self, **kwargs):
         """Magic method name for IPython to use for LaTeX formatting."""
-        return self._str_repr(formatting="latex", **kwargs)
+        return self._str_repr(formatting="latex_with_params", **kwargs)
 
     def logp_nojac(self, *args, **kwargs):
         """Return the logp, but do not include a jacobian term for transforms.

pymc3/model.py

@@ -1874,24 +1874,27 @@ def _walk_up_rv(rv, formatting="plain"):
             all_rvs.extend(_walk_up_rv(parent, formatting=formatting))
     else:
         name = rv.name if rv.name else "Constant"
-        fmt = r"\text{{{name}}}" if formatting == "latex" else "{name}"
+        fmt = r"\text{{{name}}}" if "latex" in formatting else "{name}"
         all_rvs.append(fmt.format(name=name))
     return all_rvs
 
 
 class DeterministicWrapper(tt.TensorVariable):
     def _str_repr(self, formatting="plain"):
-        if formatting == "latex":
-            return r"$\text{{{name}}} \sim \text{{Deterministic}}({args})$".format(
-                name=self.name, args=r",~".join(_walk_up_rv(self, formatting=formatting))
-            )
+        if "latex" in formatting:
+            if formatting == "latex_with_params":
+                return r"$\text{{{name}}} \sim \text{{Deterministic}}({args})$".format(
+                    name=self.name, args=r",~".join(_walk_up_rv(self, formatting=formatting))
+                )
+            return fr"$\text{{{self.name}}} \sim \text{{Deterministic}}$"
         else:
-            return "{name} ~ Deterministic({args})".format(
-                name=self.name, args=", ".join(_walk_up_rv(self, formatting=formatting))
-            )
+            if formatting == "plain_with_params":
+                args = ", ".join(_walk_up_rv(self, formatting=formatting))
+                return f"{self.name} ~ Deterministic({args})"
+            return f"{self.name} ~ Deterministic"
 
     def _repr_latex_(self):
-        return self._str_repr(formatting="latex")
+        return self._str_repr(formatting="latex_with_params")
 
     __latex__ = _repr_latex_
 

pymc3/model_graph.py

@@ -121,7 +121,7 @@ def update_input_map(key: str, val: Set[VarName]):
                     pass
         return input_map
 
-    def _make_node(self, var_name, graph):
+    def _make_node(self, var_name, graph, *, formatting: str = "plain"):
         """Attaches the given variable to a graphviz Digraph"""
         v = self.model[var_name]
 
@@ -146,7 +146,7 @@ def _make_node(self, var_name, graph):
         elif isinstance(v, SharedVariable):
             label = f"{var_name}\n~\nData"
         else:
-            label = str(v).replace(" ~ ", "\n~\n")
+            label = v._str_repr(formatting=formatting).replace(" ~ ", "\n~\n")
 
         graph.node(var_name.replace(":", "&"), label, **attrs)
 
@@ -181,7 +181,7 @@ def get_plates(self):
             plates[shape].add(var_name)
         return plates
 
-    def make_graph(self):
+    def make_graph(self, formatting: str = "plain"):
         """Make graphviz Digraph of PyMC3 model
 
         Returns
@@ -205,20 +205,20 @@ def make_graph(self):
                 # must be preceded by 'cluster' to get a box around it
                 with graph.subgraph(name="cluster" + label) as sub:
                     for var_name in var_names:
-                        self._make_node(var_name, sub)
+                        self._make_node(var_name, sub, formatting=formatting)
                     # plate label goes bottom right
                     sub.attr(label=label, labeljust="r", labelloc="b", style="rounded")
             else:
                 for var_name in var_names:
-                    self._make_node(var_name, graph)
+                    self._make_node(var_name, graph, formatting=formatting)
 
         for key, values in self.make_compute_graph().items():
             for value in values:
                 graph.edge(value.replace(":", "&"), key.replace(":", "&"))
         return graph
 
 
-def model_to_graphviz(model=None):
+def model_to_graphviz(model=None, *, formatting: str = "plain"):
     """Produce a graphviz Digraph from a PyMC3 model.
 
     Requires graphviz, which may be installed most easily with
@@ -228,6 +228,15 @@ def model_to_graphviz(model=None):
     and then `pip install graphviz` to get the python bindings.  See
     http://graphviz.readthedocs.io/en/stable/manual.html
     for more information.
+
+    Parameters
+    ----------
+    model : pm.Model
+        The model to plot. Not required when called from inside a modelcontext.
+    formatting : str
+        one of { "plain", "plain_with_params" }
     """
+    if not "plain" in formatting:
+        raise ValueError(f"Unsupported formatting for graph nodes: '{formatting}'. See docstring.")
     model = pm.modelcontext(model)
-    return ModelGraph(model).make_graph()
+    return ModelGraph(model).make_graph(formatting=formatting)