Plot multiple coordinates of a posterior variable's dimension on the same axis

Tell us about it

I think it should be possible with Arviz to plot multiple coordinates of a posterior variable on the same graph, for simple comparison. I have not seen an example that achieves this.

Basically, I would like to create select plots from az.plot_trace() for select variables/coords/dims, without the actual traces, etc:

import arviz as az
import matplotlib.pyplot as plt

data = az.load_arviz_data('non_centered_eight')

coords = {'school': ['Choate', 'Lawrenceville']}

az.plot_trace(data,
              var_names=('theta'),
              coords=coords,
              combined=True,
              divergences=False,
             )

To create the same plot myself of theta I can run:

az.plot_dist(data.posterior.sel(school='Choate')["theta"], color='C0', label='school: Choate')
az.plot_dist(data.posterior.sel(school='Lawrenceville')["theta"], color='C1', label='school: Lawrenceville')
plt.ylabel('Density')
plt.title('theta')
plt.legend(frameon=False)

I imagine a function that takes a list of variables (one variable per axis), and the coords to plot separately per variable.

az.plot_posterior_comparison(data,var_names=["theta"], coords= {'school': ['Choate', 'Lawrenceville']})
az.plot_posterior_comparison(data,var_names=["theta", "mu"], coords= {'school': ['Choate', 'Lawrenceville']})

I’m not looking for all the stats found on plot_posterior() or the ridge plot aspect of plot_forest. I’m aware of plot_dist_comparison but I think that is for comparing the prior and posterior. The name of that function might have made sense for this use.

❗ As an aside, I was surprised that dist_plot is set to color=C0 by default instead of using the matplotlib color cycle.

Thoughts on implementation

I’m not sure if this should be a standalone function, or an expanded version of dist_plot.

[aloctavodia]

The closest plot we have is https://arviz-devs.github.io/arviz/api/generated/arviz.plot_density.html

[OriolAbril]

We could have a compact argument but we’d have to be careful and see how it can be integrated with #1676 so that it handles multidimensional case gracefully.

i.e. for a 3d (in addition to chain and draw dims) do we want to compact like we do in plot_trace? or, given subject, treatment, repetition do we maybe want to “skip” repetition and “compact” subject so we see a subplot per treatment, each of them having n kdes, one per subject so we can easily compare how each subject is affected by the same treatment?