Arviz Charts are ... Different?

Questions
1

I have an old blogpost. It features a pretty plot.

Here’s the code of the model:

with pm.Model() as mod: 
    intercept = pm.Normal("intercept", 0, 2)
    time_effect = pm.Normal("time_weight_effect", 0, 2, shape=(4,))
    diet = pm.Categorical("diet", p=[0.25, 0.25, 0.25, 0.25], shape=(4,),
                          observed=dummy_rows)
    sigma = pm.HalfNormal("sigma", 2)
    sigma_time_effect = pm.HalfNormal("time_sigma_effect", 2, shape=(4,))
    weight = pm.Normal("weight", 
                       mu=intercept + time_effect.dot(diet.T)*df.time, 
                       sd=sigma + sigma_time_effect.dot(diet.T)*df.time, 
                       observed=df.weight)
    trace = pm.sample(5000, chains=1)

Here’s the traceplot:

image
image856×568 80.4 KB

This chart is very pretty. In fact … it’s beautiful! It detects the different slopes and has a different color per diet. It makes it very easy to compare! I recall making this with the builting plot_traceback from PyMC3.

The world has moved on and plotting is now done via Arviz. This is all well and good but when I try to make a traceplot it looks different.

image
image872×1448 233 KB

It’s no longer doing the pretty grouping. Does anybody know if that is still possible?

1 Like
2

Pass the flag compact=True in the call to az.plot_trace! See the docs here.

5 Likes
3

To add to the previous answer, in addition to compact you can also use combined to get all the chains combined to compute the KDE (the default also changed from a single KDE to a KDE per chain) and chain_prop={"ls": "-"} there are examples of this in one of my blogposts.

You may also notice that in the previous version the right column plot also has only one line, result that can’t be reproduced anymore. This was a bug in the previous version when combining chains. In the right column plots, only the first chain was taken into account.

1 Like