log_prob fails for Restricted Prior when sampling with SIR

[rmehta1987]

I am trying to use a restricted prior as a proposal and when creating it with the parameter SIR the process_prior function fails because it cannot calculate the log prob.

File "/python3.8/site-packages/sbi/utils/user_input_checks.py", line 354, in check_prior_attributes                                             
    prior.log_prob(theta)

File "/utils/restriction_estimator.py", line 833, in log_prob                                                       
    torch.log(self.prior_acceptance(**prior_acceptance_params))

TypeError: log(): argument 'input' (position 1) must be Tensor, not NoneType

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "posterior_train_bins.py", line 114, in <module>
    bin_proposal, num_parameters, prior_returns_numpy = process_prior(proposal)

Also does this in general when simply doing:

propr = RestrictedPrior(prior, accept, posterior=temp, sample_with="sir")
python3.8/site-packages/sbi/utils/restriction_estimator.py", line 833, in log_prob
    torch.log(self.prior_acceptance(**prior_acceptance_params))

It seems to be like there is no acceptance_rate analgous for SIR, whereas "rejection" has one explicitly defined. Is this intentional ? SIR provides a significant speed boost, so it would be nice to use. It looks like it would be possible to implement in the "sampling_importance_resampling" function.

[michaeldeistler]

Thanks for creating this issue! I can not reproduce the error though. Two questions:

1. Does the following code work for you:

import torch
from torch import eye, ones, zeros
from torch.distributions import MultivariateNormal

from sbi.inference import SNPE, DirectPosterior
from sbi.utils import RestrictedPrior, get_density_thresholder, BoxUniform

num_dim = 2
x_o = zeros((1, num_dim))

prior = BoxUniform(-ones(num_dim), ones(num_dim))
theta = prior.sample((100,))
x = theta + torch.randn(theta.shape)

inference = SNPE(prior=prior)
_ = inference.append_simulations(theta, x).train()

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, posterior=posterior, sample_with="sir")
theta = proposal.sample((1000,))

2. At what line in your code does the error occur? During .train()?

[rmehta1987]

The code you gave me works perfectly. It is when I add:
logprb = proposal.log_prob(theta)

that the error is raised.

2.) Since the process_prior(prior,...) function checks if the prior can produce log_prob and I cannot seem to get the log_prob with the restricted prior when using SIR, the error occurs during process_prior.

[michaeldeistler]

Got it, thanks!

The quick solution for now is to not run process_prior on the RestrictedPrior. If you follow the tutorial here, everything should work:

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

inference = SNPE(prior)
proposal = prior
for _ in range(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")

Note that the RestrictedPrior is not even passed as proposal in this tutorial.

I'll think a bit more about whether we should remove the log_prob requirement from process_prior (since it is only strictly required for SNLE, SNRE, and SNPE-C, but not for single-round NPE and neither for the restricted methods).

[rmehta1987]

I was hoping to use the restrictedPrior as a proposal itself in the initialization of the SNPE(prior = RestrictedPrior) (which also checks if the prior has a log_prob) as I am trying to learn two different posterior distributions. The restrictedPrior would serve as a proposal for the second distribution. For now it works if I just use sample_with="rejection" since it returns a log_prob.

Another issue with restrictedPrior I can raise separately, is there any way to pickle it ? Currently it cannot pickle the get_density_thresholder function so it raises an error.

[michaeldeistler]

Thanks for clarifying! I agree that we should allow passing the RestrictedPrior to be passed at unit. I will fix this soon, but I don’t have time this week. Until then, you can very easily work around this:

In SNPE, the prior passed at initialization is the prior, not the proposal. In TSNPE (and in single round NPE), the prior (passed at init) is not used during training at all (it is only used after training to reject samples from the approximate posterior that lie outside of the prior). Thus, you do not have to pass anything:

inference = SNPE()

or you simply pass the original prior:

Inference = SNPE(prior=prior)

The proposal is simply the distribution from the parameters theta are drawn. It does not require a log_prob method and you do not have to pass it at initialization. You will only lose the feature to reject samples from the approximate posterior that lie outside of the prior.

[rmehta1987]

That is only true for SNPE correct ? For SNLE/SNRE it requires the prior in order to calculate the posterior distribution, P(theta|X) = P(X|theta)P(theta) for SNLE and P(theta|X) = r(X,theta)P(theta) for SNRE

[michaeldeistler]

Yes, that's only true for SNPE.

[michaeldeistler]

For now, I recommend simply passing the original prior at initialization. I shouldn't make a big difference. I'll fix this next week.

[janfb]

@michaeldeistler any chance we can get this fix in for the release?
Or can you sketch a solution here quickly, then I can add it? Thanks!

[janfb]

I think the problem here is that when calling process_prior with a custom prior like the RestrictedPrior, the process function will check whether it has log_prob and sample and returns Tensors (Yes), but then wrap everything as a CustomPriorWrapper and thereby hiding the _accept_reject_fn needed for SIR.

One solution would be making RestrictedPrior a proper torch Distribution.

Or change the wrapping in user_input_checks e.g., not wrap if it is a RestrictedPrior.