Improving Graph Visualization Functionality

[FraserP117]

The present model/graph visualization capability is highly limited. As per the example in the docs, the existing implementation uses a simple GraphPlot.gplot function to output a basic spring-layout of the model graph.

@chbe-helix and I; along with those in the RxInfer.jl working group at the Active Inference Institute aim to clarify the nature of the proposed improvements to the graph visualization procedure. We anticipate a small discussion here with @bvdmitri and any other BIASlab members, in regard to the clarification of the most valuable improvement/s to this functionality.

At present, we feel that the functionality should:

1. Depict node types.
2. Layout the graph along a "taxicab" grid, in place of a spring layout. This is simply how most TFFG graphs are depicted abstractly and so is a natural choice for the layout.
3. Adequately depict very large models, perhaps by means of plate notation or ellipses to denote the continuation of a pattern/motif.

Any further specifications/requirements are more than welcome!

[bvdmitri]

Hey @FraserP117, thanks for the clear description!

I'm on board with the draft of the specification. I'll start jotting down my thoughts and detailing the internal structures over the graph. It'd be great to see my comment expanded and documented in the official documentation down the line. I know @wouterwln is swamped with other tasks right now, but it's definitely something worth considering for the future. Also, feel free to add to my comment if I've overlooked anything.

Depict node types

At a glance, we've got two main types of nodes in our codebase: factor nodes and variable nodes. Note, however, that in our papers/research we usually work with Forney-style factor graphs, which have only one type of the node, namely factor nodes and variables are represented as edges (more on that below). Nevertheless, for the codebase, nodes come with its own set of sub-types and specific properties.

Factor nodes

• Factor nodes entail:
  • Neighbors, although this info can be derived from the graph structure
  • Their "type": Deterministic or Stochastic
  • A functional form attached to them, like Gaussian or Bernoulli
  • Factorization constraints over joint distribution, for instance, if a node f encompasses 3 variables x, y, z, we could encode various factorization constraints for the variational distribution q(x, y, z), such as q(x, y)q(z) or q(x)q(y, z) or q(x)q(y)q(z)
  • Any unstructured meta information, like the approximation method used to compute messages

Variable nodes

• Variable nodes entail
  • A name, obviously
  • They can be standalone or part of a collection, with a specific index, e.g. i = 1. Collections aren't necessarily vectors; they could be matrices or any dimension, in which case the index might look like i = 1, 1, 1, and so forth. This could be useful for implementing plate notation.
  • Their "kind", like random, data, or fixed (i.e., constants)
  • They might be anonymous or not. Anonymous variables are created under the hood for expressions like y ~ Normal(f(x), 1.0), which would create an anonymous variable anonymous anonymous ~ f(x) and then call y ~ Normal(anonymous, 1.0)
  • Functional form constraints, such as desiring the posterior distribution q(s) over a single variable s to follow a Gaussian distribution or Bernoulli. A functional form constraint can take any form (e.g. any Julia object).
  • A variable can be linked to other variables. For example, if f is a deterministic factor node, then y := f(x) establishes a deterministic "link" between x and y. Understanding these links is crucial for specifying constraints over seemingly different variables, as constraints on one might affect another.
  • Any unstructured meta information, such as the approximation method used to compute the posterior

Moreover, our factor graph structure supports nested modeling capabilities, introducing the notion of "sub-models," akin to sub-graphs. This is encoded in the GraphPPL.Context structure. I think the plotting capabilities should operate on Contexts and sub-graphs seriously. The current plotting just grabs the entire graph without separating the nodes in their contexts.

Layout the graph along a "taxicab" grid, in place of a spring layout. This is simply how most TFFG graphs are depicted abstractly and so is a natural choice for the layout.

Absolutely!

Adequately depict very large models, perhaps by means of plate notation or ellipses to denote the continuation of a pattern/motif.

Indeed, I believe there should be several options for displaying a particular model's graph. For instance, we may opt for a "full" display, "plate notation," or a "sub-region," like a region around a specific variable s (which could, in turn, be displayed "fully" or in "plate notation"). We could also expand or hide submodels, a.k.a sub-graphs.
Also as I've mentioned, in BIASlab we prefer to work with Forney-style Factor graphs. They do not have "variable nodes", but instead use edges to represent variables. For variables that are connected to more than 2 nodes, Forney-style Factor Graphs introduce equality nodes. This might be also an option (maybe even the default?) for the plotting.

This was just a stream of thoughts. I might have overlooked something. Looking forward to input from @wouterwln and other team members to enrich the discussion. I'll also post this issue in our university Slack channel to invite collaboration on high-level ideas and requests.

[FraserP117]

Excellent! thank you for the detailed suggestions @bvdmitri, especially the details on a GraphPPL.Model's internal structure.

I intend to assist in the integration of all of these features in due course. I think we would ultimately like to implement the standard as laid out in Part 1 of the Realising Synthetic Active Inference Agents papers? (on account of this being quite literally the graphical specification language).

Having said that, I think I'll begin by prioritising the following functionality:

1. The depiction of node types.
2. Defaulting to a TFFG representation (and hence using equality constraint nodes).
3. The use of ellipses to denote motif continuation in large models.

I think this would result in a visual experience akin to that of Figure 5.1 in your thesis - or Figure 2.1 of @MagnusKoudahl 's. Let me know what you think about this prioritisation. Is there any other subset of your above suggestions that you feel would be more appropriate to prioritise instead? Perhaps it's better to go with plate notation instead of ellipsis? Perhaps I should make the GraphPPL.Context visualisation paramount?

Thanks again!

[bvdmitri]

Let me know what you think about this prioritisation.

I think it's a great idea! However, using ellipses to show motif continuation might be tricky. Check out my next comment.

Perhaps it's better to go with plate notation instead of ellipsis?

Plate notation would indeed be nice. I believe we should start by focusing on creating a visually appealing TFFG for smaller models. We should also reserve some space in the graph plotting system to accommodate either plate or ellipsis notation. While working on implementing these, it would be helpful to receive feedback from you on what additions are needed in the GraphPPL core to support plate or ellipsis notation. We can consider adding extra metadata during graph creation to simplify this process.

Perhaps I should make the GraphPPL.Context visualisation paramount?

That sounds reasonable, yes!

I intend to assist in the integration of all of these features in due course. I think we would ultimately like to implement the standard as laid out in Part 1 of the Realising Synthetic Active Inference Agents papers? (on account of this being quite literally the graphical specification language).

I'm fully on board with that. I particularly appreciate how Magnus and Thijs illustrated factorization constraints on the graph, as seen in Figures 5, 6, and 7 with those circles inside.

[wouterwln]

Hey @FraserP117 , thanks for opening up the discussion on the graph visualisation and concretising the requirements. I'm a big fan of having a visualization like Thijs' and Magnus' notation. However, you should realize that in base GraphPPL not every model has a set of inference constraints (we export it as an optional plugin that is always activated in RxInfer), so I would keep into consideration that the method should also work when there are no inference constraints in the node.

Perhaps I should make the GraphPPL.Context visualisation paramount?

That sounds like a good idea to me, the Context contains all node labels (factors and variables) that exist on a specific level in the hierarchy, it would be nice to visualize submodels.

I very much agree with your taxicab idea, that is usually how factor graphs are drawn anyway. The main challenge I see is designing which node goes where. I've thought about this, and came to a simple heuristic: Most people want their generative models visualized in the same way in which they define them: top to bottom, left to right. I realized that we actually save the order of construction of all nodes (the GraphPPL.NodeLabel serves as a universal identifier in the graph for all nodes and contains an id integer that is the counter in the global graph), so you have an idea of which nodes were created before others. This is only part of the story, however, since a variable node can be created first and then passed into a submodel as an interface, where we wouldn't necessarily want it drawn top left of this submodels visualization. However, this is maybe a nice first start to design an algorithm that gives a somewhat nice visualization.

We're also very much not hung up on GraphPlot, I just picked it because it seemed convenient to work with, but I don't see a lot of customization options. You can look at https://juliagraphs.org/Graphs.jl/stable/first_steps/plotting/ to get an idea of other plotting libraries, although I also don't know if they give you the customization options necessary.

As for your prioritization, I think it's fully correct, maybe I would add "visualization of constraints" as well because I think that is really separate from the TFFG visualization in general. If there's anything about GraphPPL's internals you want to know about, feel free to reach out :)

[bvdmitri]

I just also would like to emphasize that we are open to change some internals, modify our data structures or add extra functionality to the core of GraphPPL if that helps with the visualization. There is no need to constrain ourselves in the current implementation. Of course, we better avoid major refactoring, but small changes or/and additions are fine.

[FraserP117]

Wonderful. Thanks @bvdmitri for those further suggestions/clarifications, and thanks to @wouterwln for all that detail, it's greatly appreciated!

Yes, I think I'll push the implementation of plate notation above that of ellipses; in the ranking of immediate priorities. In our discussion thus far, I think we've settled on these desired features and functionalities:

1. Visualization should work regardless of the specification/absence of any inference constraints.
   1. We want a visually appealing TFFG model depiction: "full display".
      1. Aim to implement the visualization standard/s laid out in Realising Synthetic Active Inference Agents: Part 1. Most particularly, we should:
         1. Depict node types.
         2. Depict the graph with a "Taxicab" layout, using @wouterwln's heuristic for now. There will be issues when it comes to depicting submodels. We'll have to address that later perhaps, by playing around with GraphPPL.Context?
         3. Depict the form and factorisation constraints.
   2. Additionally depict the information contained in the GraphPPL.Context struct.
   3. Add the option to specify plate notation, this will be especially useful for depicting larger models.
   4. Afford "sub-region" display, like a region around a specific variable s (which could, in turn, be displayed "fully" or in "plate notation").
   5. Afford the ability to expand or hide any existing sub-models.

I have attempted to loosely indicate a feature's importance by its order in the list - more important items come first. Now this is quite a lot to tackle all at once. Consequently, I think it may be wise to prioritise everything up until section 1.i.a.c, in addition to sections 1.ii and 1.iii and in this order of importance.

As regards the use of GraphPlots, it's good to know that you're not "hung up on it". I don't mean to say that it shouldn't be used, but yes: I have also noted a lack of customisation options for GraphPlots - at least to the extent that we desire here.

[MagnusKoudahl]

@FraserP117 I'll be happy to help out with this so we get it right :) I'm not actively involved with RxInfer as much these days - but if we're doing CFFG's I'm definitely in

[FraserP117]

Fantastic! I'd be honoured to have your assistance @MagnusKoudahl. @chbe-helix and I are going to be meeting in the next day or so to begin our approach to all this. I am finishing up my tour of the GraphPPL materials in that time. When we really "get going", I'd love to discuss things and otherwise collaborate. That goes for absolutely anyone else too!

I agree, I absolutely do want to "get this right". I'll be erring on the side of annoying everybody with questions as opposed to powering ahead with what I think I understand. I also intend to make it known when I don't understand something (which will be rather frequent).

I think the biggest challenge for myself and @chbe-helix just now is mapping out all the moving parts. Not a big challenge, but one must get sufficiently oriented first.

[wmkouw]

Hi @FraserP117 , great to hear that you are so enthusiastically undertaking this project! Being able to properly generate a visual FFG representation has been on the lab's wishlist for years. :D

I would like to add something to the list. I teach the probabilistic programming sections of our master course where students use our lab's tools to automatically do Bayesian inference based on a specific probabilistic model. We used to use ForneyLab.jl which plotted graphs using Graphviz. For example:

Students are often confused by what messages are, where messages are located and where message collisions take place. @ThijsvdLaar implemented a feature that lets you isolate an edge and visualize the messages that belong to that edge:

This did a number of things:

1. it shows students which messages are relevant to a particular variable,
2. it shows what other variables a particular message depends on (via the other edges of the connected node),
3. it provides an identifier for a message that lets them plot its functional form.

It was very helpful for the students to learn message passing, helpful for us for debugging and I think it could be helpful to developers that want to get started with RxInfer as well.

RxInfer is built differently (ForneyLab is declarative, RxInfer is reactive), so this feature was not something that could be easily ported. But it should definitely be possible: user specifies a variable, algorithm performs a small graph walk to find connected nodes / 1-step removed variables (i.e. feature 1.iv in your list), and pipes subscriptions to a logger. The relevant variables, nodes and message subscriptions should be plotted in the normal way and there would need to be some filtering to make the entire thing human-readable. But that's it, I think.

To be clear: this is not a high-priority item. As mentioned in your last post, the current list is already quite extensive and will take a while. But I just wanted to have mentioned it.

I am of course more than happy to work with you on this.

[FraserP117]

Thank you @wmkouw. This is a functionality that I think would be fantastic to include. Being able to quickly visualize parts of the graph and to otherwise palpate the individual messages - and when they're passed - would be very helpful.

@chbe-helix will be meeting on Wednesday to cement our collaborative approach. We'll be back here to post that and to query any further issues that arise.

[ThijsvdLaar]

Hi @FraserP117, great to hear that you're willing to have a jab at graph visualization. I think the list of priorities looks good: first make it work, then make it right.

As @wmkouw already mentioned, the graph visualization for ForneyLab was very rudimentary. But even then, it was useful for getting visual feedback during model development, which I think is the main use case. It often helped me catch unexpected cycles (mostly in subgraphs relating to variational factors) and to track root causes of NaN values in the messages.

For me, a main grievance with the ForneyLab visualization was the layout, which made it difficult to find the node you were looking for. The taxicab grid suggestion by @wouterwln might already help here. We have an (unwritten) convention for drawing FFG's, with temporal thickness horizontally and hierarchical depth vertically. Dynamic variables are usually stored in a vector, so it might make sense to visualize these horizontally, layered vertically in order of declaration. Even without advanced features (nested, plate, etc.), this could already help a user to find the node/message they're looking for.

[bvdmitri]

@FraserP117 @chbe-helix do have you any updates regarding this issue for the upcoming RxInfer public meeting? :) Such that I can maybe prepare some slides in advance

[FraserP117]

Hello @bvdmitri ,

Yes we do. @chbe-helix and I have bounced some ideas between us, with the help of all the assistance generously provided here. We've come up with an approach that we think will satisfy the implementation requirements, while also minimising technical debt. I have a very basic "proof-of-concept" demo to show at the meeting and I'll try to get it to @bvdmitri ASAP; for inclusion into any slides - as you mention. I have not included it here just yet, although I will endeavour to do so before my night is up. What follows is an update regarding the approach that @chbe-helix and I have taken thus far.

Our main preoccupation until very recently has been to settle on an approach that would afford maximum extensibility/customizability, while also minimising technical debt and overall complexity. To that end, we have settled on the following approach - which is of course subject to revision at any time!

1. GraphViz

We think it is best to do-away with GraphPlot.jl. As mentioned by @wouterwln, GraphPlot.jl's customisation options are minimal and we have in mind a rather comprehensive set of visualization capabilities.

In place of GraphPlot.jl, we think that GraphViz is the best way to go, for two reasons. First, it is certainly a tried-and-tested package that has very much withstood the test of time. Second, GraphViz offers - perhaps - the greatest possible control/customizability over the visualization process. There are two particularly useful Julia packages which interface with GraphViz to allow the specification of an arbitrary graph with either native DOT language syntax from within Julia: GraphViz.jl, or a Julia-lang interface to GraphViz DOT syntax: GraphvizDotLang.jl. We're interested to hear more from @wmkouw and @ThijsvdLaar on the nature of the GraphViz visualizations for Forneylab.jl - owing to its use of GraphViz.

2. Display Modalities

We are very much interested in developing all the above-mentioned display capabilities - in this issue - with GraphViz. We have decided to focus on the implementation of node/edge display and the first layout style - the "Raw" style. We will hold off on the development of the other layout options, the "Display Concatenation Options" and the "Sub-Region Display", until we have satisfactorily developed node/edge-depiction and the "Raw" layout modality with GaphViz. We think this will serve as a useful prototype and "base of operations" for the full suite of visualization options to which we aspire. At the risk of being premature, we hope to develop the following list of loosely-categorised visualization options:

2.1 Depiction of Node/Edge Information

This includes specification of the node shape for factors/variables, the symbolic label for each node and shading options to denote observability.

2.2 Layout Style

1. "Raw": Non-Forney style factor graph display. This would simply represent the exact GraphPPL.Model, as it exists in memory - very much like the existing implementation. Although the "mere factor graph" representation is not particularly common, we feel that it is still useful for developers to see "exactly" what the underlying GraphPPL.Model looks like. It is also the easiest way to visualize the graph and will serve as a nice "proof-of-concept" for the development of the other layouts.
2. "FFG": This would display the GraphPPL.Model as a proper Forney-style factor graph, with equality-constraint nodes and all.
3. "CFFG": Same as "FFG", only now the constraints are also visualized - should there be any. This would implement the standards as specified in Realising Synthetic Active Inference Agents: Part 1.

2.3 Display Concatenation Options

These options would simplify the visualization of the graph in some specified way, such as using plate notation to encapsulate repeating patterns/motifs, or the choice to "show"/"hide" sub-models to some specified depth. "FFG"/"CFFG" layout options could also be considered as a concatenation option.

2.4 Sub-Region Display

This last option would allow for the visual palpation of the induced sub-graph on an arbitrary node/edge in the "raw" GraphPPL.Model or the "FFG"/"CFFG" representation.

3. Implementation Approach

In a nutshell, the task as we conceive it breaks down into the following steps:

1. A method of parsing the desired/relevant information contained in a particular GraphPPL.Model.
2. The subsequent (automatic) construction of the GraphViz DOT code to visualize that information as a GraphViz graph.
3. Options to save/display the visualization specified by the generated GraphViz DOT code.

It seems like all the relevant information to this end is contained in the model's Context structure. I'm very interested in checking my understanding on that point.

4. Implementation Considerations

At present, I have a very minimal working example of automatic DOT code generation for the Coin-toss model
. My current approach generates a string-representation of the GraphViz DOT code from a GraphPPL.Model and then uses eval() and Meta.parse() to execute the string as a command in the GraphViz.jl package.

As you all well know, evaluating arbitrary strings as code with eval(Meta.parse(my_string)) is a very "spooky" thing to be doing, so I am not married to my current approach.

There is also the above-mentioned GraphvizDotLang.jl package, which does not use the doc-string command specification as in GraphViz.jl. GraphvizDotLang.jl allows you to "Create Graphviz graphs straight from Julia [code]". Hence, it might be best to use Julia's meta-programming capabilities to auto-generate GraphvizDotLang.jl code which can then be executed to yield the final visualization. I'm certainly open to suggestions!

5. Takeaway

Our primary aim right now is to implement a satisfactory means of visualizing the raw GraphPPL.Model as a GraphViz plot. We think this means that we need to automatically generate the relevant GraphViz DOT code from the information encoded in model's Context structure. Please do tell us your thoughts on this!

Many thanks to all who have taken the time to engage thus far. I look forward to working on this with any/all of you.

[bvdmitri]

Very much appreciated for the extensive overview @FraserP117

[bvdmitri]

Something that I found out just today:

• https://research.google/blog/model-explorer/
• https://ai.google.dev/edge/model-explorer