Rolling Regression in PyMC

All the arviz plot functions return a list of matplotlib axis objects, so if you save it you can continue plotting on them as usual. I get the y-tick locations using yaxis.get_major_ticks().get_loc(), which combined with the parameter value itself lets you do a scatterplot:

ax = az.plot_forest(idata, var_names=['alpha', 'beta'], coords={'regression':np.arange(25)},
               combined=True)
y_vals = reversed(list(map(lambda i: ax[0].yaxis.get_major_ticks()[i].get_loc(), np.arange(50))))
for i, y in enumerate(y_vals):
    data = params.dropna().iloc[i % 25, i // 25]
    ax[0].scatter(data, y, marker='^', zorder=100, color='tab:red')
1 Like