run_sbc won't sample from inferred Direct Posterior

[paarth-dudani]

Describe the bug
I am trying to establish an SBC pipeline. I inferred a few independent SNPE estimates using concurrent futures. However, the run_sbc sampler won't sample from my inferred Direct Posterior.

To Reproduce

1. Versions
   Python version: 3.9.13
   SBI version: 0.23.1

2. Minimal code:

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

t = np.linspace(0,4,n)
input_array = t
initial_condition = 2.0

# 2D SBC
thetas = unifrom_beta_sigma.sample((num_sbc_samples,))
xs = simulator_2D(thetas, num_sbc_samples, t, initial_condition)

# results[0] is one of the 3 direct posteriors
num_posterior_samples = 1_000
num_workers = 1
ranks, dap_samples = run_sbc(
    thetas, xs, results[0], num_posterior_samples=num_posterior_samples, num_workers=num_workers

Code for context:

if __name__ == "__main__":
  with concurrent.futures.ProcessPoolExecutor(max_workers = 3) as executor:
    futures = [executor.submit(SNE_training_basic_multi_round, neural_net, proposal, simulator, observed_data, learning_rate, num_rounds, num_sim, initial_condition, hidden_features) for i in array_indices]

results = [futures[i].result() for i in range(len(futures))]

3. Error message
   No error message but the sampler won't sample:

Expected behavior
Fast sampling

[janfb]

Hi @paarth-dudani thanks for reporting this.

The run_sbc does nothing to the posterior than just sampling from the posterior for the sbc ranking. Two questions:

1. Have you checked that the posterior you obtained from your parallelization routine is "healthy", i.e., it's fast to sample / has no leakage etc?
2. Is the simulator_2d simulator the same one you used for training SNPE?

Cheers,
Jan

[paarth-dudani]

Hi @janfb

1. The posterior is indeed quick to sample from. I don't exactly know how to check for leakage. If you are referring to the warning given in the following link on the sbi website: FAQ, then I get no such warnings while training. The training is smooth.
2. Yes, it is the same simulator that I used for both training and the SBC procedure.

I get this for other parallel estimates as well. What should I do?

[janfb]

OK, good to know that posterior itself is fast to sample.

Then it's indeed something within the run_sbc method. Can you please try it with setting use_batched_sampling=False (True by default). I can imagen that because you have 2D data and are using an embedding network, passing the entire batch of 200 xs at once to the density estimator might be too much.

Does this help?

[paarth-dudani]

No, unfortunately it does not help.

Just as a clarification. My data is still 1D, it's just that I am trying to infer the standard deviation of the noise as well. This is the reason why my prior is 2D although my data is still 1D.

As a result, I have also not used an embedding network in this example.

As you can see, it takes too long and I don't see any output for minutes.

[patriglesias]

I am experiencing the same problem with the new version: No error message but the sampler won't sample

[janfb]

OK, interesting! Can one of you please share a minimal code example? Or at least some details on the use-case, the sbi method, the hyperparameters etc?