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 #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

• [x] 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)

[wd60622]

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

[wd60622]

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