pymc.model_graph.model_to_networkx

pymc.model_graph.model_to_networkx(model=None, *, var_names=None, formatting='plain', node_formatters=None, include_dim_lengths=True)

Produce a networkx Digraph from a PyMC model.

Requires networkx, which may be installed most easily with:

conda install networkx

Alternatively, you may install using pip with:

pip install networkx

See https://networkx.org/documentation/stable/ for more information.

Parameters:

modelModel

The model to plot. Not required when called from inside a modelcontext.

var_namesiterable of str, optional

Subset of variables to be plotted that identify a subgraph with respect to the entire model graph

formattingstr, optional

one of { “plain” }

node_formattersdict, optional

A dictionary mapping node types to functions that return a dictionary of node attributes. Check out the networkx documentation for more information how attributes are added to nodes: https://networkx.org/documentation/stable/reference/classes/generated/networkx.Graph.add_node.html

include_dim_lengthsbool

Include the dim length in the plate label. Default is True.

Examples

How to plot the graph of the model.

import numpy as np
from pymc import HalfCauchy, Model, Normal, model_to_networkx

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)

model_to_networkx(schools)

Add custom attributes to Free Random Variables and Observed Random Variables nodes.

node_formatters = {
    "Free Random Variable": lambda var: {"shape": "circle", "label": var.name},
    "Observed Random Variable": lambda var: {"shape": "square", "label": var.name},
}
model_to_networkx(schools, node_formatters=node_formatters)