Better visualizations of Memory addon internal structure

[bvdmitri]

From discussions with @ismailsenoz, we realized that AddonMemory doesn’t present results in the most user-friendly way, though its structure can be extremely helpful for debugging complex models. We have documentation on this here.

Given our growing experience with plotting graph structures, could we automate the generation of visual representations? Essentially, we should be able to create something like this:

directly from this:

@FraserP117 @kobus78 @chbe-helix, would this be something you’d be interested in working on? It seems like a valuable and highly useful addition for debugging. We would just need to define a new show_png method here, or potentially use MIME types similar to graph plotting (e.g., mime/png).

Would love to hear your thoughts!

[FraserP117]

Absolutely @bvdmitri!

We want to focus on delivering the most relevant and desirable functionality first. We have begun Phase 2: the theme of which is summary visualizations - things like plate notation and sub-model and connected component visualisation. We could prioritise the implementation of the visualisation for AddonMemory, if that sounds like a better plan?

I'll have to look into this but I suspect that the use of MIME types: mime/png or svg? is probably going to be more in line with the existing method of visualising everything? Note sure what @chbe-helix or @kobus78 think on this however.

I am personally going to be extremely preoccupied for the next week with my move to Glasgow on the 26th of February. I want to make progress on this before then however.

[chbe-helix]

I really like this direction. It seems having both formats, the text and visual, would be highly beneficial. @FraserP117 let's discuss this next Friday if you are available.

[bvdmitri]

I'll have to look into this but I suspect that the use of MIME types: mime/png or svg? is probably going to be more in line with the existing method of visualising everything?

I like mime types and if necessary it still can be displayed in text where svg/png is not supported. Similarly to graph visualizations indeed. Good luck with your move to Glasgow @FraserP117 :)

[abpolym]

Reading the issue, what was the progress since February or what is the status? Noone is assigned to this one, so does that mean someone else can jsut start hacking away on this issue on a new branch? Or is there a branch already?

[FraserP117]

Hello @abpolym.

I had started working on it, though from what @bvdmitri has said, there will be large changes coming to the AddonMemory functionality itself that will probably invalidate a lot of what I've done. Fortunately/unfortunately, I haven't done much on it.

I've therefore been hesitant to devote more time to building out my existing branch. @bvdmitri are there any more details you can share on the nature of the changes being made to the AddonMemory functionality?

I'm very keen to renew work on this and to deliver the functionality. @chbe-helix and I met the other week for the first time in a while, to renew our efforts on this. I think we need to consult with you on the status/nature of the changes before we commit to anything further on this downstream functionality.

Hope that helps.

[wouterwln]

Hey @fonsp , just tagging you in this issue since you were playing with AddonMemory today, maybe you want to collaborate on a joint effort regarding message visualization/inspection?

[abpolym]

Oh yeah that would be nice @fonsp ! I hacked a prototype Pluto notebook together this morning to visualize messages (marginals and result) from AddonMemory, maybe I can show you tomorrow if you are at the office? ❤️

[FraserP117]

Excellent!

I'd love to see your Pluto stuff @abpolym. Perhaps we could have a short video call sometime this week?

Regarding my own efforts, they've been limited to producing GraphViz renders of the underlying model that depict the result on the relevant edge of the graph. I additionally depict the information as to source and destination ports. I've not been actively working on this as of late, though I just pushed what I have to my fork of GraphPPL. See the feature/graphviz-mem-addon branch.

In principle, any/all fields of the AddonMemoryMessageMapping are amenable to visualization in this way. I've just started with result as per the doc's hand-crafted example; the one where the coloured result fields are depicted on the edges of the FFFG.

This approach is fine, but for a TensorBoard-style app we absolutely do want to go with the creation of something interactive via Pluto.jl. I'm very excited about potentially contributing to that! My extension was never intended to provide interactive visualizations. The primary intent is to afford high-quality, static images of the underlying Factor Graph, then the FFG/TFFG, in addition to debug and summary statistics. Only the former is actually implemented just now, though the feature/graphviz-mem-addon branch takes a step towards debug affordances.

To actually support the optional visualization of ReactiveMP's AddonMemory (when show_memory=true in the extension's GraphViz.load) I made ReactiveMP a weakdep of the extension. The code is incredibly rough and not at all nice to look at, it can be dramatically improved!

A simple test visualization of the result field of the AddonMemoryMessageMapping - with the Beta-Bernoulli coin toss model - would be:

using RxInfer
using GraphViz

# define the model
@model function coin_model(y, a, b)

    # We endow the θ parameter of our model with a prior
    θ ~ Beta(a, b)
    # note that, in this particular case, the `Uniform(0.0, 1.0)` prior will also work.
    # θ ~ Uniform(0.0, 1.0)

    # here, the outcome of each coin toss is governed by the Bernoulli distribution
    for i in eachindex(y)
        y[i] ~ Bernoulli(θ)
    end
end

# Create an RxInfer.ConditionedModelGenerator - conditioned on two Bernoulli trials (observations)
conditioned = coin_model(a = 2.0, b = 7.0) | (y = [ true, false], )

# Create the ProbabilisticModel
rxi_model = RxInfer.create_model(conditioned)

# perform inference:
result = infer(
    model = coin_model(a = 2.0, b = 7.0),
    data  = (y = [true, false, true, true],),
    addons = (AddonMemory(),) # make sure to pass the addon!
)

# extract the underlying GraphPPL.Model
gppl_model = RxInfer.getmodel(rxi_model)

# visualize the graph with the `result` field of `AddonMemoryMessageMapping`
graph = GraphViz.load(
    gppl_model;
    strategy         = :simple, # `:bfs` or `:simple` - for now 
    inference_result = result, # must pass the inference result
    show_memory      = true, # <-- REQUIRED
    formatter        = x -> string(x), # optional
    label_attr       = :label,  # edge label attr (I recommend sticking with `:label` and not `:xlabel`). This needs further work
    layout           = "neato", # `neato` and `fdp` work quite well
    save_to          = ".../coin_model_graph_fig.svg" # path/name of the final image
)

The resulting visualization is not very nice right now but the pluming/idea is there to be improved - if we think this is a worthwhile direction. @chbe-helix, and @abpolym, I'm keen on your thoughts! though I do agree that we should be focusing on the Pluto.jl implementation. As I say: very keen to help out on that.

[abpolym]

I will definitely look at your branch @FraserP117 ! What would be really cool is to take a representation of the FFFG like your GraphViz representation and then indeed display the results as distributions on each edge (so like the visualization on https://rxinfer.com/ but in real).

My Pluto notebook currently has a slider to select an iteration and one to select an update within an iteration, then plots the messages and results. So the idea with that was to be able to debug the inference process step by step and be able to go back and forth - and I found that Pluto is an amazing candidate to realize that! If we combine our efforts I think we can make something like an interactive RxInferBoard 😄 . Let's talk on Slack about when to call each other 😁 .

[FraserP117]

Absolutely @abpolym. Sounds like a plan!

[wouterwln]

Hi! We are moving to a cleaner Epic -> Feature -> Task issue hierarchy to better organize our backlog. This issue is currently either underspecified or not tagged appropriately.

To keep this issue open, please do the following within the next 7 days (by 25-11-2025):

1. Update/Replace: Ensure the description is clear and actionable.
2. Tag Correctly:For Tasks/Features, add the correct label (e.g., feature, task) AND include a link to the Parent Epic or Feature it belongs to.
3. For Bugs, add the Bug label. (Bugs do not require a parent link.)

Issues not updated, linked, or tagged correctly by the deadline will be closed and purged.

Thank you for helping us clean up and organize our backlog!

[wouterwln]

Hey guys, I'm closing this issue for now. I think this discussion would be better consolidated in the RxInferBoard Epic such that we might be able to break it down a bit into tasks and hopefully get more people involved (And get to a working MVP quicker!) I'm closing this for now, but I'm a big fan of a joint effort towards a beta version of RxInferBoard that allows you to inspect messages!

[FraserP117]

Excellent, I agree @wouterwln. Apologies for my delay, I'll pick up the conversation - and development efforts - in the RxInferBoard Epic.