refactor: rename inference trainer classes (#1605)

* fix: rename LikelihoodEstimator to LikelihoodEstimatorTrainer

* fix: rename PosteriorEstimator to PosteriorEstimatorTrainer

* fix: Rename RatioEstimator to RatioEstimatorTrainer

* fix: Update incorrect import for RatioEstimator

Author: abelaba

sbi/diagnostics/misspecification.py

@@ -11,7 +11,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from sbi.inference.trainers.npe.npe_base import PosteriorEstimator
+from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer
 from sbi.neural_nets.estimators import UnconditionalDensityEstimator
 from sbi.utils.metrics import check_c2st
 
@@ -113,7 +113,7 @@ def calculate_p_misspecification(
 def calc_misspecification_mmd(
     x_obs: Tensor,
     x: Tensor,
-    inference: Optional[PosteriorEstimator] = None,
+    inference: Optional[PosteriorEstimatorTrainer] = None,
     mode: str = "x_space",
     n_shuffle: int = 1_000,
     max_samples: int = 1_000,

sbi/inference/trainers/nle/mnle.py

@@ -8,14 +8,14 @@
 
 from sbi.inference.posteriors import MCMCPosterior, RejectionPosterior, VIPosterior
 from sbi.inference.potentials import likelihood_estimator_based_potential
-from sbi.inference.trainers.nle.nle_base import LikelihoodEstimator
+from sbi.inference.trainers.nle.nle_base import LikelihoodEstimatorTrainer
 from sbi.neural_nets.estimators import MixedDensityEstimator
 from sbi.sbi_types import TensorboardSummaryWriter, TorchModule
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.user_input_checks import check_prior
 
 
-class MNLE(LikelihoodEstimator):
+class MNLE(LikelihoodEstimatorTrainer):
     """Mixed Neural Likelihood Estimation (MNLE) [1].
 
     Like NLE, but designed to be applied to data with mixed types, e.g., continuous

sbi/inference/trainers/nle/nle_a.py

@@ -5,12 +5,12 @@
 
 from torch.distributions import Distribution
 
-from sbi.inference.trainers.nle.nle_base import LikelihoodEstimator
+from sbi.inference.trainers.nle.nle_base import LikelihoodEstimatorTrainer
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 
 
-class NLE_A(LikelihoodEstimator):
+class NLE_A(LikelihoodEstimatorTrainer):
     """Neural Likelihood Estimation (NLE) as in Papamakarios et al. (2019) [1].
 
     [1] Sequential Neural Likelihood: Fast Likelihood-free Inference with

sbi/inference/trainers/nle/nle_base.py

@@ -26,7 +26,7 @@
 from sbi.utils.torchutils import assert_all_finite
 
 
-class LikelihoodEstimator(NeuralInference, ABC):
+class LikelihoodEstimatorTrainer(NeuralInference, ABC):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
@@ -82,7 +82,7 @@ def append_simulations(
         from_round: int = 0,
         algorithm: str = "SNLE",
         data_device: Optional[str] = None,
-    ) -> "LikelihoodEstimator":
+    ) -> "LikelihoodEstimatorTrainer":
         r"""Store parameters and simulation outputs to use them for later training.
 
         Data are stored as entries in lists for each type of variable (parameter/data).

sbi/inference/trainers/npe/__init__.py

@@ -1,7 +1,7 @@
 from sbi.inference.trainers.npe.mnpe import MNPE  # noqa: F401
 from sbi.inference.trainers.npe.npe_a import NPE_A  # noqa: F401
 from sbi.inference.trainers.npe.npe_b import NPE_B  # noqa: F401
-from sbi.inference.trainers.npe.npe_base import PosteriorEstimator  # noqa: F401
+from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer  # noqa: F401
 from sbi.inference.trainers.npe.npe_c import NPE_C  # noqa: F401
 
 SNPE_A = NPE_A

sbi/inference/trainers/npe/npe_a.py

@@ -13,7 +13,7 @@
 from torch.distributions import Distribution, MultivariateNormal
 
 from sbi.inference.posteriors.direct_posterior import DirectPosterior
-from sbi.inference.trainers.npe.npe_base import PosteriorEstimator
+from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
 from sbi.sbi_types import TensorboardSummaryWriter, TorchModule
 from sbi.utils import torchutils
@@ -26,7 +26,7 @@
 from sbi.utils.torchutils import BoxUniform, assert_all_finite, atleast_2d
 
 
-class NPE_A(PosteriorEstimator):
+class NPE_A(PosteriorEstimatorTrainer):
     r"""Neural Posterior Estimation algorithm as in Papamakarios et al. (2016) [1].
 
     [1] *Fast epsilon-free Inference of Simulation Models with Bayesian

sbi/inference/trainers/npe/npe_b.py

@@ -8,13 +8,13 @@
 from torch.distributions import Distribution
 
 import sbi.utils as utils
-from sbi.inference.trainers.npe.npe_base import PosteriorEstimator
+from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer
 from sbi.neural_nets.estimators.shape_handling import reshape_to_sample_batch_event
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 
 
-class NPE_B(PosteriorEstimator):
+class NPE_B(PosteriorEstimatorTrainer):
     r"""Neural Posterior Estimation algorithm (NPE-B) as in Lueckmann et al. (2017) [1].
 
     [1] *Flexible statistical inference for mechanistic models of neural

sbi/inference/trainers/npe/npe_base.py

@@ -45,7 +45,7 @@
 from sbi.utils.torchutils import assert_all_finite
 
 
-class PosteriorEstimator(NeuralInference, ABC):
+class PosteriorEstimatorTrainer(NeuralInference, ABC):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
@@ -103,7 +103,7 @@ def append_simulations(
         proposal: Optional[DirectPosterior] = None,
         exclude_invalid_x: Optional[bool] = None,
         data_device: Optional[str] = None,
-    ) -> "PosteriorEstimator":
+    ) -> "PosteriorEstimatorTrainer":
         r"""Store parameters and simulation outputs to use them for later training.
 
         Data are stored as entries in lists for each type of variable (parameter/data).

sbi/inference/trainers/npe/npe_c.py

@@ -10,7 +10,7 @@
 from torch.distributions import Distribution, MultivariateNormal, Uniform
 
 from sbi.inference.posteriors.direct_posterior import DirectPosterior
-from sbi.inference.trainers.npe.npe_base import PosteriorEstimator
+from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
     reshape_to_sample_batch_event,
@@ -28,7 +28,7 @@
 from sbi.utils.torchutils import BoxUniform, assert_all_finite
 
 
-class NPE_C(PosteriorEstimator):
+class NPE_C(PosteriorEstimatorTrainer):
     """Neural Posterior Estimation algorithm (NPE-C) as in Greenberg et al. (2019). [1]
 
     [1] *Automatic Posterior Transformation for Likelihood-free Inference*,

sbi/inference/trainers/nre/nre_a.py

@@ -7,13 +7,13 @@
 from torch import Tensor, nn, ones
 from torch.distributions import Distribution
 
-from sbi.inference.trainers.nre.nre_base import RatioEstimator
+from sbi.inference.trainers.nre.nre_base import RatioEstimatorTrainer
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.torchutils import assert_all_finite
 
 
-class NRE_A(RatioEstimator):
+class NRE_A(RatioEstimatorTrainer):
     """AALR, here known as Neural Ratio Estimation algorithm (NRE-A) [1].
 
     [1] *Likelihood-free MCMC with Amortized Approximate Likelihood Ratios*, Hermans