Improve plot_summary with overlay support, colors, labels, grid, and title (#1733)

[Abraham-y]

Closes #1733
This PR improves the plot_summary function with the following changes:
New parameters:

• overlay: plots all tags on a single axes for easy comparison (e.g., training vs. validation loss)
• colors: per tag color specification
• labels: custom legend labels
• legend: toggle legend visibility
• grid: toggle grid lines
• title: subplot or figure titles

All changes are backward compatible. Existing calls to plot_summary behave identically.
Also includes a test suite and a short how to guide notebook under the visualization docs.
AI disclosure: I used Claude for assistance in planning these changes. Every line of code was reviewed and tested by me.

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.88%. Comparing base (cfed0a1) to head (6761431).

Files with missing lines	Patch %	Lines
sbi/analysis/tensorboard_output.py	90.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1814      +/-   ##
==========================================
+ Coverage   87.86%   87.88%   +0.01%     
==========================================
  Files         140      140              
  Lines       12726    12753      +27     
==========================================
+ Hits        11182    11208      +26     
- Misses       1544     1545       +1     

Flag	Coverage Δ
fast	82.75% <90.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
sbi/analysis/tensorboard_output.py	88.18% <90.00%> (+2.63%)	⬆️

[janfb]

Great addition @Abraham-y , thanks!

This already looks good overall, but I suggest a bit of refactoring, see comments below.

[janfb]

it's better practice to do an if-else and raise a suitable error here, e.g., RuntimeError or ValueError with an informative error message.

[janfb]

I think we will have a kwargs problem here: both branches of these if-else case always pass explicit color=color and label=_labels[i] keyword arguments to ax.plot() as well as **(plot_kwargs or {}).
Thus, if a user passes both plot_kwargs={"color": "red"} or plot_kwargs={"label": "foo"}, it raises TypeError: got multiple values for keyword argument 'color'.

The fix would be merging all kwargs into a single dict before passing to ax.plot(), where we implement a hierarchy: explicit colors/labels params over plot_kwargs, over tag-name defaults. I suggest implement a short helper function, e.g., _build_plot_kwargs that does this and then returns the final kwargs to be passed to ax.plot()

[janfb]

can we move these if-else cases outside of the overlay if-else case? this would reduce code duplication, plus at the moment the legend would not be rendered for the non-overlay case (or is this on purpose?).

[janfb]

typo: instead.

[janfb]

I made a couple of comments on the if-else structures below. but more generally, I think we can avoid it entirely. I suggest to construct a list of tags which can be a list with a single tag as well (for the overlay=False case). then you need just a single loop over this list for rendering the figure.

I suggest: Before creating the figure, convert the flat tags list into subplot_tags, a list of lists where each inner list is the group of tags for one subplot. In overlay mode, this is [tags] (one subplot, all tags). In non-overlay mode it's [[tag] for tag in tags] (one subplot per tag). Similarly, precompute subplot_ylabels: in overlay mode, join the ylabels if they differ (e.g. "training_loss / validation_loss"), otherwise use them directly. Does this make sense?

You would still need the helper function _build_plot_kwargs(tag_idx) that constructs the keyword arguments dict for a single ax.plot() call. It starts from a copy of plot_kwargs (or empty dict), then overrides color if colors was provided, overrides label if labels was provided, and falls back to the tag name as label if neither labels nor a plot_kwargs['label'] key exists. This avoids the duplicate-keyword problem (see comment below).

Overall, with this approach you can have a single loop and we avoid the current code duplication in the two branches.

[janfb]

let's please use plain test functions instead of these test classes, unless we really need the classes, e.g., when we share init or fixtures within classes.

[janfb]

good to this! a couple of comments:

• let's try to keep it as concise as possible but make the how-to-guide code cells executable as well, e.g., not just plain text.
• Guide 22 already covers TensorBoard setup, training, and viewing logs. This guide should focus narrowly on plot_summary usage and link to guide 22 for the setup, rather than re-explaining how to pass a Path or inference object.

I suggest the following outline / cells:

1. One-liner intro: what plot_summary does, link to guide 22 for training/tracking setup
2. One runnable setup cell: quick NPE training (reuse the pattern from guide 22: prior, simulator, train, get log_dir)
3. Basic usage: single plot_summary call, show output
4. Overlay with customization: one cell combining overlay, colors, labels, title, grid. This is the main feature to showcase
5. Drop the standalone "Using a log directory" section. just mention in the intro that you can pass either an inference object or a Path