Incorporate categorical labels in plots/summary

I’m using arviz in conjunction with pymc3 fitting a model with categorical variables. When using arviz to view/plot variables the names associated with the categories are replaced by their integer values, making information harder to interpret. It would be nice if there were a way to restore meaningful category names, perhaps during creation of the InferenceData object?

For example, plot_forest:

data_path = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/NWOGrants.csv"
nwo = pd.read_csv(data_path, delimiter=';')

gender = pd.Categorical(nwo['gender'])  # 0 is female, 1 is male

with pm.Model() as m_1:
    a = pm.Normal('a', 0, 1.5, shape=gender.categories.size)
    
    p = pm.invlogit(a[gender.codes])
    # Category labels are removed by the use of `.codes`, which provides integer labels.
    # This is necessary because theano throws a RecursionError if the categories are not integers.
    # This is the information that it would be nice to re-incorporate in arviz. 
    
    award = pm.Binomial('award', nwo.applications, p, observed=nwo.awards)
    
    trace_1 = pm.sample()

az.plot_forest(trace_1, combined=True)

It would be nice if there a way to re-incorporate the category labels lost through having to provide the categories as integers (gender.codes)

[oscarbranson]

Great, Thanks @OriolAbril. Think this can probably be closed now?

[OriolAbril]

I have been doing some tweaks to the docs, from_pymc3 will have a link to the cookbook and the cookbook will have a table of contents to ease navigation. Links point to the updated docs (hosted in my fork until merged in #1184)