Change the names of variables in Arviz Plot

Hello all,

I’m trying to find a work around for how to label the y axis labels in my arviz.plot_forest. The problem seems to relate to the code here:

az.plot_forest(
    idata_list, 
    var_names="beta_continuous", 
    figsize=(6, 32), 
    r_hat=False, 
    combined=True, 
    textsize=8, 
    hdi_prob=0.95, 
    quartiles=False
);

where beta_continuous is defined here:

# Continuous feature coefficients
            beta_continuous_offset = pm.Normal(
                "beta_continuous_offset", 0.0, 1.0, dims="continuous"
            )
            beta_continuous = pm.Deterministic(
                "beta_continuous",
                mu_beta_continuous + sigma_beta_continuous * beta_continuous_offset,
            )

The beta_continuous variable isn’t inheriting the labels from the beta_continuous_offset term so I’m getting a plot that looks like this:

If anybody knows how I can swap the [0], [1], [2],… with labels of my own choosing, I would really appreciate it. Thank you!

1 Like

Dimensions don’t propagate in PyMC (and inputs don’t align and broadcast either), that is something that happens only with the outputs when using InferenceData thanks to xarray.

Therefore, you need to annotate the dimension in beta_continuous too:

beta_continuous = pm.Deterministic(
    "beta_continuous",
    mu_beta_continuous + sigma_beta_continuous * beta_continuous_offset,
    dims="continuous"
)

This will make beta_continuous in the output to also have the same coordinate values and dimensions, and (if using a recent enough version of arviz) the plots will also use that information.

For further customization of the plotting labels you can read Label guide — ArviZ dev documentation

1 Like