refactor: unify build_posterior in base class for all trainers (#1610)

* fix: spelling update

* refactor(NeuralInference): define methods for building the posterior

* refactor(nle): update build_posterior method to use parent class implementation

* refactor(npe): update build_posterior method to use parent class implementation

* refactor(nre): update build_posterior method to use parent class implementation

* refactor(npse): implement required abstract method from parent class

* refactor(npse, fmpe): remove and replace build_posterior method with superclass method

* refactor: update build_posterior estimator annotation

* refactor: move _get_potential_function method from build_posterior to _create_posterior

* refactor: update _get_potential_function return type

* refactor: Update conditional arrangement for _create_posterior

* docs: update vector_field_potential docstring

* test(build_posterior): raise error for invalid density_estimator

* refactor(build_posterior): move estimator and prior checking into helper functions

* refactor: rearrange method order

* refactor: remove comment and update return type for abstract method train

* refactor: update method ordering for trainer classes to follow convention

* refactor(FMPE, NPSE): add vectorfield_sampling_parameters parameter

* test: update NRE_A to NRE base class

* test(FMPE, NPSE): add vectorfield_sampling_parameters to build_posterior argument

Author: abelaba

docs/tutorials/00_getting_started.ipynb

@@ -117,7 +117,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The `sbi` toolbox uses neural networks to learn the relationship between parameters and data. In this exampmle, we will use neural perform posterior estimation (NPE). To run NPE, we first instatiate a trainer, which we call `inference`:"
+    "The `sbi` toolbox uses neural networks to learn the relationship between parameters and data. In this example, we will use neural perform posterior estimation (NPE). To run NPE, we first instatiate a trainer, which we call `inference`:"
    ]
   },
   {

sbi/inference/potentials/likelihood_based_potential.py

@@ -26,7 +26,7 @@ def likelihood_estimator_based_potential(
     prior: Distribution,  # type: ignore
     x_o: Optional[Tensor],
     enable_transform: bool = True,
-) -> Tuple[Callable, TorchTransform]:
+) -> Tuple["LikelihoodBasedPotential", TorchTransform]:
     r"""Returns potential :math:`\log(p(x_o|\theta)p(\theta))` for likelihood estimator.
 
     It also returns a transformation that can be used to transform the potential into

sbi/inference/potentials/ratio_based_potential.py

@@ -1,7 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
-from typing import Callable, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 from torch import Tensor, nn
@@ -18,7 +18,7 @@ def ratio_estimator_based_potential(
     prior: Distribution,
     x_o: Optional[Tensor],
     enable_transform: bool = True,
-) -> Tuple[Callable, TorchTransform]:
+) -> Tuple["RatioBasedPotential", TorchTransform]:
     r"""Returns the potential for ratio-based methods.
 
     It also returns a transformation that can be used to transform the potential into

sbi/inference/potentials/vector_field_potential.py

@@ -37,6 +37,9 @@ def vector_field_estimator_based_potential(
         enable_transform: Whether to enable transforms. Not supported yet.
         **kwargs: Additional keyword arguments passed to
             `VectorFieldBasedPotential`.
+    Returns:
+        The potential function and a transformation that maps
+        to unconstrained space.
     """
     device = str(next(vector_field_estimator.parameters()).device)
 
@@ -60,9 +63,9 @@ def __init__(
         vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Optional[Distribution],  # type: ignore
         x_o: Optional[Tensor] = None,
+        device: Union[str, torch.device] = "cpu",
         iid_method: Literal["fnpe", "gauss", "auto_gauss", "jac_gauss"] = "auto_gauss",
         iid_params: Optional[Dict[str, Any]] = None,
-        device: Union[str, torch.device] = "cpu",
         neural_ode_backend: str = "zuko",
         neural_ode_kwargs: Optional[Dict[str, Any]] = None,
     ):
@@ -78,11 +81,11 @@ def __init__(
             vector_field_estimator: The neural network modelling the vector field.
             prior: The prior distribution.
             x_o: The observed data at which to evaluate the posterior.
+            device: The device on which to evaluate the potential.
             iid_method: Which method to use for computing the score in the iid setting.
                 We currently support "fnpe", "gauss", "auto_gauss", "jac_gauss".
             iid_params: Parameters for the iid method, for arguments see
                 `IIDScoreFunction`.
-            device: The device on which to evaluate the potential.
             neural_ode_backend: The backend to use for the neural ODE. Currently,
                 only "zuko" is supported.
             neural_ode_kwargs: Additional keyword arguments for the neural ODE.