refactor: build_posterior Posterior Configuration Using Dataclasses (#1619)

* refactor: add posterior parameters dataclass

* refactor: add posterior_parameters parameter to build_posterior method

* test: add unit test to check signature consistency

* refactor: update posterior_parameters initialization

* refactor: add posterior_parameters resolution method

* refactor(build_posterior): update posterior_parameter resolution method

* test: add VectorFieldBasedPotential to test argument

* refactor: add fields to VectorFieldPosteriorParameters and update _resolve_posterior_parameters method

* refactor: add warning and value error for conflicting posterior_parameter resolution

* test: add test for passing conflicting parameters to build_posterior

* doc(posterior_parameters): add fields to every dataclass docstring

* refactor: remove x_shape field

* refactor(mnle): add importance_sampling_parameters and ImportanceSamplingPosteriorParameters

* refactor: update str to Literal for neural_ode_backend

* refactor: modularize functions for checking posterior parameter value mismatches

* refactor: add PosteriorParameters base class and add field validations

* refactor(PosteriorParameters): update validations

* test(PosteriorParameters): add new tests

* add license header

* chore: update spacing between docstring and code implementation

* refactor: update posterior_parameters resolution flow

* test(posterior_parameters): update tests and add test fixture

* refactor(npe): validate build_posterior prior argument for rejection_sampling

* refactor(NPE): update error message for build_posterior rejection sampling

* test(NPE): check build_posterior with rejection sampling fails for empty prior

Author: abelaba

sbi/inference/posteriors/posterior_parameters.py

@@ -66,7 +66,7 @@ def __init__(
         device: Union[str, torch.device] = "cpu",
         iid_method: Literal["fnpe", "gauss", "auto_gauss", "jac_gauss"] = "auto_gauss",
         iid_params: Optional[Dict[str, Any]] = None,
-        neural_ode_backend: str = "zuko",
+        neural_ode_backend: Literal["zuko"] = "zuko",
         neural_ode_kwargs: Optional[Dict[str, Any]] = None,
     ):
         r"""

sbi/inference/potentials/vector_field_potential.py

@@ -1,8 +1,10 @@
 # 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/>
 
+import warnings
 from abc import ABC, abstractmethod
 from copy import deepcopy
+from dataclasses import asdict
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Callable, Dict, Literal, Optional, Tuple, Union
@@ -20,6 +22,15 @@
 from sbi.inference.posteriors.direct_posterior import DirectPosterior
 from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
 from sbi.inference.posteriors.mcmc_posterior import MCMCPosterior
+from sbi.inference.posteriors.posterior_parameters import (
+    DirectPosteriorParameters,
+    ImportanceSamplingPosteriorParameters,
+    MCMCPosteriorParameters,
+    PosteriorParameters,
+    RejectionPosteriorParameters,
+    VIPosteriorParameters,
+    VectorFieldPosteriorParameters,
+)
 from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
 from sbi.inference.posteriors.vector_field_posterior import VectorFieldPosterior
 from sbi.inference.posteriors.vi_posterior import VIPosterior
@@ -401,6 +412,7 @@ def build_posterior(
         sample_with: Literal[
             "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
         ],
+        posterior_parameters: Optional[PosteriorParameters],
         **kwargs,
     ) -> NeuralPosterior:
         r"""Method for building posteriors.
@@ -422,6 +434,8 @@ def build_posterior(
                 - "direct"
                 - "sde"
                 - "ode"
+            posterior_parameters: Configuration passed to the init method for the
+                posterior. Must be of type PosteriorParameters.
             **kwargs: Additional method-specific parameters.
 
         Returns:
@@ -431,12 +445,16 @@ def build_posterior(
         prior = self._resolve_prior(prior)
         estimator, device = self._resolve_estimator(estimator)
 
+        posterior_parameters = self._resolve_posterior_parameters(
+            sample_with, posterior_parameters, **kwargs
+        )
+
         self._posterior = self._create_posterior(
             estimator,
             prior,
             sample_with,
             device,
-            **kwargs,
+            posterior_parameters,
         )
 
         # Store models at end of each round.
@@ -508,6 +526,144 @@ def _resolve_estimator(
 
         return estimator, device
 
+    def _resolve_posterior_parameters(
+        self,
+        sample_with: Literal[
+            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
+        ],
+        posterior_parameters: Optional[PosteriorParameters],
+        **kwargs,
+    ) -> PosteriorParameters:
+        """
+        Resolve posterior parameters based on the sampling strategy.
+
+        If `posterior_parameters` is provided, it is returned directly.
+
+        If `posterior_parameters` is not provided, this method extracts
+        sampling-specific parameters from `kwargs` using predefined keys
+        to instantiate the appropriate posterior parameters dataclass.
+
+        Raises:
+            NotImplementedError: If an unsupported `sample_with` method is provided.
+            ValueError: If posterior_parameter and a configuration dictionary are passed
+                together.
+
+        Args:
+            sample_with: The posterior sampling method to use.
+            posterior_parameters: Optional preconstructed posterior parameter object.
+            **kwargs: Additional parameters to construct the posterior parameters.
+
+        Returns:
+            A dataclass instance containing the resolved posterior
+            parameters.
+        """
+
+        if posterior_parameters is not None:
+            self._validate_no_duplicate_parameters(**kwargs)
+            self._validate_posterior_parameters_consistency(
+                posterior_parameters, **kwargs
+            )
+        else:
+            # Resolve parameters passed through kwargs and convert
+            # into a subclass of PosteriorParameters
+            if sample_with == "direct":
+                params = kwargs.get("direct_sampling_parameters", {}) or {}
+                posterior_parameters = DirectPosteriorParameters(**params)
+            elif sample_with == "mcmc":
+                params = kwargs.get("mcmc_parameters", {}) or {}
+                posterior_parameters = MCMCPosteriorParameters(
+                    method=kwargs.get("mcmc_method", "slice_np_vectorized"), **params
+                )
+            elif sample_with in ("ode", "sde"):
+                params = kwargs.get("vectorfield_sampling_parameters", {}) or {}
+                posterior_parameters = VectorFieldPosteriorParameters(**params)
+            elif sample_with == "rejection":
+                params = kwargs.get("rejection_sampling_parameters", {}) or {}
+                posterior_parameters = RejectionPosteriorParameters(**params)
+            elif sample_with == "vi":
+                params = kwargs.get("vi_parameters", {}) or {}
+                posterior_parameters = VIPosteriorParameters(
+                    vi_method=kwargs.get("vi_method", "rKL"), **params
+                )
+            elif sample_with == "importance":
+                params = kwargs.get("importance_sampling_parameters", {}) or {}
+                posterior_parameters = ImportanceSamplingPosteriorParameters(**params)
+            else:
+                raise NotImplementedError(
+                    "Posterior parameter construction not implemented for",
+                    f"'{sample_with}'",
+                )
+
+        return posterior_parameters
+
+    def _validate_no_duplicate_parameters(self, **kwargs) -> None:
+        """
+        Ensure parameters aren't specified in both posterior_parameters and the
+        posterior parameter dictionaries in the build_posterior method.
+
+        Args:
+            **kwargs: Additional parameters to construct the posterior parameters
+        """
+
+        old_style_params = {
+            "direct_sampling_parameters",
+            "mcmc_parameters",
+            "vectorfield_sampling_parameters",
+            "rejection_sampling_parameters",
+            "vi_parameters",
+            "importance_sampling_parameters",
+        }
+
+        # Check if any old-style parameters were provided
+        provided_old_params = [
+            param for param in old_style_params if kwargs.get(param) is not None
+        ]
+
+        if provided_old_params:
+            raise ValueError(
+                f"Cannot use both old-style parameters {provided_old_params} "
+                f"and new-style posterior_parameters. Please use only one approach."
+            )
+
+    def _validate_posterior_parameters_consistency(
+        self, posterior_parameters: PosteriorParameters, **kwargs
+    ) -> None:
+        """
+        This method raises a warning for mismatches between values passed in
+        mcmc_method and MCMCPosteriorParameters.method, or vi_method and
+        VIPosteriorParameters.vi_method.
+
+        Args:
+            posterior_parameters: Configuration passed to the init method for the
+                posterior.
+            kwargs: keyword arguments passed from build_posterior method.
+        """
+
+        if not isinstance(posterior_parameters, PosteriorParameters):
+            raise TypeError(
+                "posterior_parameters must be PosteriorParameters,"
+                f" got {type(posterior_parameters).__name__}",
+            )
+        elif isinstance(posterior_parameters, MCMCPosteriorParameters):
+            mcmc_method = kwargs.get("mcmc_method")
+            if (
+                mcmc_method != "slice_np_vectorized"
+                and posterior_parameters.method != mcmc_method
+            ):
+                warnings.warn(
+                    f"Conflicting mcmc_method='{mcmc_method}' ignored in favor of "
+                    f"posterior_parameters.method='{posterior_parameters.method}'",
+                    stacklevel=2,
+                )
+        elif isinstance(posterior_parameters, VIPosteriorParameters):
+            vi_method = kwargs.get("vi_method")
+            if vi_method != "rKL" and posterior_parameters.vi_method != vi_method:
+                warnings.warn(
+                    f"Conflicting vi_method='{vi_method}' ignored in favor of "
+                    f"posterior_parameters.vi_method='{posterior_parameters.vi_method}'",
+                    stacklevel=2,
+                )
+
     def _create_posterior(
         self,
         estimator: Union[RatioEstimator, ConditionalEstimator],
@@ -516,7 +672,7 @@ def _create_posterior(
             "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
         ],
         device: Union[str, torch.device],
-        **kwargs,
+        posterior_parameters: PosteriorParameters,
     ) -> NeuralPosterior:
         """
         Create a posterior object using the specified inference method.
@@ -539,83 +695,89 @@ def _create_posterior(
                 - "ode"
             device: torch device on which to train the neural net and on which to
                 perform all posterior operations, e.g. gpu or cpu.
-            **kwargs: Additional method-specific parameters.
+            posterior_parameters: Configuration passed to the init method for the
+                posterior. Must be of type PosteriorParameters.
 
         Returns:
             NeuralPosterior object.
         """
 
-        if sample_with == "direct":
+        if isinstance(posterior_parameters, DirectPosteriorParameters):
             posterior_estimator = estimator
-            assert isinstance(posterior_estimator, ConditionalDensityEstimator), (
-                f"Expected posterior_estimator to be an instance of "
-                " ConditionalDensityEstimator, "
-                f"but got {type(posterior_estimator).__name__} instead."
-            )
+            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
+                raise TypeError(
+                    f"Expected posterior_estimator to be an instance of "
+                    " ConditionalDensityEstimator, "
+                    f"but got {type(posterior_estimator).__name__} instead."
+                )
             posterior = DirectPosterior(
                 posterior_estimator=posterior_estimator,
                 prior=prior,
                 device=device,
-                **(kwargs.get("direct_sampling_parameters") or {}),
+                **asdict(posterior_parameters),
             )
-        elif sample_with in ("sde", "ode"):
+        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
             vector_field_estimator = estimator
-            assert isinstance(
-                vector_field_estimator, ConditionalVectorFieldEstimator
-            ), (
-                f"Expected vector_field_estimator to be an instance of "
-                " ConditionalVectorFieldEstimator, "
-                f"but got {type(vector_field_estimator).__name__} instead."
-            )
+            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
+                raise TypeError(
+                    f"Expected vector_field_estimator to be an instance of "
+                    " ConditionalVectorFieldEstimator, "
+                    f"but got {type(vector_field_estimator).__name__} instead."
+                )
+            if sample_with not in ("ode", "sde"):
+                raise ValueError(
+                    "`sample_with` must be either",
+                    f" 'ode' or 'sde', got '{sample_with}'",
+                )
             posterior = VectorFieldPosterior(
-                vector_field_estimator,
-                prior,
+                vector_field_estimator=vector_field_estimator,
+                prior=prior,
                 device=device,
                 sample_with=sample_with,
-                **(kwargs.get("vectorfield_sampling_parameters") or {}),
+                **asdict(posterior_parameters),
             )
         else:
             # Posteriors requiring potential_fn and theta_transform
             potential_fn, theta_transform = self._get_potential_function(
                 prior, estimator
             )
-            if sample_with == "mcmc":
+            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                 posterior = MCMCPosterior(
                     potential_fn=potential_fn,
                     theta_transform=theta_transform,
                     proposal=prior,
-                    method=kwargs.get("mcmc_method", "slice_np_vectorized"),
                     device=device,
-                    **(kwargs.get("mcmc_parameters") or {}),
+                    **asdict(posterior_parameters),
                 )
-            elif sample_with == "rejection":
+            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                 posterior = RejectionPosterior(
                     potential_fn=potential_fn,
                     proposal=prior,
                     device=device,
-                    **(kwargs.get("rejection_sampling_parameters") or {}),
+                    **asdict(posterior_parameters),
                 )
-            elif sample_with == "vi":
+            elif isinstance(posterior_parameters, VIPosteriorParameters):
                 posterior = VIPosterior(
                     potential_fn=potential_fn,
                     theta_transform=theta_transform,
                     prior=prior,
-                    vi_method=kwargs.get("vi_method", "rKL"),
                     device=device,
-                    **(kwargs.get("vi_parameters") or {}),
+                    **asdict(posterior_parameters),
                 )
-            elif sample_with == "importance":
+            elif isinstance(
+                posterior_parameters, ImportanceSamplingPosteriorParameters
+            ):
                 posterior = ImportanceSamplingPosterior(
                     potential_fn=potential_fn,
                     proposal=prior,
                     device=device,
-                    **(kwargs.get("importance_sampling_parameters") or {}),
+                    **asdict(posterior_parameters),
                 )
             else:
                 raise NotImplementedError(
-                    f"Sampling method '{sample_with}' is not supported."
+                    "Sampling method not implemented for",
+                    f"'{posterior_parameters}'",
                 )
-
         return posterior
 
     def _converged(self, epoch: int, stop_after_epochs: int) -> bool:

sbi/inference/trainers/base.py

@@ -9,6 +9,7 @@
 
 from sbi import utils as utils
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
+from sbi.inference.posteriors.posterior_parameters import VectorFieldPosteriorParameters
 from sbi.inference.trainers.npse.vector_field_inference import (
     VectorFieldEstimatorBuilder,
     VectorFieldTrainer,
@@ -65,6 +66,7 @@ def build_posterior(
         prior: Optional[Distribution] = None,
         sample_with: Literal["ode", "sde"] = "ode",
         vectorfield_sampling_parameters: Optional[Dict[str, Any]] = None,
+        posterior_parameters: Optional[VectorFieldPosteriorParameters] = None,
     ) -> NeuralPosterior:
         r"""Build posterior from the flow matching estimator.
 
@@ -89,16 +91,19 @@ def build_posterior(
                 a numerical ODE solver.
             vectorfield_sampling_parameters: Additional keyword arguments passed to
                 `VectorFieldPosterior`.
-
+            posterior_parameters: Configuration passed to the init method for
+                VectorFieldPosterior.
 
         Returns:
             Posterior $p(\theta|x)$  with `.sample()` and `.log_prob()` methods.
         """
+
         return super().build_posterior(
             estimator=vector_field_estimator,
             prior=prior,
             sample_with=sample_with,
             vectorfield_sampling_parameters=vectorfield_sampling_parameters,
+            posterior_parameters=posterior_parameters,
         )
 
     def _build_default_nn_fn(self, **kwargs) -> VectorFieldEstimatorBuilder: