Ordering of variables after coords set

Short Description

I’m trying to use the coords/dims to set labels on posterior variables from a stan fit and the order of the dimensions gets shuffled when I set the labels (not sure what determines the order). Is there a way to (1) retain the existing order (A,B,C…) and (2) sort the order according to a custom index that I provide (e.g. pass a index array like [2,1,0,3, …] to have the plot appear in the order [C, B, A, D …])?

Sorry if this is covered in the documentation somewhere, I didn’t see it.

Without setting any coords/dims, the variables appear in the same order as one would expect from the stan code.

Setting the coords/dims with the code below changes the order to this.

Code Example or link

azdata = az.from_pystan(posterior=fit, observed_data=['Y'], coords={'betas': ['A', 'B', 'C', 'D', 'E', 'F', 'G']}, dims={'b': ['betas']})
	
axes = az.plot_forest(azdata, var_names=('b'))

Using arviz 0.3.2

[ahartikainen]

fixed by #557

There was a set making things out-of-order

ps. I recommend az.to_netcdf and az.from_netcdf for file sharing.

[ColCarroll]

Thanks for the prompt reply @daeh! (also, kudos for the sweet pathlib use!)

This can be recreated with the eight schools model in pymc3. I should be able to dig deeper tomorrow morning. Definitely a bug, but thank goodness that the dimensions map correctly (0 -> ‘a’, 1 -> ‘b’, etc)!

import arviz as az
import numpy as np
import pymc3 as pm

data = {
    "J": 8,
    "y": np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0]),
    "sigma": np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0]),
}

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sd=5)
    tau = pm.HalfCauchy("tau", beta=5)
    eta = pm.Normal("eta", mu=0, sd=1, shape=data["J"])
    theta = pm.Deterministic("theta", mu + tau * eta)
    pm.Normal("obs", mu=theta, sd=data["sigma"], observed=data["y"])
    trace = pm.sample()

az.plot_forest(trace, var_names='theta')

data = az.from_pymc3(trace=trace, coords={'schools': list('abcdefgh')}, dims={'theta': ['schools']})
az.plot_forest(data, var_names='theta')