%load_ext autoreload
%autoreload 2

P-value Function Diagnostics

INTRO & SETTINGS

The p-values calibration method fits a monotonic probabilistic classifier that estimates the rejection probability \(P(T \le \tau \mid \theta)\) – i.e. a p-value/CDF function of the test statistic, conditional on \(\theta\). Notebooks 2 and 4 use this fitted function to build confidence sets, but never ask how good the fit itself is.

This notebook is about that second question: diagnosing the fit of the p-value function, independent of the confidence sets built from it. We reuse the same model as notebook 2 (Posterior + calibration_method='p-values' on GaussianMean) so the diagnostics here are directly comparable to that notebook’s construction.

# SETTINGS

LIKELIHOOD_COV = 0.01
PRIOR_LOC = 0
PRIOR_COV = 0.1

PARAM_DIM = 2
DATA_DIM = 2
BATCH_SIZE = 1
PARAM_SPACE_BOUNDS = {'low': -1.5, 'high': 1.5}
PARAM_GRID_SIZE = 1_000

CONFIDENCE_LEVEL = 0.90

B = 20_000
B_PRIME = 10_000
MONTE_CARLO_SIZE = 2_000  # MC draws per grid point for the diagnostics themselves

SIMULATE

import torch

from lf2i.simulator.gaussian import GaussianMean

gm = GaussianMean(
    likelihood_cov=LIKELIHOOD_COV,
    prior='gaussian',
    prior_kwargs={'loc': PRIOR_LOC, 'cov': PRIOR_COV},
    poi_space_bounds=PARAM_SPACE_BOUNDS,
    poi_grid_size=PARAM_GRID_SIZE,
    poi_dim=PARAM_DIM,
    data_dim=DATA_DIM,
    batch_size=BATCH_SIZE,
)

FIT THE P-VALUE FUNCTION

from lf2i.inference import LF2I
from lf2i.test_statistics import Posterior
from sbi.inference import SNPE

posterior_ts = Posterior(poi_dim=PARAM_DIM, estimator=SNPE())
lf2i = LF2I(test_statistic=posterior_ts)

x_obs = gm(param=torch.tensor([[0.5, -0.3]])).reshape(1, DATA_DIM)

_ = lf2i.inference(
    x=x_obs,
    evaluation_grid=gm.poi_grid,
    confidence_level=CONFIDENCE_LEVEL,
    calibration_method='p-values',
    calibration_model='nn',
    simulator=gm,
    b=B,
    b_prime=B_PRIME,
)

DIAGNOSE THE FIT

For each \(\theta\) on the evaluation grid, monte_carlo_pvalue_diagnostics draws fresh Monte Carlo samples, evaluates the test statistic, and compares the fitted p-value/CDF model’s predictions to the empirical CDF at that \(\theta\) – returning, per grid point, a normalized CRPS (0 = perfect fit, 1 = no better than ignoring \(\theta\) and using the marginal CDF) and pinball losses at a range of quantile levels.

from lf2i.diagnostics.monte_carlo_methods import monte_carlo_pvalue_diagnostics

evaluation_grid_out, estimation_errors = monte_carlo_pvalue_diagnostics(
    test_statistic=posterior_ts,
    calibration_model=lf2i.calibration_model,
    simulator=gm,
    evaluation_grid=gm.poi_grid,
    monte_carlo_size=MONTE_CARLO_SIZE,
)
list(estimation_errors.keys())

Calibration score heatmap

from lf2i.plot.calibration_diagnostics import calibration_score_plot

calibration_score_plot(
    parameters=gm.poi_grid.numpy(),
    scores=estimation_errors['crps'],
    score_label='CRPS (normalized)',
    param_dim=PARAM_DIM,
    title='Normalized CRPS of the fitted p-value function across the parameter grid',
)

CDF comparison at a single theta

from lf2i.plot.calibration_diagnostics import plot_cdf_comparison

plot_cdf_comparison(
    test_statistic=posterior_ts,
    calib_model=lf2i.calibration_model[f'{CONFIDENCE_LEVEL:.2f}'],
    theta_eval=torch.tensor([[0.5, -0.3]]),
    simulator=gm,
    monte_carlo_size=MONTE_CARLO_SIZE,
    title='Fitted vs. empirical CDF at theta = (0.5, -0.3)',
)

Combined diagnostic panel

from lf2i.plot.calibration_diagnostics import calibration_cdf_panel

calibration_cdf_panel(
    evaluation_grid=gm.poi_grid,
    estimation_errors=estimation_errors,
    test_statistic=posterior_ts,
    calibration_model=lf2i.calibration_model,
    simulator=gm,
    param_dim=PARAM_DIM,
    score_key='crps',
    monte_carlo_size=MONTE_CARLO_SIZE,
)