You can do this by manipulating the
idata.posterior object, which is an xarray Dataset. If you only care about changing the order in which the plots appear, then you can index the data based on school using the Dataset.sel method.
Here’s an example where I select schools in the opposite order (sorry about the missing cells, I’m sure no one wants to see my typos!):
In : import arviz as az
In : inf = az.load_arviz_data("non_centered_eight")
In : inf.posterior
Dimensions: (chain: 4, draw: 500, school: 8)
* chain (chain) int64 0 1 2 3
* draw (draw) int64 0 1 2 3 4 5 6 7 8 ... 492 493 494 495 496 497 498 499
* school (school) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
mu (chain, draw) float64 ...
theta_t (chain, draw, school) float64 ...
tau (chain, draw) float64 ...
theta (chain, draw, school) float64 ...
In : _ = az.plot_trace(inf.posterior, var_names=["theta"])
# will plot image 1 below
In : order = inf.posterior.school.values[::-1]
In : order
array(['Mt. Hermon', "St. Paul's", 'Lawrenceville', 'Hotchkiss',
'Phillips Exeter', 'Phillips Andover', 'Deerfield', 'Choate'],
In : _ = az.plot_trace(inf.posterior.sel(school=order), var_names=["theta"])
# will plot image 2 below
Cell  plots this:
Cell  plots this:
In general, arviz integrates very nicely with xarray—both are, in my experience, very well-designed and quite flexible, so hats off to the developers of both If you want to plot subsets of your posterior, to rename the labels, reorder things, etc., you can usually do it by manipulating the underlying xarray Dataset and passing that straight to arviz.
With that said, this was challenging to do at first—mostly because learning the language of xarray Datasets and the new abstractions they use took some trial and error. I think this might be a useful example to add to the arviz docs to highlight how you can do some operations in xarray to make your plotting life easier.