Expected Coverage in TSNPE

[hoangdzung]

Hi everyone,

I am currently working with TSNPE and using the following code snippet from the tutorial:

from sbi.inference import NPE
from sbi.utils import RestrictedPrior, get_density_thresholder

inference = NPE(prior)
proposal = prior
for _ in range(num_rounds):
    theta = proposal.sample((num_sims,))
    x = simulator(theta)
    _ = inference.append_simulations(theta, x).train(force_first_round_loss=True)
    posterior = inference.build_posterior().set_default_x(x_o)

    accept_reject_fn = get_density_thresholder(posterior, quantile=1e-4)
    proposal = RestrictedPrior(prior, accept_reject_fn, sample_with="rejection")

I would like to understand where the expected coverage (Algorithm 2 in the paper) is computed in the code. Could you point me to the relevant implementation or function that handles this?
Is there any automated mechanism that performs early stopping?

Thank you for your time!

[michaeldeistler]

Hi there,

the snippet above does not include the computation of expected coverage. To compute it, run

from sbi.diagnostics import run_sbc
from sbi.analysis.plot import sbc_rank_plot

num_sbc_samples = 200  # choose a number of sbc runs, should be ~100s

# generate ground truth parameters and corresponding simulated observations for SBC.
thetas = prior.sample((num_sbc_samples,))
xs = simulator(thetas)

# run SBC: for each inference we draw 1000 posterior samples.
num_posterior_samples = 1_000
num_workers = 1
ranks, dap_samples = run_sbc(
    thetas,
    xs,
    posterior,
    num_posterior_samples=num_posterior_samples,
    num_workers=num_workers,
    reduce_fns=posterior.log_prob,
)
f, ax = sbc_rank_plot(ranks, 1_000, plot_type="cdf")

Note that this diagnostic is only correct after the first round of TSNPE. Later rounds will require truncation of the posterior log_prob which is not directly implemented in sbi.

For more background on diagnostics, see this tutorial on SBC.

And yes, .train() performs early stopping based on a held-out validation set.

Hope that helps
Michael

[hoangdzung]

Thanks a lot for your quick and helpful reply. I will check this out.