Running TARP with BNRE and MCMC

[diego-gonher]

Hi there!

I have been trying use TARP to evaluate the performance of the inferences obtained using the implemented BNRE algorithm in this package. To be able to directly compare it as close as possible to some work that I've done in the past, I'd like to perform all the inferences using HMC. Thus, after training the estimator, I'm trying to run TARP with the following code:

# set MCMC parameters to use Pyro NUTS.
mcmc_parameters = dict(num_chains=4,
                       warmup_steps=250,
                       init_strategy="proposal",
                       method="nuts_pyro")


# create the posterior object
bnre_posterior = MCMCPosterior(potential_fn, 
                               proposal=prior, 
                               theta_transform=parameter_transform, 
                               **mcmc_parameters)

# run TARP
expected_coverage, ideal_coverage = run_tarp(
    thetas,
    xs,
    bnre_posterior,
    references=None,  # optional, defaults to uniform samples across parameter space.
    num_posterior_samples=1_000)

But I get the following error:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-10-dd3077f6147a> in <cell line: 4>()
      2 from sbi.analysis import plot_tarp
      3 
----> 4 expected_coverage, ideal_coverage = run_tarp(
      5     thetas,
      6     xs,

2 frames
/usr/local/lib/python3.10/dist-packages/sbi/diagnostics/tarp.py in run_tarp(thetas, xs, posterior, references, num_posterior_samples, num_workers, show_progress_bar, distance, num_bins, z_score_theta)
     62     num_tarp_samples, dim_theta = thetas.shape
     63 
---> 64     posterior_samples = get_posterior_samples_on_batch(
     65         xs,
     66         posterior,

/usr/local/lib/python3.10/dist-packages/sbi/utils/diagnostics_utils.py in get_posterior_samples_on_batch(xs, posterior, sample_shape, num_workers, show_progress_bar, use_batched_sampling)
     36         # has shape (num_samples, batch_size, dim_parameters)
     37         if use_batched_sampling:
---> 38             posterior_samples = posterior.sample_batched(
     39                 sample_shape, x=xs, show_progress_bars=show_progress_bar
     40             )

/usr/local/lib/python3.10/dist-packages/sbi/inference/posteriors/mcmc_posterior.py in sample_batched(self, sample_shape, x, method, thin, warmup_steps, num_chains, init_strategy, init_strategy_parameters, num_workers, mp_context, show_progress_bars)
    420 
    421         assert (
--> 422             method == "slice_np_vectorized"
    423         ), "Batched sampling only supported for vectorized samplers!"
    424 

AssertionError: Batched sampling only supported for vectorized samplers!

I've tried to look in the documentation and the source code, but I am still not sure how to solve this. Does anyone have any suggestions? Thank you in advance, and I hope you're having a great day!

[michaeldeistler]

Hi there! Thanks for reporting!

The reason you are seeing this is that we use sample_batched() within TARP to vectorize across observations, which makes the diagnostic tool much faster. However, not all samplers support vectorized sampling (e.g., HMC), which is why it breaks.

If there is a possibility to do so, I would recommend you to switch to our slice-sampler for MCMC, for several reasons:

1. We found it to be very reliable
2. It will not give this bug :)
3. Even if we fix the bug by looping over observations with HMC, it will be much (at least an order of magnitude) slower than the vectorized slice samplers.

Hope this helps!
Michael

[diego-gonher]

I see, thank you so much Michael for the clarification! I got it to work using ''slice_np_vectorized" instead of "nuts_pyro", which is perfect for now!

[janfb]

Another workaround would be using the internal _run_tarp method. This method takes the generated posterior samples as argument. Thus, to do so, you would need to generate them yourself, just like it is done in the user-facing run_tarp method:

sbi/sbi/diagnostics/tarp.py

Lines 62 to 83 in 4b3d409

	num_tarp_samples, dim_theta = thetas.shape
	posterior_samples = get_posterior_samples_on_batch(
	xs,
	posterior,
	(num_posterior_samples,),
	num_workers,
	show_progress_bar=show_progress_bar,
	)
	assert posterior_samples.shape == (
	num_posterior_samples,
	num_tarp_samples,
	dim_theta,
	), f"Wrong posterior samples shape for TARP: {posterior_samples.shape}"
	# Sample reference points uniformly if not provided
	if references is None:
	references = get_tarp_references(thetas)
	return _run_tarp(
	posterior_samples, thetas, references, distance, num_bins, z_score_theta
	)

The only difference would be that you pass the use_batched_sampling=False to get_posterior_samples_on_batch. This will be slow, but it will allow you to use an HMC sampler within the MCMCPosterior.

We will expose the use_batched_sampling=False in a future release.

I hope this helps!
Best,
Jan

[janfb]

You can also pass num_workers=<your number of CPU cores> to get_posterior_samples_on_batch to speed things up.

[diego-gonher]

Thank you so much for this Jan, this is extremely helpful! Does this mean that I could apply this approach to use TARP with other, more "standard" techniques, as long as I generate the posterior samples beforehand?

[janfb]

Yes, TARP is agnostic to the source of posterior samples. It just needs the a batch of test thetas, the corresponding posterior samples, and potentially a set of reference samples. The posterior samples could have been generated from a ground posterior, or via PyMC, or via SBI...

[diego-gonher]

Amazing, this is very useful since I want to compare the results of using sbi compared to more traditional methods in my subfield. I managed to run TARP using this approach for that! Thank you so much again!