Fix: harmonise visualize style -- axes, returns, attributes by Doresic · Pull Request #1709 · ICB-DCM/pyPESTO

Doresic · 2026-05-04T16:10:45Z

This PR lays the foundation for a coherent pyPESTO visualization API. It primarily normalizes the contracts that callers depend on, so that downstream style work can be applied uniformly, with a few small review-driven follow-up changes to helper placement, layout handling, and observable-mapping axes handling.

What changed

Shared helpers in pypesto/visualize/misc.py
Shared helpers used by all plotters:

get_ax(ax, size) — returns the provided axes or creates a new figure/axes pair
get_axes_array(axes, nrows, ncols, size) — returns the provided 2-D array or creates a new subplot grid; validates shape if an existing array is passed
plot_diagonal_marginal(ax, values, diag_kind) — shared KDE/histogram for scatter-matrix diagonals
process_deprecated_kwarg(...) — generic rename shim with DeprecationWarning
make_grid_shape(n_panels) — shared near-square subplot-grid helper

Type hints (all visualize files)

tuple → tuple[float, float] for size parameters
ax: Axes = None → ax: Axes | None = None throughout
Added -> matplotlib.axes.Axes and -> np.ndarray return annotations where missing

Single-panel ax contract
All single-panel public functions now accept ax: matplotlib.axes.Axes | None and return matplotlib.axes.Axes.

Multi-panel axes grid contract
All multi-panel public functions now accept axes: np.ndarray | None and return a 2-D np.ndarray of Axes. optimization_scatter rewritten from seaborn PairGrid to pure matplotlib: per-panel KDE diagonals, fval-encoded viridis_r colorbar, correct lb_full/ub_full[parameter_indices[i]] bound indexing.

sampling_parameter_cis: alpha → confidence_levels

Old: alpha: Sequence[int] (integer percentages, e.g. [95, 68])
New: confidence_levels: Sequence[float] (fractions, e.g. [0.95, 0.68])
alpha kept as a deprecated kwarg; auto-converts with DeprecationWarning

profiles: confidence_level convenience parameter
New optional confidence_level: float | None = None alongside ratio_min; converts via chi2_quantile_to_ratio.

Layout handling
Where figures are created internally, they now use constrained layout instead of forcing tight_layout() afterward.

Tests

helper tests moved from test/visualize/test_style.py to test/visualize/test_misc.py
close_fig moved to test/conftest.py so it can be reused across visualization tests

What this PR does NOT touch

Default figure sizes, font sizes, or any rcParams
Spine removal, colorbar style unification (planned for follow-up)
model_fit.py, ordinal_categories.py internals

- Add pypesto/visualize/_style.py with shared helpers get_ax and process_deprecated_kwarg. The latter is the canonical pattern for the visualization signature harmonization that follows. - Drop the module-level cmap global in sampling.py; inline at the single use site. - Rename par_indices → parameter_indices in sampling_parameter_traces, sampling_1d_marginals, and the internal get_data_to_plot. Public functions accept par_indices as a deprecated alias (DeprecationWarning).

- ax: canonicalize to `matplotlib.axes.Axes | None = None` everywhere (was `plt.Axes | None`, `matplotlib.axes.Axes` without `| None`, or untyped `ax=None` in profiles.py) - size: fix `tuple[float]` → `tuple[float, float]` in optimizer_convergence.py and waterfall_lowlevel; bare `tuple` → `tuple[float, float]` in optimizer_history_lowlevel - title/suptitle: `str = None` → `str | None = None` in sampling functions - Add `import matplotlib.axes` to waterfall.py, optimizer_convergence.py, optimizer_history.py, profiles.py

- switch remaining single-panel plotters and lowlevel helpers to the shared `get_ax` helper - remove ad hoc pyplot-based axes creation in favor of the shared path - ensure plots still fully configure passed-in axes, including labels and other axis-level metadata - make `projection_scatter_umap_original` explicitly accept and return an axes object - preserve backward compatibility for `optimization_run_properties_one_plot` by keeping positional argument order intact - keep tests broad and figure-focused, with targeted checks only where behavior actually changed

- `get_axes_array(axes, nrows, ncols, size)` — create or normalize a 2-D axes grid; raises on shape mismatch - `plot_diagonal_marginal(ax, values, diag_kind)` — KDE/hist marginal extracted from `optimization_scatter` so `sampling_scatter` can share it Functions updated: `optimization_scatter`, `sampling_scatter`, `sampling_1d_marginals`, `sampling_prediction_trajectories`, `visualize_estimated_observable_mapping`, `plot_linear_observable_mappings_from_pypesto_result`, `plot_splines_from_pypesto_result`, `plot_splines_from_inner_result`, `ensemble_scatter_lowlevel`, `projection_scatter_umap_original`, `optimization_run_properties_one_plot`. Also removes all `sns.set(style="ticks")` calls from `sampling.py` — these were the root cause of global matplotlib style mutations that made plots look inconsistent depending on call order.

`sampling_parameter_cis` accepted `alpha: Sequence[int]` (e.g. `[95]`) as a credibility level in integer percentages. This was mis-named (not matplotlib transparency) and used inconsistent units compared to every other CI parameter in the codebase (all others use 0–1 fractions).

`profiles` accepted only `ratio_min` (a likelihood-ratio threshold) to filter profile points. Users wanting to cut off at a standard confidence level had to call `chi2_quantile_to_ratio` manually first. New additive kwarg `confidence_level: float | None = None` — pass e.g. `confidence_level=0.95` and it is converted to the equivalent `ratio_min` via `chi2_quantile_to_ratio`. The two parameters are mutually exclusive; passing both raises `ValueError`. `ratio_min` is unchanged and keeps its existing meaning and default.

codecov-commenter · 2026-05-04T16:15:21Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 83.65123% with 60 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.29%. Comparing base (37e21df) to head (5eb1381).

Files with missing lines	Patch %	Lines
pypesto/visualize/observable_mapping.py	60.00%	22 Missing ⚠️
pypesto/visualize/parameters.py	84.21%	12 Missing ⚠️
pypesto/visualize/misc.py	87.17%	10 Missing ⚠️
pypesto/visualize/sampling.py	90.00%	8 Missing ⚠️
pypesto/visualize/dimension_reduction.py	85.18%	4 Missing ⚠️
pypesto/visualize/profiles.py	72.72%	3 Missing ⚠️
pypesto/visualize/optimization_stats.py	95.00%	1 Missing ⚠️
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #1709      +/-   ##
===========================================
+ Coverage    83.99%   84.29%   +0.30%     
===========================================
  Files          164      164              
  Lines        14492    14636     +144     
===========================================
+ Hits         12172    12337     +165     
+ Misses        2320     2299      -21

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

vwiela · 2026-05-06T07:40:19Z

-    size: tuple[float] | None = (12, 6),
-):
+    size: tuple[float, float] | None = (12, 6),
+) -> plt.Axes:


Should this rather be matplotlib.axes.Axes as in the other files?

Yepp will change

vwiela · 2026-05-06T07:40:29Z

-    size: tuple[float] | None = (16, 10),
-):
+    size: tuple[float, float] | None = (16, 10),
+) -> plt.Axes:


Yepp will change

vwiela · 2026-05-06T07:50:29Z

-        list(cmap(v))[:LEN_RGB]
+        list(matplotlib.cm.viridis(v))[:LEN_RGB]
        for v in np.linspace(cmap_min, cmap_max, n_variables)
    ]


This whole function is pretty tailored to this one specific colormap.
Since we are touching visualizations anyways, we could think about making this more verstaile. So a user could use different colormaps and cut-offs (0.3, 0.85) and we leave the current one as defaults.

Good idea
I agree
Paul suggested below some similar additions of allowing the user pass settings to the functions.
I would tackle all of these together in the next PR that'll focus on style of all plots + the possibility of used adding style changes to the functions or even whole module. There could be a very cool module-wide-changes thing one can do and define.
A simple thing could be just allowing the user to pass a dictionary of kwargs that will get distributed acorss all plotting functions to change sizes shapes etc. -- up to some limit, adding all would probably be too much.
And then having a specification that's general across all visualisation functions in _style.py of general sizes/font sizes etc that will be used by default and imported from the _style as constants.

vwiela · 2026-05-06T07:52:35Z

+    ax: plt.Axes | None = None,
    **kwargs,
-):
+) -> plt.Axes:


Shouldn't this be matplotlib.axes.Axes?

Either is fine.

I'll put it to matplotlib.axes.Axes just for consistency.

PaulJonasJost

Some great ideas among this. Not entirely sure what the big picture is, you want to move towards with this initial PR as the AI summary hints here. Would perhaps be important to discuss this once in the near future.

PaulJonasJost · 2026-05-06T09:46:45Z

+Cross-cutting pieces used by multiple plotters live here so each plotter
+does not reinvent them.


Suggested change

Cross-cutting pieces used by multiple plotters live here so each plotter

does not reinvent them.

Superfluous

PaulJonasJost · 2026-05-06T09:57:55Z

    legends: str | list[str] | None = None,
    plot_type: str = "line",
-) -> matplotlib.axes.Axes:
+) -> np.ndarray:


Yes it is an array, yes it says that the array consist of axes, yet i find the previous matplotlib.axes.Axematplotlib.axes.Axes still more informative

Hmmm yeah I see the point. matplotlib.axes.Axes is more informative, but in this case it was wrong to use it.
The code defines axes, a np.ndarray of axes. In both cases.

if plot_type!="both" it creates a (1,1) array of Axes objects

if plot_type=="both" it creates a (1,2) array of Axes objects
so with the current code of the visualisation, matplotlib.axes.Axes would always have been wrong.

What we can do is make it return really just a matplotlib.axes.Axes object in the first case, and then say the return type is either this or that depending on the settings. But I would vote for current situation instead, to make each function always have one type of output instead of the output type changing depending on input parameters.

I wouldn't consider an incorrect type more informative 🤔

PaulJonasJost · 2026-05-06T10:05:28Z

+    n = len(parameter_indices)
+    x_labels = [result.problem.x_names[i] for i in parameter_indices]

-    ax = sns.pairplot(


If I see this correctly, we have changed something like 10 lines of code for 120, for an improvement that to me is not immediately visible.

If the goal is to move to purely matplotlib, that might be a reason, but I don't see why we would push for that necessarily.

Summary of this specific change (my opinion, to be discussed): This does not improve the code, rather the opposite.

Yeah, true
Initial reason for changing this code was to go away from the different return object (in the idea of harmonisation) -- a Seaborn grid object -- compared to the object we return in most cases -- either Axes or a np.ndarray of them.
Furthermore, a second PR later on would focus on harmonisation of general looks of the figures across the module. In the sense that the examples we have in docs have figures with similar fonts, color schemes, and general looks. So moving only to matplotlib would make that simpler. For instance, it's nice to have all plots of this type, a scatter plot with diagonal distributions, have a colorbar on the side to map objective function value to the scatter plot dots. This is a bit tough or would be very different from other ones if we stuck with seaborn.

So I would vote for this. I do see it's more lines just since matplotlib doesn't have an integrated visualisation of this type of a plot, but this allows us a nicer freedom of playing with it to make "our style" across the module.

PaulJonasJost · 2026-05-06T10:06:30Z

@@ -408,10 +424,10 @@ def profile_lowlevel(
    ax: plt.Axes | None = None,


No changes here to the Axes? If we make a point of harmonising it, please check all functions.

ty, missed some

PaulJonasJost · 2026-05-06T10:13:06Z

+    return axes_array
+
+
+def plot_diagonal_marginal(


This is not really a style thing right? In general, why isn't this whole file put into misc.py or in a utility.py which we rather use?

True, was thinking about that
feels nicer in misc.py

Yeah I think all of these functions fit better to misc.py
There will be a purpose to the style file but it's not there yet -- next PR stuff that actually deal with style.
In the sense that it will contain constants of sizes, fonts etc that would be used as style defaults across the vis. module.
Makes more sense yup

PaulJonasJost · 2026-05-06T10:25:12Z

+            color="C2",
+            linewidth=1.8,
            label="Spline function",
+            zorder=2,


I again feel like with more hardcoded style parameters, we move away from a flexible visualisation approach.

similar to above

PaulJonasJost · 2026-05-06T10:29:09Z

        stepsize=stepsize,
        full_trace=full_trace,
-        par_indices=par_indices,
+        parameter_indices=parameter_indices,


Needs a depracation hint as in lines 1359

PaulJonasJost · 2026-05-06T10:30:05Z

-    ax = sns.pairplot(
-        params_fval.drop(["logPosterior", "iteration"], axis=1),


Same argument as in the other scatterplot. Is this change useful?!

I think so yeah
Reasoning above

PaulJonasJost · 2026-05-06T10:32:32Z

+    if alpha is not None:
+        if confidence_levels is not None:
+            raise ValueError(
+                "Pass either `confidence_levels` or the deprecated `alpha`, not both."
+            )
+        import warnings
+
+        warnings.warn(
+            "`alpha` is deprecated; use `confidence_levels` instead. "
+            "Note: units have changed — pass fractions in (0, 1) "
+            "(e.g. `confidence_levels=[0.95]`) instead of integer percentages "
+            "(e.g. `alpha=[95]`). Your values have been divided by 100 automatically.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        confidence_levels = [a / 100 for a in alpha]
+
+    if confidence_levels is None:
+        confidence_levels = [0.95]


This seems to be a target that should also be covered by process_deprecated_kwarg right?

Yeah, but this one is a bit different. The kwarg deprecation is only changing the naming of some argument.
This is much more specific deprecation -- the value users provide has to be changed as well. So I made a custom one for here to point this out to not cause very plotting outcomes for people without warning.

PaulJonasJost · 2026-05-06T10:35:54Z

+    ax.set_ylabel("Count")
+
+
+def process_deprecated_kwarg(


I feel this is an amazing idea. So amazing, this should not be confined to visualization alone and also perhaps not to specific instances, rather have this as some kind of decorator on top of a function. Instead of passing the value we can have some kind of transformation function (either identity or in this case for example y = x/100), this would be great for maintaining this and can be easily removed at the next major release as discussed in our versioning policy.

Yeah, good idea
For now I made it only in this module, so I kept it at this point.
But definitely a good idea for a separate whole-package-wide PR introducing this concept for further use
Apart from a transformation function, then also a custom message should also be possible to add. Cause without it, the warning would be too vague in case of an added transformation. Or maybe just messages depending on the transformation.
Definitely a good addition for whole pyPESTO

dweindl · 2026-05-06T12:07:37Z

+    value:
+        The resolved value, or ``None`` if neither was given.
+    """
+    if deprecated_value is None:


This will miss f(deprecated_kwarg=None), right?

Hmmm that's true
I didn't think about this, in case a user wanted to pass None, this would not catch it.
Currently I don't think I have examples of where I use this deprecation in which a user would ever want to pass None. But there might be

Should we then in this case

have a specific _UNSET value that tracks whether something is set or not for variables that are set or not that would then be different from None to catch None's too?

just have a dosctring acknowledging this limitation for future usage of this function (mostly for us) to make sure the case where we use it is not one of these cases?

I'd go for the first option, which is probably even faster...

dweindl · 2026-05-06T12:08:20Z

Doc incompatible with return type annotation.

dweindl · 2026-05-06T12:09:59Z

+    ax: plt.Axes | None = None,
    **kwargs,
-):
+) -> plt.Axes:


Either is fine.

dweindl · 2026-05-06T12:10:27Z

Doc incompatible with return type annotation.

dweindl · 2026-05-06T12:16:25Z

    pass


+def _observable_mapping_grid_shape(n_panels: int) -> tuple[int, int]:


This isn't really related to observable mappings. Rename and move to general plotting utils?

Yep true, makes sense, we use this at many places. Will add to misc and use it anywhere we want to define a grid with a certain number of plots

dweindl · 2026-05-06T12:23:55Z

    legends: str | list[str] | None = None,
    plot_type: str = "line",
-) -> matplotlib.axes.Axes:
+) -> np.ndarray:


I wouldn't consider an incorrect type more informative 🤔

dweindl · 2026-05-06T12:38:44Z

@@ -1178,18 +1210,19 @@ def sampling_parameter_traces(
    fig.tight_layout()


If a user is allowed to pass fig/ax, we should remove all calls to tight_layout(). The user may have chosen a different layout engine which would be replaced here - that's annoying.

very true, didn't have this in mind
hmmm should this be "no tight layout in case an ax is passed" or in general?

No experience with this specific instance, but overall, I think, tight_layout will disappear.
I'd suggest using plt.subplots(..., layout="constrained") in get_ax et al., and just removing all tight_layout() calls.

dweindl · 2026-05-06T12:43:11Z

+    custom = get_ax(size=(4.0, 3.0))
+    assert tuple(custom.get_figure().get_size_inches()) == (4.0, 3.0)
+
+    plt.close("all")


There's close_fig to be used as decorator in test_visualizy.py. Move that to conftest.py so it can also be used here. Replace all direct plt.close("all") calls to ensure figures are closed also in case of test failure.

- move shared visualize helpers from _style.py into misc.py - add hide_unused_axes and reuse it across plotting functions that trim unused subplot panels - switch internal figure creation to constrained layout and stop forcing tight_layout afterwards - keep deprecated-kwarg handling with _UNSET-based detection of omitted vs explicit None - update observable mapping axes preparation to use the shared helpers - move helper tests from test_style.py to test_misc.py - move close_fig to test/conftest.py and ensure figure cleanup also happens on test failure

Doresic · 2026-05-06T15:21:31Z

Thank you a lot for all the comments!
They should be addressed now
Feel free to take another look @dweindl @PaulJonasJost @vwiela

dweindl · 2026-05-07T10:51:25Z

They should be addressed now

Didn't check everything in all details, but my points have been addressed, thanks.

Doresic added 7 commits April 29, 2026 17:25

add return types, remove deprecation test

bf37126

Small update of type

dfe3f49

Doresic marked this pull request as ready for review May 5, 2026 12:38

Doresic requested review from PaulJonasJost, stephanmg and vwiela as code owners May 5, 2026 12:38

Doresic added 3 commits May 5, 2026 16:49

merge dev

ac85a82

small fix for docs test

d5be0f3

Merge branch 'develop' into fix-harmonise-visualize-style

8acc7b9

vwiela reviewed May 6, 2026

View reviewed changes

vwiela requested changes May 6, 2026

View reviewed changes

PaulJonasJost requested changes May 6, 2026

View reviewed changes

PaulJonasJost requested a review from dweindl May 6, 2026 10:39

dweindl reviewed May 6, 2026

View reviewed changes

small additional change

5eb1381

		Cross-cutting pieces used by multiple plotters live here so each plotter
		does not reinvent them.

		@@ -408,10 +424,10 @@ def profile_lowlevel(
		ax: plt.Axes \| None = None,

		ax = sns.pairplot(
		params_fval.drop(["logPosterior", "iteration"], axis=1),

		pass


		def _observable_mapping_grid_shape(n_panels: int) -> tuple[int, int]:

		@@ -1178,18 +1210,19 @@ def sampling_parameter_traces(
		fig.tight_layout()

Conversation

Doresic commented May 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changed

What this PR does NOT touch

Uh oh!

codecov-commenter commented May 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PaulJonasJost left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Doresic commented May 4, 2026 •

edited

Loading

codecov-commenter commented May 4, 2026 •

edited

Loading

dweindl May 6, 2026 •

edited

Loading

Doresic commented May 6, 2026 •

edited

Loading