Sampler interface

[michaeldeistler]

Main API remains unchanged

inference = SNLE(prior)  # prior is no longer required (can also be passed to `build_posterior`)
_ = inference.append_simulations(theta, x).train()
posterior = inference.build_posterior()  # posterior is now of type `MCMCPosterior`

samples = posterior.sample((100,), x=xo)

The sampler interface

The .build_posterior() method is a wrapper around the sampler interface:

inference = SNLE()  # no more prior needed

# likelihood_model will have unified API once we move to pyroflows
likelihood_model = inference.append_simulations(theta, x).train()

potential_fn, theta_transform = likelihood_estimator_based_potential(likelihood_model, prior, x_o)
posterior = MCMCPosterior(potential_fn, proposal=prior, theta_transform=theta_transform)

samples = posterior.sample((100,))

Other available samplers are:

posterior = RejectionPosterior(potential_fn, proposal=prior)
posterior = DirectPosterior(posterior_model, prior)  # only applicable to SNPE

For devs: how we deal with default_x()

When not using the sampler interface:

• x_o is always passed as before: either via posterior.set_default_x(xo) or via posterior.sample((1,), x=xo).
• set_default_x() saves the value as the .default_x property of the posterior classes.

When using the sampler interface:

• The posterior can be instantiated without knowledge of x_o (to avoid API changes). Therefore, the x_o argument to the likelihood_estimator_based_potential() method can be None (type is Optional[Tensor]), but it has no default value. The following works:

potential_fn, theta_transform = likelihood_estimator_based_potential(likelihood_model, prior, x_o=None)
posterior = MCMCPosterior(potential_fn, theta_transform=theta_transform, proposal=prior)
samples = posterior.sample((100,), x=ones((1,2))

• x_o can also be passed to the likelihood_estimator_based_potential(). In that case, the default_x property of the posterior is directly set to whatever was used in the potential_fn. The following works:

potential_fn, theta_transform = likelihood_estimator_based_potential(likelihood_model, prior, x_o=ones((1,2))
posterior = MCMCPosterior(potential_fn, theta_transform=theta_transform, proposal=prior)
samples = posterior.sample((100,))
print(posterior.default_x) # -> tensor([[1.0, 1.0]])

• The potential_fn.x_o is always the most recently used data. It is not the same as posterior.default_x

API changes

.sample(..., sample_with="mcmc") is no longer supported

The posteriors are now sampler specific, e.g. MCMCPosterior. Thus, one can not change whether one samples the posterior with MCMC or rejection sampling (or vi) at .sample(). We give an explicit error in that case. The user has to rerun .build_posterior(sample_with="mcmc").

.sample_conditional() is no longer part of the posterior

.sample_conditional requires using the sampler-interface. The new API is:

from sbi.analysis import parameter_conditional_potential

likelihood_model = inference.append_simulations(theta, x).train()
pot_fn, theta_tf = likelihood_estimator_based_potential(likelihood_model, prior, xo)

cond_pot_fn, cond_tf, cond_prior = conditional_potential(pot_fn, theta_tf, prior)

cond_posterior = MCMCPosterior(cond_pot_fn, cond_prior, cond_tf)
cond_samples = cond_posterior.sample((100,))

Posterior is no longer aware of how many rounds it has been trained on

The posterior no longer has the attribute _num_trained_rounds. This affects the print(posterior) statement: before, it printed whether the posterior is amortized or not. This is no longer done.

Changes in code-base structure

The main change is that the posterior classes are no longer specific to the inference method (e.g. LikelihoodBasedPosterior), but instead to the sampler (e.g. MCMCPosterior). What is specific to the inference method are the potentials (e.g. likelihood_potential()). These potentials all lie in a new folder sbi/inference/potentials.

Notes

• loading old posteriors under the new version is not supported
• the tutorials still have to be adapted

[codecov-commenter]

Codecov Report

Merging #573 (f78cf2a) into main (1ea7246) will decrease coverage by 0.12%.
The diff coverage is 68.35%.

❗ Current head f78cf2a differs from pull request most recent head 5bf8e57. Consider uploading reports for the commit 5bf8e57 to get more accurate results

@@            Coverage Diff             @@
##             main     #573      +/-   ##
==========================================
- Coverage   66.92%   66.79%   -0.13%     
==========================================
  Files          56       67      +11     
  Lines        4193     4186       -7     
==========================================
- Hits         2806     2796      -10     
- Misses       1387     1390       +3     

Flag	Coverage Δ
unittests	66.79% <68.35%> (-0.13%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
sbi/analysis/plot.py	80.00% <ø> (+6.66%)	⬆️
sbi/analysis/sensitivity_analysis.py	15.48% <0.00%> (ø)
sbi/inference/snle/snle_a.py	100.00% <ø> (ø)
sbi/inference/snre/snre_a.py	50.00% <ø> (ø)
sbi/inference/snre/snre_b.py	100.00% <ø> (ø)
sbi/samplers/mcmc/init_strategy.py	38.88% <ø> (ø)
sbi/samplers/mcmc/mcmc.py	71.79% <ø> (ø)
sbi/samplers/mcmc/slice.py	98.57% <ø> (ø)
sbi/samplers/mcmc/slice_numpy.py	78.49% <ø> (ø)
sbi/samplers/mcmc/slice_numpy_vectorized.py	100.00% <ø> (ø)
... and 46 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1ea7246...5bf8e57. Read the comment docs.

[jan-matthis]

Great work, this should be very helpful down the road!

I think the most crucial design decision to discuss is this one: "Currently, x can not be changed after instantiation -> one has to rerun .build_posterior() with the new x.". Should it stay this way, or do we want to support this?

As I understand it, the main difficulty doing this within the proposed design is that the posterior class would need to rebuild or change its potential function to support this. I was wondering whether it would be advantageous to pass a potential function builder rather than the fully-built potential to the init of the posterior instead. Since all potential function builders are treated the same way (https://github.com/mackelab/sbi/blob/sampler/sbi/inference/potentials/likelihood_based_potential.py#L36-L41, https://github.com/mackelab/sbi/blob/sampler/sbi/inference/potentials/posterior_based_potential.py#L39-L44, https://github.com/mackelab/sbi/blob/sampler/sbi/inference/potentials/ratio_based_potential.py#L36-L41), couldn't calling the builder function happen inside the posterior class instead, e.g., in a method of the posterior base class?

About the proposed change to the simple interface: I would strongly favor keeping things down to a single line of code -- e.g., by keeping support for setting x_o on sample. Alternatively, infer could return both the estimator as well as a fully-built posterior.

With respect to naming things:

• I'd prefer sticking to likelihood estimator, posterior estimator, etc. rather than model
• I'd suggest renaminglikelihood_potential to likelihood_estimation_potential or likelihood_estimator_potential to avoid confusion

[michaeldeistler]

Thanks for the feedback JM!

I think the most crucial design decision to discuss is this one: "Currently, x can not be changed after instantiation -> one has to rerun .build_posterior() with the new x.". Should it stay this way, or do we want to support this?

I'm also still unsure about whether this is ok or not. One thing to keep in mind though: at least for SNPE, the posterior_model will have a pyro-API (after we use flowtorch/pyroflows). So x can be swapped out easily in that object.

As I understand it, the main difficulty doing this within the proposed design is that the posterior class would need to rebuild or change its potential function to support this.

My main reason for having the potential_fn as a Callable is that it can be plugged into any MCMC sampler (even those that are not supported by sbi). A potential_fn_provider (as suggested by you) could also support this, but it is a bit more cryptic and might be a bit confusing to users.

couldn't calling the builder function happen inside the posterior class instead, e.g., in a method of the posterior base class?

yes, that would be possible with the potential_fn_provider.

About the proposed change to the simple interface: I would strongly favor keeping things down to a single line of code -- e.g., by keeping support for setting x_o on sample. Alternatively, infer could return both the estimator as well as a fully-built posterior.

if we go for a potential_fn_provider this could easily be supported.

• I'd prefer sticking to likelihood estimator, posterior estimator, etc. rather than model

Agreed!

• I'd suggest renaminglikelihood_potential to likelihood_estimation_potential or likelihood_estimator_potential to avoid confusion

I like it, but it's quite verbose. How about likelihood_based_potential (not much shorter either though)?

[jan-matthis]

if we go for a potential_fn_provider this could easily be supported.

Exactly :)

This is what got me thinking about passing the builder.

I like it, but it's quite verbose. How about likelihood_based_potential (not much shorter either though)?

I think being verbose is okay if it helps avoid confusion.

To me, the essential bit to get across is that we are building potential functions (always the same quantity -- likelihood x prior) based on different estimators (estimators of posterior/likelihood/likelihood-to-evidence ratio etc.). likelihood_potential, posterior_potential sounded a bit as if these would estimate different quantities -- rather than the same quantity on the basis of different estimators. So that's why I think having estimator or estimation as part of the naming would be helpful to avoid confusion.

My main reason for having the potential_fn as a Callable is that it can be plugged into any MCMC sampler (even those that are not supported by sbi). A potential_fn_provider (as suggested by you) could also support this, but it is a bit more cryptic and might be a bit confusing to users.

We could still support this case by keeping the functions for the two-step build process. I think, ultimately support for external samplers would need to be documented in a tutorial or FAQ entry, where we'd explain this. I think being able to keep a single line infer example and avoiding API changes as much as possible would be great.

[michaeldeistler]

I think you convinced me regarding the potential_fn_provider. But let's discuss tomorrow and then we can implement it :)

[janfb]

Yes, this looks great already!

I agree with keeping the names likelihood_estimator etc.

I would vote for a more descriptive name for the function that returns the potential function, e.g., build_potential_fn, or, if we change the API accordingly, potential_fn_provider.

More detailed comments in the review below.

[janfb]

Great effort!

I added some comments and questions.

Overall, I suggest to rename {likelihood, posterior, ratio}_model back to ``{likelihood, posterior, ratio}_estimator`, and to simplify to construction and naming of the potential functions as commented.

[jan-matthis]

Looks great -- think the interface turned out very well! So much fewer API changes now. Only left a few brief comments. If there remains something to be discussed, let me know :)

[janfb]

Overall great work! Looks good to. me now, except for some small comments, some of which we can move to issues and future PRs.

[janfb]

I see, but then one could just have this method like it is here and in DirectPosterior (there are almost the same no?) and let the child method call the base method with appropriate default args? e.g., init_method="prior" here, and init_method="posterior" in the DirectPosterior?

[janfb]

could also be sth for an issue and future PR