Problems with proposal prior and SNLE

Hello, I'm new to the field... I have a one dimensional model theta and x are 1D, and I'm trying to use SNLE with maf and mdn. What I'd like to do is to train the estimator on my data (theta was sampled from a uniform prior between 0 and 1) and then get a posterior with a more restrictive prior such as a narrow gaussian (of course centered within 0 and 1) to visualize the effect of the prior on my posterior estimate.

Please can someone tell me if I'm doing something wrong ?

Running on numpy: v1.21.6 Running on torch: v1.11.0 Running on sbi: v0.19.0

Here is my code:

  lk = utils.get_nn_models.likelihood_nn('mdn', z_score_theta='independent', z_score_x='independent', hidden_features=50,  num_components=10)
  inference = SNLE(density_estimator=lk)
  inference = inference.append_simulations(theta, x)
  density_estimator = inference.train()
  posterior = inference.build_posterior(density_estimator, prior=prior, sample_with="rejection")
  posterior_samples = posterior.sample((10000,), x=torch.as_tensor([x_0]))

When I use a uniform prior it works but not when I use a gaussian:

prior = utils.BoxUniform(low=.1*torch.ones(num_dim), 
                         high=.9*torch.ones(num_dim))

prior = torch.distributions.Normal(torch.tensor(np.array([.5], dtype='float32')), torch.tensor(np.array([1], dtype='float32')))

If I use sample_with='rejection' I get this error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/var/folders/_b/qt7bmb5154l22n5k1rjgvbym0000gn/T/ipykernel_17287/484018902.py in <module>
      5     density_estimator = inference.train()
      6     posterior = inference.build_posterior(density_estimator, prior=prior, sample_with="rejection")
----> 7     posterior_samples = posterior.sample((10000,), x=torch.as_tensor([pci_base]))
      8 
      9     # plot posterior samples

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/posteriors/rejection_posterior.py in sample(self, sample_shape, x, max_sampling_batch_size, num_samples_to_find_max, num_iter_to_find_max, m, sample_with, show_progress_bars)
    158             num_iter_to_find_max=num_iter_to_find_max,
    159             m=m,
--> 160             device=self._device,
    161         )
    162 

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/samplers/rejection/rejection.py in rejection_sample(potential_fn, proposal, theta_transform, num_samples, show_progress_bars, warn_acceptance, max_sampling_batch_size, num_samples_to_find_max, num_iter_to_find_max, m, device)
     78         learning_rate=0.01,
     79         num_to_optimize=max(1, int(num_samples_to_find_max / 10)),
---> 80         show_progress_bars=False,
     81     )
     82 

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/utils/sbiutils.py in gradient_ascent(potential_fn, inits, theta_transform, num_iter, num_to_optimize, learning_rate, save_best_every, show_progress_bars, interruption_note)
    815 
    816             optimizer.zero_grad()
--> 817             probs = potential_fn(theta_transform.inv(optimize_inits)).squeeze()
    818             loss = -probs.sum()
    819             loss.backward()

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/samplers/rejection/rejection.py in potential_over_proposal(theta)
     68     # Define a potential as the ratio between target distribution and proposal.
     69     def potential_over_proposal(theta):
---> 70         return potential_fn(theta) - proposal.log_prob(theta)
     71 
     72     # Search for the maximum of the ratio.

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/potentials/likelihood_based_potential.py in __call__(self, theta, track_gradients)
     89             theta=theta.to(self.device),
     90             net=self.likelihood_estimator,
---> 91             track_gradients=track_gradients,
     92         )
     93 

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/potentials/likelihood_based_potential.py in _log_likelihoods_over_trials(x, theta, net, track_gradients)
    121     # at neural network evaluation time.
    122     theta_repeated, x_repeated = match_theta_and_x_batch_shapes(
--> 123         theta=atleast_2d(theta), x=atleast_2d(x)
    124     )
    125     assert (

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/utils/sbiutils.py in match_theta_and_x_batch_shapes(theta, x)
    541     x_repeated = x.repeat_interleave(theta_batch_size, dim=0)
    542     # Repeat theta as ABCABC.
--> 543     theta_repeated = theta.repeat(x_batch_size, 1)
    544 
    545     # Double check: batch size for log prob evaluation must match.

RuntimeError: Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor

If I use sample_with='mcmc' I get this error:

`RuntimeError                              Traceback (most recent call last)
/var/folders/_b/qt7bmb5154l22n5k1rjgvbym0000gn/T/ipykernel_17287/1838947700.py in <module>
      4 density_estimator = inference.train()
      5 posterior = inference.build_posterior(density_estimator, prior=prior, sample_with="mcmc")
----> 6 posterior_samples = posterior.sample((10000,), x=torch.as_tensor([pci_base]))

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/posteriors/mcmc_posterior.py in sample(self, sample_shape, x, method, thin, warmup_steps, num_chains, init_strategy, init_strategy_parameters, init_strategy_num_candidates, mcmc_parameters, mcmc_method, sample_with, num_workers, show_progress_bars)
    278             num_workers,
    279             show_progress_bars,
--> 280             **init_strategy_parameters,
    281         )
    282         num_samples = torch.Size(sample_shape).numel()

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/posteriors/mcmc_posterior.py in _get_initial_params(self, init_strategy, num_chains, num_workers, show_progress_bars, **kwargs)
    420         else:
    421             initial_params = torch.cat(
--> 422                 [init_fn() for _ in range(num_chains)]  # type: ignore
    423             )
    424 

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/posteriors/mcmc_posterior.py in <listcomp>(.0)
    420         else:
    421             initial_params = torch.cat(
--> 422                 [init_fn() for _ in range(num_chains)]  # type: ignore
    423             )
    424 

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/inference/posteriors/mcmc_posterior.py in <lambda>()
    354         elif init_strategy == "resample":
    355             return lambda: resample_given_potential_fn(
--> 356                 proposal, potential_fn, transform=transform, **kwargs
    357             )
    358         elif init_strategy == "latest_sample":

~/INS_Code/sbi_env/lib/python3.7/site-packages/sbi/samplers/mcmc/init_strategy.py in resample_given_potential_fn(proposal, potential_fn, transform, num_candidate_samples, num_batches, **kwargs)
    110         probs /= probs.sum()
    111 
--> 112         idxs = torch.multinomial(probs, 1, replacement=False)
    113         # Return transformed sample.
    114         return transform(init_param_candidates[idxs, :])  # type: ignore

RuntimeError: number of categories cannot exceed 2^24

sbi-dev / sbi

Problems with proposal prior and SNLE #737