Allow customizing style of model_graph nodes

[wd60622]

Description

Allow users to override the default graphviz by passing mapping from node type to function that creates the node kwargs.

The default behavior is the same but now the user can override based on the node type defined:

• Data
• Free Random Variable
• Observed Random Variable
• Deterministic
• Potential

User needs to define function from node variable to kwargs passed to graphviz / networkx

def node_formatter(var: TensorVariable) -> dict[str, Any]:
    return {"label": var.name}

then pass a mapping from node type to node_formatter.

node_mapping = {"Data": node_formatter}
pm.model_to_graphviz(model, node_formatters=node_mapping)

Example

import pymc as pm 

with pm.Model() as model: 
    a = pm.Normal("a")
    b = pm.Normal("b", mu=a)
    c = pm.Normal("c", mu=a)
    d = pm.Normal("d", mu=b + c, observed=0)

default = pm.model_to_graphviz(model)

simple_random_variable_formatter = {
    "Free Random Variable": lambda var: {"shape": "circle", "label": var.name}, 
}
pm.model_to_graphviz(model, node_formatters=simple_random_variable_formatter)

fancy_formatter = {
    "Free Random Variable": lambda var: {"shape": "polygon", "sides": "7", "label": var.name, "style": "dashed"}, 
    "Observed Random Variable": lambda var: {"shape": "circle", "label": var.name, "style": "solid"},
}
pm.model_to_graphviz(model, node_formatters=fancy_formatter)

Default:

Simple:

Fancy:

Related Issue

• Closes #
• Related to Allow more minimal graphviz representation #6716

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7302.org.readthedocs.build/en/7302/

[wd60622]

@drbenvincent might be interested in this example

with pm.Model() as model: 
    c = pm.Normal("c")
    z = pm.Normal("z", mu=c)
    y = pm.Normal("y", mu=c + z)

node_formatters = {
    "Basic Random Variable": lambda var: {"shape": "circle", "label": var.name},
}
model_graph.model_to_graphviz(model, node_formatters=node_formatters)

[ricardoV94]

Looks sweet. Left some minor suggestions.

[ricardoV94]

Small optimization, most models don't have potentials, and most have many free_RVs, and deterministics, so I would suggest changing the order of this to be something like if deterministics elif free rvs elif observed rvs elif data elif potentials instead

[wd60622]

Thank you for the review @ricardoV94
Will add these changes and add some tests

[wd60622]

Any thoughts on to tests while accounting for graphviz putting the nodes in random order in the string? @ricardoV94

[wd60622]

Think it might have something to do with our side actual and the set ordering. graphviz python source code for reference

[wd60622]

but not 100% sure about that

[wd60622]

Some random / unrelated tests seem to be randomly failing. Latest run passed