Add Model.to_graphviz shortcut (#6865)

* add methods

* add test structure

* fix docstrings

* remove networkx method

* imptove tests

* rename method name

* add mock test

* Update tests/test_model.py

Co-authored-by: Michael Osthege <[email protected]>

* update docsting

Co-authored-by: Ricardo Vieira <[email protected]>

---------

Co-authored-by: Michael Osthege <[email protected]>
Co-authored-by: Ricardo Vieira <[email protected]>

Author: 3 people

pymc/model.py

@@ -24,6 +24,7 @@
     Any,
     Callable,
     Dict,
+    Iterable,
     List,
     Literal,
     Optional,
@@ -65,6 +66,7 @@
 from pymc.initial_point import make_initial_point_fn
 from pymc.logprob.basic import transformed_conditional_logp
 from pymc.logprob.utils import ParameterValueError
+from pymc.model_graph import VarName, model_to_graphviz
 from pymc.pytensorf import (
     PointFunc,
     SeedSequenceSeed,
@@ -1879,6 +1881,53 @@ def debug_parameters(rv):
         elif not verbose:
             print_("You can set `verbose=True` for more details")
 
+    def to_graphviz(
+        self, *, var_names: Optional[Iterable[VarName]] = None, formatting: str = "plain"
+    ):
+        """Produce a graphviz Digraph from a PyMC model.
+
+        Requires graphviz, which may be installed most easily with
+            conda install -c conda-forge python-graphviz
+
+        Alternatively, you may install the `graphviz` binaries yourself,
+        and then `pip install graphviz` to get the python bindings.  See
+        http://graphviz.readthedocs.io/en/stable/manual.html
+        for more information.
+
+        Parameters
+        ----------
+        var_names : iterable of variable names, optional
+            Subset of variables to be plotted that identify a subgraph with respect to the entire model graph
+        formatting : str, optional
+            one of { "plain" }
+
+        Examples
+        --------
+        How to plot the graph of the model.
+
+        .. code-block:: python
+
+            import numpy as np
+            from pymc import HalfCauchy, Model, Normal, model_to_graphviz
+
+            J = 8
+            y = np.array([28, 8, -3, 7, -1, 1, 18, 12])
+            sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18])
+
+            with Model() as schools:
+
+                eta = Normal("eta", 0, 1, shape=J)
+                mu = Normal("mu", 0, sigma=1e6)
+                tau = HalfCauchy("tau", 25)
+
+                theta = mu + tau * eta
+
+                obs = Normal("obs", theta, sigma=sigma, observed=y)
+
+            schools.to_graphviz()
+        """
+        return model_to_graphviz(model=self, var_names=var_names, formatting=formatting)
+
 
 # this is really disgusting, but it breaks a self-loop: I can't pass Model
 # itself as context class init arg.

tests/test_model.py

@@ -17,6 +17,8 @@
 import unittest
 import warnings
 
+from unittest.mock import MagicMock, patch
+
 import arviz as az
 import cloudpickle
 import numpy as np
@@ -47,6 +49,7 @@
 from pymc.logprob.basic import conditional_logp, transformed_conditional_logp
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Point, ValueGradFunction, modelcontext
+from pymc.model_graph import model_to_graphviz
 from pymc.util import _FutureWarningValidatingScratchpad
 from pymc.variational.minibatch_rv import MinibatchRandomVariable
 from tests.models import simple_model
@@ -1653,3 +1656,28 @@ def test_model_logp_fast_compile():
 
     with pytensor.config.change_flags(mode="FAST_COMPILE"):
         assert m.point_logps() == {"a": -1.5}
+
+
+class TestModelGraphs:
+    @staticmethod
+    def school_model(J: int) -> pm.Model:
+        y = np.array([28, 8, -3, 7, -1, 1, 18, 12])
+        sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18])
+        with pm.Model(coords={"school": np.arange(J)}) as schools:
+            eta = pm.Normal("eta", 0, 1, dims="school")
+            mu = pm.Normal("mu", 0, sigma=1e6)
+            tau = pm.HalfCauchy("tau", 25)
+            theta = mu + tau * eta
+            pm.Normal("obs", theta, sigma=sigma, observed=y, dims="school")
+        return schools
+
+    @pytest.mark.parametrize(
+        argnames="var_names", argvalues=[None, ["mu", "tau"]], ids=["all", "subset"]
+    )
+    def test_graphviz_call_function(self, var_names) -> None:
+        model = self.school_model(J=8)
+        with patch("pymc.model.model_to_graphviz") as mock_model_to_graphviz:
+            model.to_graphviz(var_names=var_names)
+            mock_model_to_graphviz.assert_called_once_with(
+                model=model, var_names=var_names, formatting="plain"
+            )