Kruschke diagrams as an option in model_to_graphviz or with a separate model_to_kruschke

[opherdonchin]

Kruschke plots, as used in Doing Bayesian Data Analysis, and as described in this blog post are preferred by many users to the style of diagramming produced by model_to_graphviz. I've looked a bit at the code, and I don't think it should be hard to modify it to create Kruschke diagrams.

I'm scared to start on my own. I'd like a collaborator or a mentor to help me go through the process step by step.

I'm interested in hearing from people who would be interested and also on general thoughts both about the idea and about how to go about the actual programming.

Thanks!
Opher

[ricardoV94]

@opherdonchin the biggest challenge is how to include plots somewhere in graphviz / another visualization library. Getting draws / a representation of the distribution would be easy.

If you can figure that out, you'll be more than halfway there

[ColCarroll]

We talked about doing something similar in arviz (as a way of comparing prior / posterior), but I guess doing it in pymc would give you access to the graph structure. @aloctavodia might have some thoughts from preliz work.

The reason graphviz has been used is that it is the only reasonable graph layout engine. Since graphviz allows images to be attached to nodes, a reasonable first pass at a design might be:

1. generate a reasonable pdf/pmf for each distribution
2. generate and save these using matplotlib to a tmpdir
3. use the existing graphviz machinery to display these images instead of (or as well as!) the latex
4. iterate?

[aloctavodia]

Kruschke diagrams use "predefined" distributions/shapes, for instance, all Beta distributions are represented with Beta(2, 2) (or something similar) independent of the actual parameters in the model. Not sure if you are suggesting a different route, i.e. dynamically computing the shapes/images, instead of using predefined ones. Using predefined images is not only simpler, it's usually preferred as it will be easier to get the "structure" of the model at first glance.

For manually doing Krusche diagrams I have been using this libre office template https://www.sumsar.net/blog/2013/10/diy-kruschke-style-diagrams/
The PreliZ's distribution gallery is another possible source of predefined images https://preliz.readthedocs.io/en/latest/gallery_content.html

[ricardoV94]

Do the colors have a meaning?

[aloctavodia]

No, they are randomly assigned from a set of 4 colours or so. Using all the same colour looks kind of dull. But I not entirely happy with randomess either. Suggestions are welcome. Maybe we could use it to convey information about boundaries, that are already represented, but a colour will make it more prominent.

[opherdonchin]

Wouldn't it be enough to give the colors a relatively basic meaning like "computed node" or "distribution node" or "data"? That would be useful, I think.

Eventually, the code could allow users to define colors when defining the model and let that override the default. However, I would expect that would mean updating model_to_graphviz as well.

[opherdonchin]

Certainly, we could take templates from either Preliz or the LibreOffice templates. The challenge I see here are how best to link distributions to images. I don't know if this is a big refactor or if it is desirable. Ideally, someone creating their own new distribution should be able to specify an image to use for it in the graph. Perhaps they should be allowed to override the default image defined for the existing distributions.

A second challenge is to represent computed nodes. The computation isn't usually shown in traditional diagrams, but is generally part of the Krushke diagrams. Can we turn a computed node in a PyTensor graph into Latex? (I found a pymc.printing in v5.5 but I didn't see it in the current version) Manually drawn Krushcke diagrams also have the arrow from the priors of parameters in the computation pointing to the relevant variable. Having the computation and the appropriately directed arrow makes the diagrams much more accessible.

Finally, Kruschke diagrams are not well specified when it comes to dimensions and shapes, at least as far as I know. When I teach, I use subscripts indicating dimensions on variable names. Perhaps this would work? I also add them to the arrows, but maybe that's not necessary.