Rename SN(P)E -> N(P)E in tests, tutorials, and docs (#1239)

* refactor: move all inference methods into `trainers` folder (#1238)

* move all inference methods into `trainers` folder

* remove the "S" from all inference methods

* Rename FMPE

* All imports are fixed now

* Add dummy files to snpe/snre/snle folders to raise an explicit error

* formatting

* feedback from review

* Rename SN(P)E -> N(P)E in tests

* Rename SN(P)E -> N(P)E in all tutorials

* Rename SN(P)E -> N(P)E in all documentation

* Remove `infer` method from docs

* Remove simulate_for_sbi from tutorial

* update implemented methods

* rename in tests

* fixup

* modify tutorial

docs/docs/faq/question_01_leakage.md

@@ -1,7 +1,7 @@
 # What should I do when my 'posterior samples are outside the prior support' in SNPE?
 
-When working with **multi-round** SNPE, you might have experienced the following
-warning:
+When working with **multi-round** NPE (i.e., SNPE), you might have experienced the
+following warning:
 
 ```python
 Only x% posterior samples are within the prior support. It may take a long time to
@@ -10,19 +10,37 @@ collect the remaining 10000 samples. Consider interrupting (Ctrl-C) and switchin
 ```
 
 The reason for this issue is described in more detail
-[here](https://arxiv.org/abs/2002.03712) and
+[here](https://arxiv.org/abs/2210.04815),
+[here](https://arxiv.org/abs/2002.03712), and
 [here](https://arxiv.org/abs/1905.07488). The following fixes are possible:
 
+- use truncated proposals for SNPE (TSNPE)
+```python
+from sbi.inference import NPE
+from sbi.utils import RestrictedPrior, get_density_thresholder
+
+inference = NPE(prior)
+proposal = prior
+for _ in range(num_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")
+```
+
 - sample with MCMC: `samples = posterior((num_samples,), x=x_o, sample_with_mcmc=True)`.
 This approach will make sampling slower, but samples will not "leak".
 
-- resort to single-round SNPE and (if necessary) increase your simulation budget.
+- resort to single-round NPE and (if necessary) increase your simulation budget.
 
 - if your prior is either Gaussian (torch.distributions.MultivariateNormal) or Uniform
 (sbi.utils.BoxUniform), you can avoid leakage by using a mixture density network as
 density estimator. I.e., set `density_estimator='mdn'` when creating the `SNPE`
 inference object. When running inference, there should be a print statement "Using
 SNPE-C with non-atomic loss".
 
-- use a different algorithm, e.g., SNRE and SNLE. Note, however, that these algorithms
-can have different issues and potential pitfalls.
+- use a different algorithm, e.g., Sequential NRE and Sequential NLE. Note, however,
+that these algorithms can have different issues and potential pitfalls.

docs/docs/faq/question_02_nans.md

@@ -16,15 +16,15 @@ Importantly, in order for neural network training work well, the floating point
 number should still be in a reasonable range, i.e., maybe a few standard
 deviations outside of 'good' values.
 
-If you are running **multi-round** SNPE, however, things can go fully wrong if
+If you are running **multi-round** NPE (SNPE), however, things can go fully wrong if
 invalid data are encountered. In that case, you will get the following warning
 
 ```python
 When invalid simulations are excluded, multi-round SNPE-C can leak into the regions
 where parameters led to invalid simulations. This can lead to poor results.
 ```
 
-Hence, if you are running multi-round SNPE and a significant fraction of
+Hence, if you are running multi-round NPE and a significant fraction of
 simulations returns at least one invalid number, we strongly recommend manually
 replacing the value in your simulation code as described above (or resorting to
-single-round SNPE, or using a different `sbi` method entirely).
+single-round NPE, or using a different `sbi` method entirely).

docs/docs/faq/question_04_gpu.md

@@ -10,7 +10,7 @@ Yes, we support GPU training. When creating the inference object in the flexible
 interface, you can pass the `device` as an argument, e.g.,
 
 ```python
-inference = SNPE(prior, device="cuda", density_estimator="maf")
+inference = NPE(prior, device="cuda", density_estimator="maf")
 ```
 
 The device is set to `"cpu"` by default. But it can be set to anything, as long
@@ -33,7 +33,7 @@ device="cuda:0"), covariance_matrix=torch.eye(2, device="cuda:0"))
 
 Whether or not you reduce your training time when training on a GPU depends on
 the problem at hand. We provide a couple of default density estimators for
-`SNPE`, `SNLE` and `SNRE`, e.g., a mixture density network
+`NPE`, `NLE` and `NRE`, e.g., a mixture density network
 (`density_estimator="mdn"`) or a Masked Autoregressive Flow
 (`density_estimator="maf"`). For these default density estimators, we do **not**
 expect a speed-up. This is because the underlying neural networks are relatively

docs/docs/faq/question_06_resume_training.md

@@ -5,7 +5,7 @@ Many clusters have a time limit, and `sbi` might exceed this limit. You can
 circumvent this problem by stopping and resuming training:
 
 ```python
-inference = SNPE(prior=prior)
+inference = NPE(prior=prior)
 inference = inference.append_simulations(theta, x)
 inference.train(max_num_epochs=300)  # Pick `max_num_epochs` such that it does not exceed the runtime.
 

docs/docs/faq/question_07_custom_prior.md

@@ -81,7 +81,7 @@ Note that in `custom_prior_wrapper_kwargs` you can pass additinal arguments for
 the wrapper, e.g., `validate_args` or `arg_constraints` see the `Distribution`
 documentation for more details.
 
-If you are using `sbi` < 0.17.2 and use `SNLE` the code above will produce a
+If you are using `sbi` < 0.17.2 and use `NLE` the code above will produce a
 `NotImplementedError` (see [#581](https://github.com/mackelab/sbi/issues/581)).
-In this case, you need to update to a newer version of `sbi` or use `SNPE`
+In this case, you need to update to a newer version of `sbi` or use `NPE`
 instead.

docs/docs/index.md

@@ -5,7 +5,7 @@ interface:
 
 ```python
 import torch
-from sbi.inference import SNPE
+from sbi.inference import NPE
 
 # define shifted Gaussian simulator.
 def simulator(θ): return θ + torch.randn_like(θ)
@@ -15,7 +15,7 @@ def simulator(θ): return θ + torch.randn_like(θ)
 x = simulator(θ)
 
 # choose sbi method and train
-inference = SNPE()
+inference = NPE()
 inference.append_simulations(θ, x).train()
 
 # do inference given observed data
@@ -103,8 +103,8 @@ The methods then proceed by
    full space of parameters consistent with the data and the prior, i.e. the
    posterior distribution. The posterior assigns high probability to parameters
    that are consistent with both the data and the prior, and low probability to
-   inconsistent parameters. While SNPE directly learns the posterior
-   distribution, SNLE and SNRE need an extra MCMC sampling step to construct a
+   inconsistent parameters. While NPE directly learns the posterior
+   distribution, NLE and NRE need an extra MCMC sampling step to construct a
    posterior.
 4. If needed, an initial estimate of the posterior can be used to adaptively
    generate additional informative simulations.

docs/docs/reference/inference.md

@@ -1,48 +1,48 @@
 # Inference
 
-## Algorithms
+## Training algorithms
 
-::: sbi.inference.snpe.snpe_a.SNPE_A
+::: sbi.inference.trainers.npe.npe_a.NPE_A
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.snpe.snpe_c.SNPE_C
+::: sbi.inference.trainers.npe.npe_c.NPE_C
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.fmpe.fmpe_base.FMPE
+::: sbi.inference.trainers.fmpe.fmpe.FMPE
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.npse.npse.NPSE
+::: sbi.inference.trainers.npse.npse.NPSE
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.snle.snle_a.SNLE_A
+::: sbi.inference.trainers.nle.nle_a.NLE_A
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.snre.snre_a.SNRE_A
+::: sbi.inference.trainers.nre.nre_a.NRE_A
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.snre.snre_b.SNRE_B
+::: sbi.inference.trainers.nre.nre_b.NRE_B
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.snre.snre_c.SNRE_C
+::: sbi.inference.trainers.nre.nre_c.NRE_C
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
 
-::: sbi.inference.snre.bnre.BNRE
+::: sbi.inference.trainers.nre.bnre.BNRE
     selection:
       filters: [ "!^_", "^__", "!^__class__" ]
       inherited_members: true
@@ -59,9 +59,7 @@
 
 ## Helpers
 
-::: sbi.inference.base.infer
-
-::: sbi.inference.base.simulate_for_sbi
+::: sbi.inference.trainers.base.simulate_for_sbi
 
 ::: sbi.utils.user_input_checks.process_prior
 

sbi/inference/trainers/npe/npe_a.py

@@ -230,10 +230,10 @@ def train(
     def correct_for_proposal(
         self,
         density_estimator: Optional[TorchModule] = None,
-    ) -> "SNPE_A_MDN":
+    ) -> "NPE_A_MDN":
         r"""Build mixture of Gaussians that approximates the posterior.
 
-        Returns a `SNPE_A_MDN` object, which applies the posthoc-correction required in
+        Returns a `NPE_A_MDN` object, which applies the posthoc-correction required in
         SNPE-A.
 
         Args:
@@ -272,8 +272,8 @@ def correct_for_proposal(
                 """
             proposal = self._proposal_roundwise[-1]
 
-        # Create the SNPE_A_MDN
-        wrapped_density_estimator = SNPE_A_MDN(
+        # Create the NPE_A_MDN
+        wrapped_density_estimator = NPE_A_MDN(
             flow=density_estimator,  # type: ignore
             proposal=proposal,
             prior=self._prior,
@@ -372,7 +372,7 @@ def _expand_mog(self, eps: float = 1e-5):
                     param.grad = None  # let autograd construct a new gradient
 
 
-class SNPE_A_MDN(ConditionalDensityEstimator):
+class NPE_A_MDN(ConditionalDensityEstimator):
     """Generates a posthoc-corrected MDN which approximates the posterior.
 
     TODO: Adapt this class to the new `DensityEstimator` interface. Maybe even to a
@@ -639,7 +639,7 @@ def _proposal_posterior_transformation(
             covariances_post, means_pp, precisions_pp, means_d, precisions_d
         )
 
-        logits_post = SNPE_A_MDN._logits_posterior(
+        logits_post = NPE_A_MDN._logits_posterior(
             means_post,
             precisions_post,
             covariances_post,
@@ -655,7 +655,7 @@ def _proposal_posterior_transformation(
 
     def _set_state_for_mog_proposal(self) -> None:
         """
-        Set state variables of the SNPE_A_MDN instance every time `set_proposal()`
+        Set state variables of the NPE_A_MDN instance every time `set_proposal()`
         is called, i.e. every time a posterior is build using
         `SNPE_A.build_posterior()`.
 

tests/analysis_test.py

@@ -5,7 +5,7 @@
 import torch
 
 from sbi.analysis import ActiveSubspace, conditional_corrcoeff, conditional_pairplot
-from sbi.inference import SNPE
+from sbi.inference import NPE
 from sbi.utils import BoxUniform, get_1d_marginal_peaks_from_kde
 from sbi.utils.torchutils import process_device
 
@@ -34,7 +34,7 @@ def simulator(theta):
     theta = prior.sample((300,)).to(device)
     x = simulator(theta)
 
-    inf = SNPE(prior=prior, device=device)
+    inf = NPE(prior=prior, device=device)
     _ = inf.append_simulations(theta, x).train(max_num_epochs=10)
 
     observation = torch.zeros(num_dim)

tests/base_test.py

@@ -5,7 +5,7 @@
 import torch
 
 from sbi import utils
-from sbi.inference import SNPE, infer
+from sbi.inference import NPE, infer
 
 
 def test_infer():
@@ -16,12 +16,12 @@ def test_infer():
     def simulator(parameter_set):
         return 1.0 + parameter_set + torch.randn(parameter_set.shape) * 0.1
 
-    posterior = infer(simulator, prior, method="SNPE_A", num_simulations=10)
+    posterior = infer(simulator, prior, method="NPE_A", num_simulations=10)
     assert posterior is not None, "Most basic use of 'infer' failed"
     posterior = infer(
         simulator,
         prior,
-        method="SNPE_A",
+        method="NPE_A",
         num_simulations=10,
         init_kwargs={"num_components": 5},
         train_kwargs={"max_num_epochs": 2},
@@ -35,7 +35,7 @@ def test_get_dataloaders(training_batch_size):
     N = 1000
     validation_fraction = 0.1
 
-    inferer = SNPE()
+    inferer = NPE()
     inferer.append_simulations(torch.ones(N), torch.zeros(N))
     _, val_loader = inferer.get_dataloaders(
         0,

tests/embedding_net_test.py

@@ -9,7 +9,7 @@
 from torch.distributions import MultivariateNormal
 
 from sbi import utils
-from sbi.inference import FMPE, SNLE, SNPE, SNRE
+from sbi.inference import FMPE, NLE, NPE, NRE
 from sbi.neural_nets import classifier_nn, flowmatching_nn, likelihood_nn, posterior_nn
 from sbi.neural_nets.embedding_nets import (
     CNNEmbedding,
@@ -25,7 +25,7 @@
 
 
 @pytest.mark.mcmc
-@pytest.mark.parametrize("method", ["SNPE", "SNLE", "SNRE", "FMPE"])
+@pytest.mark.parametrize("method", ["NPE", "NLE", "NRE", "FMPE"])
 @pytest.mark.parametrize("num_dim", [1, 2])
 @pytest.mark.parametrize("embedding_net", ["mlp"])
 def test_embedding_net_api(
@@ -49,19 +49,19 @@ def test_embedding_net_api(
     else:
         raise NameError(f"{embedding_net} not supported.")
 
-    if method == "SNPE":
+    if method == "NPE":
         density_estimator = posterior_nn("maf", embedding_net=embedding)
-        inference = SNPE(
+        inference = NPE(
             prior, density_estimator=density_estimator, show_progress_bars=False
         )
-    elif method == "SNLE":
+    elif method == "NLE":
         density_estimator = likelihood_nn("maf", embedding_net=embedding)
-        inference = SNLE(
+        inference = NLE(
             prior, density_estimator=density_estimator, show_progress_bars=False
         )
-    elif method == "SNRE":
+    elif method == "NRE":
         classifier = classifier_nn("resnet", embedding_net_x=embedding)
-        inference = SNRE(prior, classifier=classifier, show_progress_bars=False)
+        inference = NRE(prior, classifier=classifier, show_progress_bars=False)
     elif method == "FMPE":
         vectorfield_net = flowmatching_nn(model="mlp", embedding_net=embedding)
         inference = FMPE(
@@ -106,7 +106,7 @@ def test_embedding_api_with_multiple_trials(
     )
 
     density_estimator = posterior_nn("maf", embedding_net=embedding_net)
-    inference = SNPE(prior, density_estimator=density_estimator)
+    inference = NPE(prior, density_estimator=density_estimator)
 
     _ = inference.append_simulations(theta, x).train(max_num_epochs=5)
 
@@ -165,7 +165,7 @@ def simulator1d(theta):
             1, num_channels, *[1 for _ in range(len(input_shape))]
         )
 
-    trainer = SNPE(prior=prior, density_estimator=estimator_provider)
+    trainer = NPE(prior=prior, density_estimator=estimator_provider)
     trainer.append_simulations(theta, x).train(max_num_epochs=2)
     posterior = trainer.build_posterior().set_default_x(xo)
 
@@ -217,7 +217,7 @@ def test_npe_with_with_iid_embedding_varying_num_trials(trial_factor=50):
         z_score_x="none",  # turn off z-scoring because of NaN encodings.
         z_score_theta="independent",
     )
-    inference = SNPE(prior, density_estimator=density_estimator)
+    inference = NPE(prior, density_estimator=density_estimator)
 
     # do not exclude invalid x, as we padded with nans.
     _ = inference.append_simulations(theta, x, exclude_invalid_x=False).train(