refactor: bake RatioEstimator into the general ConditionalEstimator API (#1652)

* refactor: update RatioEstimator class to inherit from ConditionalEstimator

* refactor: update type hints from Union to ConditionalEstimator type

* reafactor(nre): remove RatioEstimatorBuilder protocol and update type hints

* refactor(npe, nle): update estimator builder type hints to not include RatioEstimator

* refactor: remove VectorFieldEstimatorBuilder and update estimator type hints to use DensityEstimatorBuilder

* docs: remove empty line

* test: add FMPE and NPSE for testing valid and invalid estimator builders

* chore: add loss function to RatioEstimator

* chore(nle, npe): remove ConditionalVectorFieldEstimator builder type

* test: update parameters for checking invalid builder

* chore(fmpe): assign default value for prior

* test(npse, fmpe): correct test methods

* chore: rename DensityEstimatorBuilder to ConditionalEstimatorBuilder

* chore(nle, npe): update density_estimator type annotation

* test: check embedding_net property is not None

* chore: update density_estimator type annotation for mnle and mnpe

* Update sbi/diagnostics/misspecification.py

Co-authored-by: Jan <janfb@users.noreply.github.com>

* chore(nle): update import structure

* chore: update formatting

---------

Co-authored-by: Jan <janfb@users.noreply.github.com>

Author: abelaba

sbi/diagnostics/misspecification.py

@@ -155,6 +155,12 @@ def calc_misspecification_mmd(
                 "in that case the MMD is computed in the x-space.",
                 stacklevel=2,
             )
+        if inference._neural_net.embedding_net is None:
+            raise AttributeError(
+                "embedding_net attribute is None but is required for misspecification"
+                " detection."
+            )
+
         z_obs = inference._neural_net.embedding_net(x_obs).detach()
         z = inference._neural_net.embedding_net(x).detach()
     else:

sbi/inference/trainers/base.py

@@ -49,7 +49,6 @@
     ConditionalEstimator,
     ConditionalVectorFieldEstimator,
 )
-from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.sbi_types import TorchTransform
 from sbi.utils import (
     check_prior,
@@ -314,7 +313,7 @@ def train(
     def _get_potential_function(
         self,
         prior: Distribution,
-        estimator: Union[RatioEstimator, ConditionalEstimator],
+        estimator: ConditionalEstimator,
     ) -> Tuple[BasePotential, TorchTransform]:
         """Subclass-specific potential creation"""
         ...
@@ -416,7 +415,7 @@ def get_dataloaders(
 
     def build_posterior(
         self,
-        estimator: Optional[Union[RatioEstimator, ConditionalEstimator]],
+        estimator: Optional[ConditionalEstimator],
         prior: Optional[Distribution],
         sample_with: Literal[
             "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
@@ -501,8 +500,8 @@ def _resolve_prior(self, prior: Optional[Distribution]) -> Distribution:
         return prior
 
     def _resolve_estimator(
-        self, estimator: Optional[Union[RatioEstimator, ConditionalEstimator]]
-    ) -> Tuple[Union[RatioEstimator, ConditionalEstimator], str]:
+        self, estimator: Optional[ConditionalEstimator]
+    ) -> Tuple[ConditionalEstimator, str]:
         """
         Resolves the estimator and determines its device.
 
@@ -525,9 +524,9 @@ def _resolve_estimator(
             # If internal net is used device is defined.
             device = self._device
         else:
-            if not isinstance(estimator, (ConditionalEstimator, RatioEstimator)):
+            if not isinstance(estimator, ConditionalEstimator):
                 raise TypeError(
-                    "estimator must be ConditionalEstimator or RatioEstimator,"
+                    "estimator must be ConditionalEstimator,"
                     f" got {type(estimator).__name__}",
                 )
             # Otherwise, infer it from the device of the net parameters.
@@ -759,7 +758,7 @@ def _validate_posterior_parameters_consistency(
 
     def _create_posterior(
         self,
-        estimator: Union[RatioEstimator, ConditionalEstimator],
+        estimator: ConditionalEstimator,
         prior: Distribution,
         sample_with: Literal[
             "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"

sbi/inference/trainers/nle/mnle.py

@@ -14,7 +14,7 @@
 )
 from sbi.inference.trainers.nle.nle_base import LikelihoodEstimatorTrainer
 from sbi.neural_nets.estimators import MixedDensityEstimator
-from sbi.neural_nets.estimators.base import DensityEstimatorBuilder
+from sbi.neural_nets.estimators.base import ConditionalEstimatorBuilder
 from sbi.sbi_types import TensorBoardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 
@@ -34,7 +34,10 @@ class MNLE(LikelihoodEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        density_estimator: Union[Literal["mnle"], DensityEstimatorBuilder] = "mnle",
+        density_estimator: Union[
+            Literal["mnle"],
+            ConditionalEstimatorBuilder[MixedDensityEstimator],
+        ] = "mnle",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
         summary_writer: Optional[TensorBoardSummaryWriter] = None,
@@ -49,11 +52,11 @@ def __init__(
             density_estimator: If it is a string, it must be "mnle" to use the
                 preconfiugred neural nets for MNLE. Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
+                `.sample()` and must return a `MixedDensityEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nle/nle_a.py

@@ -6,7 +6,10 @@
 from torch.distributions import Distribution
 
 from sbi.inference.trainers.nle.nle_base import LikelihoodEstimatorTrainer
-from sbi.neural_nets.estimators.base import DensityEstimatorBuilder
+from sbi.neural_nets.estimators.base import (
+    ConditionalDensityEstimator,
+    ConditionalEstimatorBuilder,
+)
 from sbi.sbi_types import TensorBoardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 
@@ -23,7 +26,8 @@ def __init__(
         self,
         prior: Optional[Distribution] = None,
         density_estimator: Union[
-            Literal["nsf", "maf", "mdn", "made"], DensityEstimatorBuilder
+            Literal["nsf", "maf", "mdn", "made"],
+            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
         ] = "maf",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
@@ -39,11 +43,11 @@ def __init__(
             density_estimator: If it is a string, use a pre-configured network of the
                 provided type (one of nsf, maf, mdn, made). Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
+                `.sample()` and must return a `ConditionalDensityEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nle/nle_base.py

@@ -26,7 +26,7 @@
 from sbi.inference.trainers.base import NeuralInference
 from sbi.neural_nets import likelihood_nn
 from sbi.neural_nets.estimators import ConditionalDensityEstimator
-from sbi.neural_nets.estimators.base import DensityEstimatorBuilder
+from sbi.neural_nets.estimators.base import ConditionalEstimatorBuilder
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
 )
@@ -40,7 +40,8 @@ def __init__(
         self,
         prior: Optional[Distribution] = None,
         density_estimator: Union[
-            Literal["nsf", "maf", "mdn", "made"], DensityEstimatorBuilder
+            Literal["nsf", "maf", "mdn", "made"],
+            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
         ] = "maf",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
@@ -57,11 +58,11 @@ def __init__(
             density_estimator: If it is a string, use a pre-configured network of the
                 provided type (one of nsf, maf, mdn, made). Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
+                `.sample()` and must return a `ConditionalDensityEstimator`.
 
         See docstring of `NeuralInference` class for all other arguments.
         """

sbi/inference/trainers/npe/mnpe.py

@@ -15,7 +15,7 @@
 )
 from sbi.inference.trainers.npe.npe_c import NPE_C
 from sbi.neural_nets.estimators import MixedDensityEstimator
-from sbi.neural_nets.estimators.base import DensityEstimatorBuilder
+from sbi.neural_nets.estimators.base import ConditionalEstimatorBuilder
 from sbi.sbi_types import TensorBoardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 
@@ -32,7 +32,10 @@ class MNPE(NPE_C):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        density_estimator: Union[Literal["mnpe"], DensityEstimatorBuilder] = "mnpe",
+        density_estimator: Union[
+            Literal["mnpe"],
+            ConditionalEstimatorBuilder[MixedDensityEstimator],
+        ] = "mnpe",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
         summary_writer: Optional[TensorBoardSummaryWriter] = None,
@@ -47,11 +50,11 @@ def __init__(
             density_estimator: If it is a string, it must be "mnpe" to use the
                 preconfigured neural nets for MNPE. Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
+                `.sample()` and must return a `MixedDensityEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/npe/npe_a.py

@@ -18,7 +18,7 @@
 )
 from sbi.neural_nets.estimators.base import (
     ConditionalDensityEstimator,
-    DensityEstimatorBuilder,
+    ConditionalEstimatorBuilder,
 )
 from sbi.sbi_types import TensorBoardSummaryWriter
 from sbi.utils import torchutils
@@ -53,7 +53,8 @@ def __init__(
         self,
         prior: Optional[Distribution] = None,
         density_estimator: Union[
-            Literal["mdn_snpe_a"], DensityEstimatorBuilder
+            Literal["mdn_snpe_a"],
+            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
         ] = "mdn_snpe_a",
         num_components: int = 10,
         device: str = "cpu",
@@ -71,14 +72,15 @@ def __init__(
             density_estimator: If it is a string (only "mdn_snpe_a" is valid), use a
                 pre-configured mixture of densities network. Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`. Note that until the last round only a single (multivariate)
-                Gaussian component is used for training (seeAlgorithm 1 in [1]). In the
-                last round, this component is replicated `num_components` times, its
-                parameters are perturbed with a very small noise, and then the last
+                `.sample()` and must return a `ConditionalDensityEstimator`.
+                Note that until the last round only a single (multivariate) Gaussian
+                component is used for training (seeAlgorithm 1 in [1]). In the last
+                round, this component is replicated `num_components` times,
+                its parameters are perturbed with a very small noise, and then the last
                 training round is done with the expanded Gaussian mixture as estimator
                 for the proposal posterior.
             num_components: Number of components of the mixture of Gaussians in the

sbi/inference/trainers/npe/npe_b.py

@@ -11,7 +11,10 @@
 from sbi.inference.trainers.npe.npe_base import (
     PosteriorEstimatorTrainer,
 )
-from sbi.neural_nets.estimators.base import DensityEstimatorBuilder
+from sbi.neural_nets.estimators.base import (
+    ConditionalDensityEstimator,
+    ConditionalEstimatorBuilder,
+)
 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
@@ -37,7 +40,8 @@ def __init__(
         self,
         prior: Optional[Distribution] = None,
         density_estimator: Union[
-            Literal["nsf", "maf", "mdn", "made"], DensityEstimatorBuilder
+            Literal["nsf", "maf", "mdn", "made"],
+            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
         ] = "maf",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
@@ -52,11 +56,11 @@ def __init__(
             density_estimator: If it is a string, use a pre-configured network of the
                 provided type (one of nsf, maf, mdn, made). Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
+                `.sample()` and must return a `ConditionalDensityEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/npe/npe_base.py

@@ -15,9 +15,7 @@
 from torch.utils.tensorboard.writer import SummaryWriter
 from typing_extensions import Self
 
-from sbi.inference.posteriors import (
-    DirectPosterior,
-)
+from sbi.inference.posteriors import DirectPosterior
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.posteriors.posterior_parameters import (
     DirectPosteriorParameters,
@@ -34,7 +32,7 @@
 )
 from sbi.neural_nets import posterior_nn
 from sbi.neural_nets.estimators import ConditionalDensityEstimator
-from sbi.neural_nets.estimators.base import DensityEstimatorBuilder
+from sbi.neural_nets.estimators.base import ConditionalEstimatorBuilder
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
     reshape_to_sample_batch_event,
@@ -59,7 +57,8 @@ def __init__(
         self,
         prior: Optional[Distribution] = None,
         density_estimator: Union[
-            Literal["nsf", "maf", "mdn", "made"], DensityEstimatorBuilder
+            Literal["nsf", "maf", "mdn", "made"],
+            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
         ] = "maf",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
@@ -76,11 +75,11 @@ def __init__(
             density_estimator: If it is a string, use a pre-configured network of the
                 provided type (one of nsf, maf, mdn, made). Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
+                `.sample()` and must return a `ConditionalDensityEstimator`.
 
         See docstring of `NeuralInference` class for all other arguments.
         """

sbi/inference/trainers/npe/npe_c.py

@@ -15,7 +15,7 @@
 )
 from sbi.neural_nets.estimators.base import (
     ConditionalDensityEstimator,
-    DensityEstimatorBuilder,
+    ConditionalEstimatorBuilder,
 )
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
@@ -73,7 +73,8 @@ def __init__(
         self,
         prior: Optional[Distribution] = None,
         density_estimator: Union[
-            Literal["nsf", "maf", "mdn", "made"], DensityEstimatorBuilder
+            Literal["nsf", "maf", "mdn", "made"],
+            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
         ] = "maf",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
@@ -88,12 +89,11 @@ def __init__(
             density_estimator: If it is a string, use a pre-configured network of the
                 provided type (one of nsf, maf, mdn, made). Alternatively, a function
                 that builds a custom neural network, which adheres to
-                `DensityEstimatorBuilder` protocol can be provided. The function will
-                be called with the first batch of simulations (theta, x), which can
+                `ConditionalEstimatorBuilder` protocol can be provided. The function
+                will be called with the first batch of simulations (theta, x), which can
                 thus be used for shape inference and potentially for z-scoring. The
                 density estimator needs to provide the methods `.log_prob` and
-                `.sample()`.
-
+                `.sample()` and must return a `ConditionalDensityEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.