Adding Ridgeplot to Arviz-Plots

[imperorrp]

First commit for adding ridge plot. As per issue #7, the already implemented forest plot structure was taken as base logic for this plot and the conventions for adding a new plot (https://arviz-plots.readthedocs.io/en/latest/contributing/new_plot.html) have been followed. This plot plots kde ridges instead of credible intervals and point estimates.

I also added two top-level arguments to this plot since I thought they might be useful-

• rescale_kde (rescales kde by raising the 'y' coord values in the density dataset to a power, set to 1 by default. Using exponent here to amplify the peaks so they're more prominently visible because there's a tendency for kdes to appear very flat when plot_ridge() has a lot of subplots stacked)

• plot_ridge_base (boolean to choose whether or not to plot a baseline for each ridge. I think this helps in legibility if the ridges overlap each other or have high peaks, especially if the kde was rescaled)

Also modified the line_xy visual element function to accept an additional y argument (set to None by default) which is added to da.sel(plot_axis='y') if it is not None, so that each ridge is plotted at the correct height corresponding to its labels on the left column. Since the functionality is the same as previously if y is not provided, this ensures the function works where used already (like in plot_dist()) without deprecation.

---

📚 Documentation preview 📚: https://arviz-plots--57.org.readthedocs.build/en/57/

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 90.68627% with 19 lines in your changes missing coverage. Please review.

Please upload report for BASE (main@6f8b750). Learn more about missing BASE report.
Report is 160 commits behind head on main.

Files with missing lines	Patch %	Lines
src/arviz_plots/plots/ridgeplot.py	90.32%	18 Missing ⚠️
src/arviz_plots/visuals/__init__.py	93.75%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #57   +/-   ##
=======================================
  Coverage        ?   85.75%           
=======================================
  Files           ?       17           
  Lines           ?     1951           
  Branches        ?        0           
=======================================
  Hits            ?     1673           
  Misses          ?      278           
  Partials        ?        0           

☔ 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.

[imperorrp]

Some examples of the plot output with the centered_eight data.

Without the extra top-level args for rescaling/adding a baseline:

azp.plot_ridge(data)

azp.plot_ridge(data, var_names=["mu", "tau"])

---

With rescaling:

azp.plot_ridge(data, rescale_kde=0.5)

azp.plot_ridge(data, var_names=["mu", "tau"], rescale_kde=0.5)

azp.plot_ridge(data, rescale_kde=1.5)

azp.plot_ridge(data, var_names=["mu", "tau"], rescale_kde=1.5)

---

With base line added:

azp.plot_ridge(data, plot_ridge_base=True)

azp.plot_ridge(data, var_names=["mu", "tau"], plot_ridge_base=True)

(Looks a bit cluttered when baselines are added when there's a lot of ridges but for the latter I think helps in seeing the density, especially if some ridges are overlapping.)

[OriolAbril]

There should be two visual elements that generate the ridge: the edge (already implemented and working) and the face (filling the area of the kde). This will match the behaviour of current plot_ridge (via plot_forest), see for example: https://python.arviz.org/en/stable/examples/plot_forest_ridge.html.

As always either or both can be turned off. Regarding defaults, I think having the edge have the default color is great, for the area I think using the same color as the default would mean having an data-ink ratio too low, we could use color="white", alpha=0.4 or something like that to help with visualizing overlap between ridges but keeping a "clean" look (that would look quite weird for backgrounds other than white though like arviz-darkgrid style), or we could use plot_kwargs.setdefault("face", False) to skip it by default but giving users the option to turn it on.

API question (needs discussion before any implementation): I think it would probably be nice to have a shortcut/syntactic sugar to have both edge and face use the same properties, but I am not sure how much worth it it would be.

[OriolAbril]

I think combined=True might be a better default from plot_ridge

[OriolAbril]

I would also remove this one

[OriolAbril]

Given the natural meaning of "ridge", I think it makes sense as the plot name and as the name of the whole axes at the right of the figure, but similarly to plot_forest not so much for a single visual element. Here we could have edge and face as the visual elements.

[OriolAbril]

This type of transformation should not be an option. kdes are already hard enough to interpret (even if they look simple and we believe we can read them with ease that is not the case). I would not add this as this makes it basically impossible to actually read anything from the kde other than the range and the position of the peak.

By rescaling I meant multiplying all kdes by a common scale factor in order to preserve all proportions in the curve but having the tallest one peak at height 1 (or something else, this could be user modified, similarly to the overlap control argument there is in current plot_forest).

Suggested change

	for var in density.data_vars:
	density[var].loc[{"plot_axis": "y"}] = density[var].sel(plot_axis="y") ** rescale_kde
	for var in density.data_vars:
	density.loc[{"plot_axis": "y"}] = density.sel(plot_axis="y") / density.sel(plot_axis="y").max().to_array().max() * ridge_height

where ridge_height regulates the height of the ridge (that is, the height of the tallest peak in the ridge). If ridge_height=1 (or lower) there should be no overlap at all (when combined=True and there aren't multiple models) with higher values there will be some area overlap (>1 but smallish should mean area overlap but no edge overlap generally, >>1 will generally mean both are and edge overlap). Once both area and edge plotting are working, we should try several values like 1, 1.5, 2 and see with the example arviz datasets how each looks.

[OriolAbril]

This should use the same pattern as annotate_label or point_estimate_text. That is

1. Both x and y are arguments of the function, both default to None
2. the actual x and y values are obtained with x, y = _process_da_x_y(da, x, y) (combining values from the function arguments and the different dataarray coordinate values as needed).

[imperorrp]

Renamed the ridge artist to edge and added the face artist. I had to modify the fill_between_y visual element function a bit as well so that the 'y' aesthetic argument passed to it by pc.map is considered and also updated the line_xy visual element with the pattern those other visual element functions use.

Also moved the kde rescaling code to just after statistical computation so that a new dataset face_density can be created and later fed to the fill_between_y function.

About the default color for face- I've just used the same one as for edge for now but with the alpha=0.4 added. Have left the bit about face/edge using the same properties for now as well, for after discussion like you said.

Also added some tests, with style drawn from the existing ones for plot_dist and plot_forest. Tried adding some hypothesis tests but there were issues with that so I haven't committed those changes yet.

Here's how the plots look like now:

When combined==True (by default)

When combined==False

[OriolAbril]

Tried adding some hypothesis tests but there were issues with that so I haven't committed those changes yet.

Try manually generating a plot_ridge without the edge and another without the face. The hypothesis test if you copied it from the plot_forest one is failing at least due to that

[OriolAbril]

I have added a couple comments for the visuals, to make sure we can combine x and y as aesthetics along with the data provided by the plots.

[OriolAbril]

Looking at the code again, I think for maximum flexibility we should modify the _process_da_x_y so the function keeps only the logic of combining the x and y coordinate values in da (if any) with the x and y arguments (if any). Thus all following possibilities are valid:

(da)       # with x, y in plot_axis
(da, x, y) # with x, y in plot axis, returns da.sel(x) + x, da.sel(y) + y
(da, x, y) # without x, y in plot axis, returns x, y
(da, x)    # with x, y in plot axis, returns da.sel(x) + x, da.sel(y)
(da, x)    # without x, y in plot axis, returns x, da
# basically the same changing x for y too

This means that if we (functions in plots module) provide the data via da, both x and y can be used as aesthetics and will work as expected.

Then have a different function _ensure_scalar to make sure things that are supposed to be scalars are actually scalars and not 0-dimensional arrays, dataarrays with a single element or anything of the sort.

Overall:

def _process_da_x_y(da, x, y):
    """Process da, x and y arguments into x and y values."""
    da_has_x = "plot_axis" in da.dims and "x" in da.plot_axis
    da_has_y = "plot_axis" in da.dims and "y" in da.plot_axis
    if da_has_x:
        x = da.sel(plot_axis="x") if x is None else da.sel(plot_axis="x") + x
    if da_has_y:
        y = da.sel(plot_axis="y") if y is None else da.sel(plot_axis="y") + y
    if x is None and y is None:
        raise ValueError("Unable to find values for x and y.")
    if x is None:
        x = da
    elif y is None:
        y = da
    return x, y

def _ensure_scalar(*args):
   return tuple(arg.item() if hasattr(arg, "item") else arg for arg in args)

With that, line_xy should use x, y = _process_da_x_y(da, x, y) directly, and the text related functions that currently use it should use x, y = _ensure_scalar(_process_da_x_y(da, x, y))

[OriolAbril]

Side comment is that if y_bottom is None and y would evaluate y as a boolean, thus, y_bottom=None, y=0 would not enter the if. Not relevant with the other behavioural comment though

[OriolAbril]

I think I would modify this to make sure things behave as expected. x and y are mostly used so we can have aesthetics x and y that control shifting the visual elements within a plot. Therefore, I would expect that if I pass y, everything is shifted along the vertical axes. With these checks, if I give both y_top and y this won't happen.

I would use

if y is not None:
    y_top += y
    y_bottom += y

once the if "kwarg" in da.dims: finishes as then both y_top and y_bottom should be array_like elements and not None.

[imperorrp]

Tried adding some hypothesis tests but there were issues with that so I haven't committed those changes yet.

Try manually generating a plot_ridge without the edge and another without the face. The hypothesis test if you copied it from the plot_forest one is failing at least due to that

Yeah I just checked this- since the current code calls the density calculation function only if edge_kwargs is not False, it doesn't work when edge is turned off with plot_kwargs. That might be the reason behind the test failing. I could move the density computation outside to fix this, with perhaps a check for if edge_kwargs or face_kwargs.

[imperorrp]

Modified density calculation to occur if either one of edge_kwargs or face_kwargs is not False, indicating that it is required to be calculated for plotting to be possible. Also added the visual element function changes and hypothesis tests code, which appears to pass now for plot_ridge

[imperorrp]

Current plot outputs:

With both edge and face:

With edge only:

With face only:

[OriolAbril]

This should be outside the catch_warnings context. The catch_warnings is used to prevent a warning being printed when we have subsets of the data that are all nan (which happens when we have multiple models with different variables/coordinate values). It is not something that means anything, so we don't want users to see warnings that are not actionable. If it appears when there is a single model then it is an actual issue with the model/sampling, so we do want to print it hence why the filterwarnings is only active when multiple models are present.

[OriolAbril]

This should be np.size(point) now

[OriolAbril]

Suggested change

	if x is not None or y is not None: # if either x or y is provided as an arg
	x, y = _process_da_x_y(da, x, y)
	else:
	x = da.sel(plot_axis="x")
	y = da.sel(plot_axis="y")
	x, y = _process_da_x_y(da, x, y)