Simulated comparison of plugin and debased estimators for L1 calibration error.

ece-estimators-sim.py

"""
Simulated comparison of estimators for L1 calibration error:

1. Plugin estimator, i.e. ECE [Guo et al. 2017]
2. De-biased estimator [Kumar et al. 2019]

In each simulation, we draw y_per_bin, p_per_bin, and fix n_per_bin to a constant.
We draw samples y_hat_per_bin and compare MSE of the two estimators as a function of bin size.
"""
import numpy as np
import pandas as pd
import seaborn as sns
from collections import defaultdict
from tqdm import tqdm
from matplotlib import pyplot as plt
from dfply import gather


def run_simulation(bin_size=1000, n_bins=15, n_mc_samples=1000, verbose=False):

    # fix n_per_bin and p_per_bin; treat y_per_bin as random variable to be estimated
    n_per_bin = np.full(n_bins, bin_size)
    y_per_bin = np.random.rand(n_bins)
    p_per_bin = y_per_bin + 0.05 * np.random.randn(n_bins)
    y_hat_per_bin = np.zeros(n_bins)

    # draw samples for observed y_hat_per_bin ~ binomial
    for i in range(n_bins):
        y_hat_per_bin[i] = (np.random.rand(n_per_bin[i]) < y_per_bin[i]).mean()

    true_ece = np.average(np.abs(y_per_bin - p_per_bin), weights=n_per_bin)
    plugin_ece = np.average(np.abs(y_hat_per_bin - p_per_bin), weights=n_per_bin)

    # monte carlo samples for de-biased estimator
    mc_samples = np.random.randn(n_mc_samples, n_bins)
    mc_samples *= (y_hat_per_bin * (1 - y_hat_per_bin) / n_per_bin) ** 0.5
    mc_samples += y_hat_per_bin
    mc_ece = np.mean(np.average(np.abs(p_per_bin - mc_samples), axis=1, weights=n_per_bin))

    debiased_ece = plugin_ece - (mc_ece - plugin_ece)
    proposed_ece = plugin_ece - mc_ece

    if verbose:
        print(f"True ECE: {true_ece:.4f}")
        print(f"Plugin ECE: {plugin_ece:.4f}")
        print(f"Monte Carlo ECE: {mc_ece:.4f}")
        print(f"Debiased ECE: {debiased_ece:.4f}")

    return {
        "plugin": (plugin_ece - true_ece) ** 2,
        "debiased": (debiased_ece - true_ece) ** 2,
    }

if __name__ == "__main__":

    df = defaultdict(list)

    for bin_size in (50, 100, 200, 500, 1000, 1500, 2000):
        for _ in tqdm(range(1000)):
            for k, v in run_simulation(bin_size).items():
                df[k].append(v)
            df["bin_size"].append(bin_size)

    df = pd.DataFrame(df)
    df = df >> gather("estimator", "mse", ["plugin", "debiased"])

    plt.figure(figsize=(8, 3))
    sns.set_style("white")
    sns.lineplot(data=df, x="bin_size", y="mse", hue="estimator",
                 palette=sns.color_palette("gray", 2))
    plt.tight_layout()
    plt.legend()
    plt.show()