Plot a calibration plot with error bars. Optionally overlay v in each bin.

calibration_plot.py

from scipy import stats
import numpy as np
import matplotlib as plt

def beta_errors(num, denom):
    return stats.beta.interval(.95, num+1, denom-num+1)


def calibration_curve_error_bars(a, p, n_bins=10):
    pmin, pmax = p.min(), p.max()
    pmin -= 1e-10
    binstarts = np.linspace(pmin, pmax, n_bins+1)
    bincentres = binstarts[:-1] + (binstarts[1] - binstarts[0])/2.0
    numerators = np.zeros(n_bins)
    denomonators = np.zeros(n_bins)
    for b in range(n_bins):
        idx_bin = (p > binstarts[b]) & (p <= binstarts[b+1])
        denomonators[b] = idx_bin.sum()
        numerators[b] = a[idx_bin].sum()

    errors = beta_errors(numerators, denomonators)
    return bincentres, numerators, denomonators, errors


def plot_calibration_curve_error_bars(a, p, n_bins=10, ax=None, alpha=1.0, label='',
    add_n=False):
    x, n, d, err = calibration_curve_error_bars(a, p, n_bins)
    if ax is None:
        fig, ax = plt.subplots(1, 1)
    
    nan_idx = (d != 0)
    x_nanan = x[nan_idx]
    n_nanan = n[nan_idx]
    d_nanan = d[nan_idx]
    err_nanan = err[0][nan_idx], err[1][nan_idx]
    ax.errorbar(x_nanan, n_nanan/d_nanan, yerr=[n_nanan/d_nanan-err_nanan[0], err_nanan[1]-n_nanan/d_nanan],alpha=alpha,label=label)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_ylabel("Fraction of positives",color='blue')
    ax.set_xlabel("Mean predicted value")
    if add_n:
        ax2 = ax.twinx()
        ax2.step(x_nanan,d_nanan,color='green',alpha=0.5)
        ax2.set_ylabel('N', color='green')
        ax2.tick_params(axis='y', labelcolor='green')
    ax.plot([0, 1], [0, 1], "k:")