arviz-devs / arviz-plots

ArviZ modular plotting
https://arviz-plots.readthedocs.io
Apache License 2.0
2 stars 1 forks source link

Add Rootograms #52

Open aloctavodia opened 3 months ago

aloctavodia commented 3 months ago

Rootograms are a visual model checking tool for count models. You can read more here and https://fromthebottomoftheheap.net/2016/06/07/rootograms/ here https://arxiv.org/abs/1605.01311

output

The black dots (and dashed lines) represent the observations and the bars are the predictions from the model. The predictions "hang" from the observations. A perfect model will be one where the hanging bars touch the dashed grey lines at zero. Instead of the frequency, we plot its square root. This makes it easier to compare observed and expected frequencies even for low frequencies.

The previous figure is from an example in BAP3. We should create a more Bayesian version and represent the uncertainty for the bars.

def rotogram(idata, ax):
    max_ = 17
    bins = np.array(range(0, max_))
    dims = idata.posterior_predictive.dims
    n_samples = dims["chain"] * dims["draw"]
    pred =  (np.histogram(idata.posterior_predictive["satellite"].values.ravel(),  bins=bins)[0] / n_samples)**0.5
    observed = np.histogram(crab["satellite"].values, bins=bins)[0]**0.5

    ax.bar(bins[:-1], observed, 0.5, bottom=pred-observed, color="C2")
    ax.plot(bins[:-1], pred,  "k.--")
    ax.hlines(0, -0.5, max_-1.5, linestyle="--", color="C1")

fig, axes = plt.subplots(2, 2, sharey=True, figsize=(12, 6))
for ax, idata, model in zip(axes.ravel(), 
                            [idata_crab_p, idata_crab_hp, idata_crab_nb, idata_crab_hnb],
                            [model_crab_p, model_crab_hp, model_crab_nb, model_crab_hnb]):
    rotogram(idata, ax)
    ax.set_title(model.family.name)

axes[1,1].set_xlabel("Satellites")
axes[1,0].set_xlabel("Satellites")
fig.text(-0.03, 0.5, "$\sqrt{Frequency}$", va="center", rotation="vertical", fontsize=14)