Inference

Training algorithms

NPE_A

Bases: PosteriorEstimatorTrainer

Neural Posterior Estimation algorithm as in Papamakarios et al. (2016) [1].

[1] Fast epsilon-free Inference of Simulation Models with Bayesian Conditional Density Estimation, Papamakarios et al., NeurIPS 2016. https://arxiv.org/abs/1605.06376

Like all NPE methods, this method trains a deep neural density estimator to directly approximate the posterior. Also like all other NPE methods, in the first round, this density estimator is trained with a maximum-likelihood loss.

This class implements NPE-A. NPE-A trains across multiple rounds with a maximum-likelihood loss. This will make training converge to the proposal posterior instead of the true posterior. To correct for this, SNPE-A applies a post-hoc correction after training. This correction is performed analytically and requires Mixture of Gaussians (MoG) density estimators.

Note

In multi-round SNPE-A, the number of MoG components grows multiplicatively with each round: if the proposal has L components and the density estimator has K components, the corrected posterior has L×K components. For many rounds, consider using SNPE-C (APT) instead, which handles multi-round inference more efficiently.

Example:

::

import torch
from sbi.inference import NPE_A
from sbi.utils import BoxUniform

# 1. Setup simulator, prior, and observation
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
x_o = torch.randn(1, 3)  # Observed data

def simulator(theta):
    return theta + torch.randn_like(theta) * 0.1

# 2. Multi-round inference
inference = NPE_A(prior=prior, num_components=5)
proposal = prior

for round_idx in range(5):
    theta = proposal.sample((100,))
    x = simulator(theta)
    density_estimator = inference.append_simulations(theta, x).train()
    posterior = inference.build_posterior(density_estimator)
    proposal = posterior.set_default_x(x_o)

# 3. Sample from final posterior
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/npe/npe_a.py

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
        Conditional Density Estimation*, Papamakarios et al., NeurIPS 2016.
        https://arxiv.org/abs/1605.06376

    Like all NPE methods, this method trains a deep neural density estimator to
    directly approximate the posterior. Also like all other NPE methods, in the
    first round, this density estimator is trained with a maximum-likelihood loss.

    This class implements NPE-A. NPE-A trains across multiple rounds with a
    maximum-likelihood loss. This will make training converge to the proposal
    posterior instead of the true posterior. To correct for this, SNPE-A applies a
    post-hoc correction after training. This correction is performed analytically
    and requires Mixture of Gaussians (MoG) density estimators.

    Note:
        In multi-round SNPE-A, the number of MoG components grows multiplicatively
        with each round: if the proposal has L components and the density estimator
        has K components, the corrected posterior has L×K components. For many
        rounds, consider using SNPE-C (APT) instead, which handles multi-round
        inference more efficiently.

    Example:
    --------

    ::

        import torch
        from sbi.inference import NPE_A
        from sbi.utils import BoxUniform

        # 1. Setup simulator, prior, and observation
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        x_o = torch.randn(1, 3)  # Observed data

        def simulator(theta):
            return theta + torch.randn_like(theta) * 0.1

        # 2. Multi-round inference
        inference = NPE_A(prior=prior, num_components=5)
        proposal = prior

        for round_idx in range(5):
            theta = proposal.sample((100,))
            x = simulator(theta)
            density_estimator = inference.append_simulations(theta, x).train()
            posterior = inference.build_posterior(density_estimator)
            proposal = posterior.set_default_x(x_o)

        # 3. Sample from final posterior
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        density_estimator: Union[
            Literal["mdn_snpe_a"],
            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
        ] = "mdn_snpe_a",
        num_components: int = 10,
        device: str = "cpu",
        logging_level: Union[int, str] = "WARNING",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize NPE-A [1].

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Any
                object with `.log_prob()`and `.sample()` (for example, a PyTorch
                distribution) can be used.
            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
                `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()` and must return a `ConditionalDensityEstimator`.
            num_components: Number of components of the mixture of Gaussians.
                Note: In multi-round SNPE-A, the number of components grows
                multiplicatively with each round due to the analytical correction
                (L components in proposal × K components in density = L*K posterior
                components).
            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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during training.
        """

        # Catch invalid inputs.
        if not ((density_estimator == "mdn_snpe_a") or callable(density_estimator)):
            raise TypeError(
                "The `density_estimator` passed to SNPE_A needs to be a "
                "callable or the string 'mdn_snpe_a'!"
            )

        self._num_components = num_components

        # WARNING: sneaky trick ahead. We proxy the parent's `train` here,
        # requiring the signature to have `num_components`, save it for use below, and
        # continue. It's sneaky because we are using the object (self) as a namespace
        # to pass arguments between functions, and that's implicit state management.
        kwargs = del_entries(
            locals(),
            entries=("self", "__class__", "num_components"),
        )
        super().__init__(**kwargs)

    def train(
        self,
        training_batch_size: int = 200,
        learning_rate: float = 5e-4,
        validation_fraction: float = 0.1,
        stop_after_epochs: int = 20,
        max_num_epochs: int = 2**31 - 1,
        clip_max_norm: Optional[float] = 5.0,
        calibration_kernel: Optional[Callable] = None,
        resume_training: bool = False,
        retrain_from_scratch: bool = False,
        show_train_summary: bool = False,
        dataloader_kwargs: Optional[Dict] = None,
    ) -> ConditionalDensityEstimator:
        r"""Return density estimator that approximates the proposal posterior.

        [1] _Fast epsilon-free Inference of Simulation Models with Bayesian Conditional
            Density Estimation_, Papamakarios et al., NeurIPS 2016,
            https://arxiv.org/abs/1605.06376.

        Training is performed with maximum likelihood on samples from the latest round,
        which leads the algorithm to converge to the proposal posterior.

        Args:
            training_batch_size: Training batch size.
            learning_rate: Learning rate for Adam optimizer.
            validation_fraction: The fraction of data to use for validation.
            stop_after_epochs: The number of epochs to wait for improvement on the
                validation set before terminating training.
            max_num_epochs: Maximum number of epochs to run. If reached, we stop
                training even when the validation loss is still decreasing. Otherwise,
                we train until validation loss increases (see also `stop_after_epochs`).
            clip_max_norm: Value at which to clip the total gradient norm in order to
                prevent exploding gradients. Use None for no clipping.
            calibration_kernel: A function to calibrate the loss with respect to the
                simulations `x`. See Lueckmann, Gonçalves et al., NeurIPS 2017.
            resume_training: Can be used in case training time is limited, e.g. on a
                cluster. If `True`, the split between train and validation set, the
                optimizer, the number of epochs, and the best validation log-prob will
                be restored from the last time `.train()` was called.
            retrain_from_scratch: Whether to retrain the conditional density
                estimator for the posterior from scratch each round. Not supported for
                SNPE-A.
            show_train_summary: Whether to print the number of epochs and validation
                loss and leakage after the training.
            dataloader_kwargs: Additional or updated kwargs to be passed to the training
                and validation dataloaders (like, e.g., a collate_fn)

        Returns:
            Density estimator that approximates the distribution $p(\theta|x)$.
        """

        assert not retrain_from_scratch, """Retraining from scratch is not supported in
            SNPE-A yet. The reason for this is that, if we reininitialized the density
            estimator, the z-scoring would change, which would break the posthoc
            correction. This is a pure implementation issue."""

        kwargs = del_entries(
            locals(),
            entries=(
                "self",
                "__class__",
            ),
        )

        # SNPE-A always discards the prior samples.
        kwargs["discard_prior_samples"] = True
        kwargs["force_first_round_loss"] = True

        if len(self._data_round_index) == 0:
            raise RuntimeError(
                "No simulations found. You must call .append_simulations() "
                "before calling .train()."
            )

        self._round = max(self._data_round_index)

        # Always use the specified number of components
        self._build_neural_net = partial(
            self._build_neural_net, num_components=self._num_components
        )

        density_estimator = super().train(**kwargs)

        return density_estimator

    def _get_proposal_mog(
        self,
        proposal: Union["NPE_A_Posterior", MultivariateNormal, MoG, Any],
    ) -> MoG:
        """Extract MoG parameters from a proposal distribution.

        Supports multiple proposal types:
        - NPE_A_Posterior: extracts corrected MoG via get_mog_params()
        - MultivariateNormal: converts to single-component MoG
        - MoG: uses directly
        - Any object with get_mog_params(x) method

        Args:
            proposal: The proposal distribution from the previous round.

        Returns:
            MoG parameters from the proposal.

        Raises:
            ValueError: If NPE_A_Posterior proposal doesn't have default_x set.
            TypeError: If proposal type is not supported.
        """
        if isinstance(proposal, NPE_A_Posterior):
            default_x = proposal.default_x
            if default_x is None:
                raise ValueError(
                    "Proposal posterior must have a default_x set for SNPE-A "
                    "correction. Call posterior.set_default_x(x_o) before using "
                    "as proposal."
                )
            if default_x.shape[0] != 1:
                raise ValueError(
                    f"SNPE-A requires default_x batch size of 1, got "
                    f"{default_x.shape[0]}. SNPE-A only supports single "
                    "observations for correction."
                )
            return proposal.get_mog_params(default_x)

        if isinstance(proposal, MultivariateNormal):
            mean: Tensor = proposal.mean.to(self._device)  # type: ignore[assignment]
            cov: Tensor = proposal.covariance_matrix.to(  # type: ignore[assignment]
                self._device
            )
            return MoG.from_gaussian(mean.unsqueeze(0), cov.unsqueeze(0))

        if isinstance(proposal, MoG):
            return proposal.to(self._device)

        # Case 4: Any object with get_mog_params method
        if hasattr(proposal, "get_mog_params"):
            # Try to get default_x if available
            default_x = getattr(proposal, "default_x", None)
            if default_x is None:
                raise ValueError(
                    "Proposal has get_mog_params() but no default_x set. "
                    "Call proposal.set_default_x(x_o) before using as proposal."
                )
            if default_x.shape[0] != 1:
                raise ValueError(
                    f"SNPE-A requires default_x batch size of 1, got "
                    f"{default_x.shape[0]}."
                )
            mog = proposal.get_mog_params(default_x)
            if not isinstance(mog, MoG):
                raise TypeError(
                    f"Proposal's get_mog_params() must return MoG, "
                    f"got {type(mog).__name__}."
                )
            return mog.to(self._device)

        # Unsupported type
        raise TypeError(
            f"For multi-round SNPE-A, proposal must be one of: NPE_A_Posterior, "
            f"MultivariateNormal, MoG, or an object with get_mog_params() method. "
            f"Got {type(proposal).__name__}. For custom proposals, construct "
            f"NPE_A_Posterior directly with your proposal_mog parameter."
        )

    def _compute_z_scored_prior_mog(
        self, density_estimator: MixtureDensityEstimator
    ) -> Optional[MoG]:
        """Compute the prior as a MoG in z-scored space (if applicable).

        For SNPE-A correction, the prior needs to be in the same coordinate system
        as the density estimator's output. When z-scoring is applied to inputs,
        the density estimator outputs MoG parameters in z-scored space, so the
        prior must also be transformed to z-scored space.

        For uniform priors (BoxUniform), returns None since uniform priors have
        zero precision (infinite covariance) and are handled specially in the
        correction formula.

        Mathematical background:
            For z-score transform: z = (theta - shift) / scale
            If theta ~ N(mu, Sigma), then:
            z ~ N((mu - shift) / scale, Sigma / (scale ⊗ scale))
            where (scale ⊗ scale)_ij = scale_i * scale_j

        Args:
            density_estimator: The MixtureDensityEstimator (to get z-score parameters).

        Returns:
            MoG representation of the z-scored prior for Gaussian priors,
            or None for uniform priors.
        """
        # Uniform priors have zero precision, return None
        if isinstance(self._prior, BoxUniform):
            return None

        if not isinstance(self._prior, MultivariateNormal):
            raise TypeError(
                f"Prior must be MultivariateNormal or BoxUniform, "
                f"got {type(self._prior).__name__}"
            )

        # Get prior parameters
        prior_mean = self._prior.mean
        prior_cov = self._prior.covariance_matrix

        # Apply z-score transform if enabled
        if density_estimator.has_input_transform:
            shift = density_estimator._transform_shift
            scale = density_estimator._transform_scale

            # Z-scored mean: (mu - shift) / scale
            z_mean = (prior_mean - shift) / scale

            # Z-scored covariance: Sigma_z[i,j] = Sigma[i,j] / (scale_i * scale_j)
            scale_outer = scale.unsqueeze(-1) * scale.unsqueeze(-2)
            z_cov = prior_cov / scale_outer
        else:
            z_mean = prior_mean
            z_cov = prior_cov

        # Validate covariance is positive definite
        try:
            torch.linalg.cholesky(z_cov)
        except RuntimeError as e:
            raise ValueError(
                "Z-scored prior covariance is not positive definite. "
                "This may indicate numerical issues with the z-score transform. "
                f"Original error: {e}"
            ) from e

        # Convert to MoG
        return MoG.from_gaussian(z_mean, z_cov)

    def build_posterior(
        self,
        density_estimator: Optional[ConditionalDensityEstimator] = None,
        prior: Optional[Distribution] = None,
        sample_with: Literal["direct"] = "direct",
        **kwargs,
    ) -> NPE_A_Posterior:
        r"""Build posterior from the neural density estimator.

        Returns an NPE_A_Posterior that applies the SNPE-A correction formula:
            p(θ|x) ∝ q(θ|x) × prior(θ) / proposal(θ)

        Note:
            NPE_A only supports `sample_with="direct"`. The corrected posterior is a
            Mixture of Gaussians (MoG) which can be sampled directly and efficiently.
            MCMC, VI, rejection, and importance sampling methods do not provide
            benefits over direct MoG sampling and are therefore not supported.

        Args:
            density_estimator: The density estimator that the posterior is based on.
                If `None`, use the latest neural density estimator that was trained.
            prior: Prior distribution.
            sample_with: Must be "direct". Other sampling methods are not supported.
            **kwargs: Additional arguments passed to NPE_A_Posterior.

        Returns:
            NPE_A_Posterior with the SNPE-A correction applied.

        Raises:
            ValueError: If sample_with is not "direct".
        """
        if sample_with != "direct":
            raise ValueError(
                f"NPE_A only supports sample_with='direct', got '{sample_with}'. "
                "The corrected posterior is a Mixture of Gaussians which can be "
                "sampled directly and efficiently. MCMC, VI, rejection, and "
                "importance sampling do not provide benefits over direct MoG sampling."
            )

        if prior is None:
            assert self._prior is not None, (
                "You did not pass a prior. You have to pass the prior either at "
                "initialization `inference = NPE_A(prior)` or to "
                "`.build_posterior(prior=prior)`."
            )
            prior = self._prior

        # Resolve and validate density estimator
        if density_estimator is None:
            density_estimator = deepcopy(self._neural_net)

        if not isinstance(density_estimator, MixtureDensityEstimator):
            raise TypeError(
                "NPE_A requires MixtureDensityEstimator, "
                f"got {type(density_estimator).__name__}. "
                "Use density_estimator='mdn_snpe_a' when initializing NPE_A."
            )

        # Compute correction parameters
        proposal = self._proposal_roundwise[-1]
        is_first_round = proposal is self._prior or proposal is None

        if is_first_round:
            proposal_mog = None
            prior_mog = None
        else:
            proposal_mog = self._get_proposal_mog(proposal)
            prior_mog = self._compute_z_scored_prior_mog(density_estimator)

        # Build the posterior
        self._posterior = NPE_A_Posterior(
            posterior_estimator=density_estimator,
            prior=prior,
            proposal_mog=proposal_mog,
            prior_mog=prior_mog,
            device=self._device,
            **kwargs,
        )

        return self._posterior

    def _log_prob_proposal_posterior(
        self,
        theta: Tensor,
        x: Tensor,
        masks: Tensor,
        proposal: Optional[Any],
    ) -> Tensor:
        """Return the log-probability of the proposal posterior.

        For SNPE-A this is the same as `self._neural_net.log_prob(theta, x)` in
        `_loss()` to be found in `snpe_base.py`.

        Args:
            theta: Batch of parameters θ.
            x: Batch of data.
            masks: Mask that is True for prior samples in the batch in order to train
                them with prior loss.
            proposal: Proposal distribution.

        Returns: Log-probability of the proposal posterior.
        """
        return self._neural_net.log_prob(theta, x)

__init__(prior=None, density_estimator='mdn_snpe_a', num_components=10, device='cpu', logging_level='WARNING', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize NPE-A [1].

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. Any object with .log_prob()and .sample() (for example, a PyTorch distribution) can be used.	None
density_estimator	Union[Literal['mdn_snpe_a'], ConditionalEstimatorBuilder[ConditionalDensityEstimator]]	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 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() and must return a ConditionalDensityEstimator.	'mdn_snpe_a'
num_components	int	Number of components of the mixture of Gaussians. Note: In multi-round SNPE-A, the number of components grows multiplicatively with each round due to the analytical correction (L components in proposal × K components in density = L*K posterior components).	10
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'WARNING'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during training.	True

Source code in sbi/inference/trainers/npe/npe_a.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    density_estimator: Union[
        Literal["mdn_snpe_a"],
        ConditionalEstimatorBuilder[ConditionalDensityEstimator],
    ] = "mdn_snpe_a",
    num_components: int = 10,
    device: str = "cpu",
    logging_level: Union[int, str] = "WARNING",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize NPE-A [1].

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. Any
            object with `.log_prob()`and `.sample()` (for example, a PyTorch
            distribution) can be used.
        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
            `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()` and must return a `ConditionalDensityEstimator`.
        num_components: Number of components of the mixture of Gaussians.
            Note: In multi-round SNPE-A, the number of components grows
            multiplicatively with each round due to the analytical correction
            (L components in proposal × K components in density = L*K posterior
            components).
        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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during training.
    """

    # Catch invalid inputs.
    if not ((density_estimator == "mdn_snpe_a") or callable(density_estimator)):
        raise TypeError(
            "The `density_estimator` passed to SNPE_A needs to be a "
            "callable or the string 'mdn_snpe_a'!"
        )

    self._num_components = num_components

    # WARNING: sneaky trick ahead. We proxy the parent's `train` here,
    # requiring the signature to have `num_components`, save it for use below, and
    # continue. It's sneaky because we are using the object (self) as a namespace
    # to pass arguments between functions, and that's implicit state management.
    kwargs = del_entries(
        locals(),
        entries=("self", "__class__", "num_components"),
    )
    super().__init__(**kwargs)

build_posterior(density_estimator=None, prior=None, sample_with='direct', **kwargs)

Build posterior from the neural density estimator.

Returns an NPE_A_Posterior that applies the SNPE-A correction formula

p(θ|x) ∝ q(θ|x) × prior(θ) / proposal(θ)

Note

NPE_A only supports sample_with="direct". The corrected posterior is a Mixture of Gaussians (MoG) which can be sampled directly and efficiently. MCMC, VI, rejection, and importance sampling methods do not provide benefits over direct MoG sampling and are therefore not supported.

Parameters:

Name	Type	Description	Default
density_estimator	Optional[ConditionalDensityEstimator]	The density estimator that the posterior is based on. If None, use the latest neural density estimator that was trained.	None
prior	Optional[Distribution]	Prior distribution.	None
sample_with	Literal['direct']	Must be “direct”. Other sampling methods are not supported.	'direct'
**kwargs		Additional arguments passed to NPE_A_Posterior.	{}

Returns:

Type	Description
NPE_A_Posterior	NPE_A_Posterior with the SNPE-A correction applied.

Raises:

Type	Description
ValueError	If sample_with is not “direct”.

Source code in sbi/inference/trainers/npe/npe_a.py

def build_posterior(
    self,
    density_estimator: Optional[ConditionalDensityEstimator] = None,
    prior: Optional[Distribution] = None,
    sample_with: Literal["direct"] = "direct",
    **kwargs,
) -> NPE_A_Posterior:
    r"""Build posterior from the neural density estimator.

    Returns an NPE_A_Posterior that applies the SNPE-A correction formula:
        p(θ|x) ∝ q(θ|x) × prior(θ) / proposal(θ)

    Note:
        NPE_A only supports `sample_with="direct"`. The corrected posterior is a
        Mixture of Gaussians (MoG) which can be sampled directly and efficiently.
        MCMC, VI, rejection, and importance sampling methods do not provide
        benefits over direct MoG sampling and are therefore not supported.

    Args:
        density_estimator: The density estimator that the posterior is based on.
            If `None`, use the latest neural density estimator that was trained.
        prior: Prior distribution.
        sample_with: Must be "direct". Other sampling methods are not supported.
        **kwargs: Additional arguments passed to NPE_A_Posterior.

    Returns:
        NPE_A_Posterior with the SNPE-A correction applied.

    Raises:
        ValueError: If sample_with is not "direct".
    """
    if sample_with != "direct":
        raise ValueError(
            f"NPE_A only supports sample_with='direct', got '{sample_with}'. "
            "The corrected posterior is a Mixture of Gaussians which can be "
            "sampled directly and efficiently. MCMC, VI, rejection, and "
            "importance sampling do not provide benefits over direct MoG sampling."
        )

    if prior is None:
        assert self._prior is not None, (
            "You did not pass a prior. You have to pass the prior either at "
            "initialization `inference = NPE_A(prior)` or to "
            "`.build_posterior(prior=prior)`."
        )
        prior = self._prior

    # Resolve and validate density estimator
    if density_estimator is None:
        density_estimator = deepcopy(self._neural_net)

    if not isinstance(density_estimator, MixtureDensityEstimator):
        raise TypeError(
            "NPE_A requires MixtureDensityEstimator, "
            f"got {type(density_estimator).__name__}. "
            "Use density_estimator='mdn_snpe_a' when initializing NPE_A."
        )

    # Compute correction parameters
    proposal = self._proposal_roundwise[-1]
    is_first_round = proposal is self._prior or proposal is None

    if is_first_round:
        proposal_mog = None
        prior_mog = None
    else:
        proposal_mog = self._get_proposal_mog(proposal)
        prior_mog = self._compute_z_scored_prior_mog(density_estimator)

    # Build the posterior
    self._posterior = NPE_A_Posterior(
        posterior_estimator=density_estimator,
        prior=prior,
        proposal_mog=proposal_mog,
        prior_mog=prior_mog,
        device=self._device,
        **kwargs,
    )

    return self._posterior

train(training_batch_size=200, learning_rate=0.0005, validation_fraction=0.1, stop_after_epochs=20, max_num_epochs=2 ** 31 - 1, clip_max_norm=5.0, calibration_kernel=None, resume_training=False, retrain_from_scratch=False, show_train_summary=False, dataloader_kwargs=None)

Return density estimator that approximates the proposal posterior.

[1] Fast epsilon-free Inference of Simulation Models with Bayesian Conditional Density Estimation, Papamakarios et al., NeurIPS 2016, https://arxiv.org/abs/1605.06376.

Training is performed with maximum likelihood on samples from the latest round, which leads the algorithm to converge to the proposal posterior.

Parameters:

Name	Type	Description	Default
training_batch_size	int	Training batch size.	200
learning_rate	float	Learning rate for Adam optimizer.	0.0005
validation_fraction	float	The fraction of data to use for validation.	0.1
stop_after_epochs	int	The number of epochs to wait for improvement on the validation set before terminating training.	20
max_num_epochs	int	Maximum number of epochs to run. If reached, we stop training even when the validation loss is still decreasing. Otherwise, we train until validation loss increases (see also stop_after_epochs).	2 ** 31 - 1
clip_max_norm	Optional[float]	Value at which to clip the total gradient norm in order to prevent exploding gradients. Use None for no clipping.	5.0
calibration_kernel	Optional[Callable]	A function to calibrate the loss with respect to the simulations x. See Lueckmann, Gonçalves et al., NeurIPS 2017.	None
resume_training	bool	Can be used in case training time is limited, e.g. on a cluster. If True, the split between train and validation set, the optimizer, the number of epochs, and the best validation log-prob will be restored from the last time .train() was called.	False
retrain_from_scratch	bool	Whether to retrain the conditional density estimator for the posterior from scratch each round. Not supported for SNPE-A.	False
show_train_summary	bool	Whether to print the number of epochs and validation loss and leakage after the training.	False
dataloader_kwargs	Optional[Dict]	Additional or updated kwargs to be passed to the training and validation dataloaders (like, e.g., a collate_fn)	None

Returns:

Type	Description
ConditionalDensityEstimator	Density estimator that approximates the distribution \(p(\theta|x)\).

Source code in sbi/inference/trainers/npe/npe_a.py

def train(
    self,
    training_batch_size: int = 200,
    learning_rate: float = 5e-4,
    validation_fraction: float = 0.1,
    stop_after_epochs: int = 20,
    max_num_epochs: int = 2**31 - 1,
    clip_max_norm: Optional[float] = 5.0,
    calibration_kernel: Optional[Callable] = None,
    resume_training: bool = False,
    retrain_from_scratch: bool = False,
    show_train_summary: bool = False,
    dataloader_kwargs: Optional[Dict] = None,
) -> ConditionalDensityEstimator:
    r"""Return density estimator that approximates the proposal posterior.

    [1] _Fast epsilon-free Inference of Simulation Models with Bayesian Conditional
        Density Estimation_, Papamakarios et al., NeurIPS 2016,
        https://arxiv.org/abs/1605.06376.

    Training is performed with maximum likelihood on samples from the latest round,
    which leads the algorithm to converge to the proposal posterior.

    Args:
        training_batch_size: Training batch size.
        learning_rate: Learning rate for Adam optimizer.
        validation_fraction: The fraction of data to use for validation.
        stop_after_epochs: The number of epochs to wait for improvement on the
            validation set before terminating training.
        max_num_epochs: Maximum number of epochs to run. If reached, we stop
            training even when the validation loss is still decreasing. Otherwise,
            we train until validation loss increases (see also `stop_after_epochs`).
        clip_max_norm: Value at which to clip the total gradient norm in order to
            prevent exploding gradients. Use None for no clipping.
        calibration_kernel: A function to calibrate the loss with respect to the
            simulations `x`. See Lueckmann, Gonçalves et al., NeurIPS 2017.
        resume_training: Can be used in case training time is limited, e.g. on a
            cluster. If `True`, the split between train and validation set, the
            optimizer, the number of epochs, and the best validation log-prob will
            be restored from the last time `.train()` was called.
        retrain_from_scratch: Whether to retrain the conditional density
            estimator for the posterior from scratch each round. Not supported for
            SNPE-A.
        show_train_summary: Whether to print the number of epochs and validation
            loss and leakage after the training.
        dataloader_kwargs: Additional or updated kwargs to be passed to the training
            and validation dataloaders (like, e.g., a collate_fn)

    Returns:
        Density estimator that approximates the distribution $p(\theta|x)$.
    """

    assert not retrain_from_scratch, """Retraining from scratch is not supported in
        SNPE-A yet. The reason for this is that, if we reininitialized the density
        estimator, the z-scoring would change, which would break the posthoc
        correction. This is a pure implementation issue."""

    kwargs = del_entries(
        locals(),
        entries=(
            "self",
            "__class__",
        ),
    )

    # SNPE-A always discards the prior samples.
    kwargs["discard_prior_samples"] = True
    kwargs["force_first_round_loss"] = True

    if len(self._data_round_index) == 0:
        raise RuntimeError(
            "No simulations found. You must call .append_simulations() "
            "before calling .train()."
        )

    self._round = max(self._data_round_index)

    # Always use the specified number of components
    self._build_neural_net = partial(
        self._build_neural_net, num_components=self._num_components
    )

    density_estimator = super().train(**kwargs)

    return density_estimator

NPE_B

Bases: PosteriorEstimatorTrainer

Neural Posterior Estimation algorithm (NPE-B) as in Lueckmann et al. (2017) [1].

NPE-B (also known as SNPE-B) trains a neural network to directly approximate the posterior \(p(\theta|x)\) using an importance-weighted loss. Unlike NPE-A, this importance weighting ensures convergence to the true posterior in multi-round inference, and it is not limited to Gaussian proposals. NPE-B can use flexible density estimators like normalizing flows.

For single-round inference, NPE-A, NPE-B, and NPE-C are equivalent and use plain NLL loss.

[1] Flexible statistical inference for mechanistic models of neural dynamics, Lueckmann, Gonçalves et al., NeurIPS 2017. https://arxiv.org/abs/1711.01861

Example:

::

import torch
from sbi.inference import NPE_B
from sbi.utils import BoxUniform

# 1. Setup simulator, prior, and observation
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
x_o = torch.randn(1, 3)  # Observed data

def simulator(theta):
    return theta + torch.randn_like(theta) * 0.1

# 2. Multi-round inference
inference = NPE_B(prior=prior)
proposal = prior

for round_idx in range(5):
    theta = proposal.sample((100,))
    x = simulator(theta)
    density_estimator = inference.append_simulations(theta, x).train()
    posterior = inference.build_posterior(density_estimator)
    proposal = posterior.set_default_x(x_o)

# 3. Sample from final posterior
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/npe/npe_b.py

class NPE_B(PosteriorEstimatorTrainer):
    r"""Neural Posterior Estimation algorithm (NPE-B) as in Lueckmann et al. (2017) [1].

    NPE-B (also known as SNPE-B) trains a neural network to directly approximate
    the posterior $p(\theta|x)$ using an importance-weighted loss. Unlike NPE-A,
    this importance weighting ensures convergence to the true posterior in multi-round
    inference, and it is not limited to Gaussian proposals. NPE-B can use flexible
    density estimators like normalizing flows.

    For single-round inference, NPE-A, NPE-B, and NPE-C are equivalent and use
    plain NLL loss.

    [1] *Flexible statistical inference for mechanistic models of neural
        dynamics*, Lueckmann, Gonçalves et al., NeurIPS 2017.
        https://arxiv.org/abs/1711.01861

    Example:
    --------

    ::

        import torch
        from sbi.inference import NPE_B
        from sbi.utils import BoxUniform

        # 1. Setup simulator, prior, and observation
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        x_o = torch.randn(1, 3)  # Observed data

        def simulator(theta):
            return theta + torch.randn_like(theta) * 0.1

        # 2. Multi-round inference
        inference = NPE_B(prior=prior)
        proposal = prior

        for round_idx in range(5):
            theta = proposal.sample((100,))
            x = simulator(theta)
            density_estimator = inference.append_simulations(theta, x).train()
            posterior = inference.build_posterior(density_estimator)
            proposal = posterior.set_default_x(x_o)

        # 3. Sample from final posterior
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        density_estimator: Union[
            Literal["nsf", "maf", "mdn", "made"],
            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
        ] = "maf",
        device: str = "cpu",
        logging_level: Union[int, str] = "WARNING",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize NPE-B.

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them.
            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
                `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()` 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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during training.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

    def _log_prob_proposal_posterior(
        self,
        theta: Tensor,
        x: Tensor,
        masks: Tensor,
        proposal: Optional[Any],
    ) -> Tensor:
        """
        Return importance-weighted log probability (Lueckmann, Goncalves et al., 2017).

        Args:
            theta: Batch of parameters θ.
            x: Batch of data.
            masks: Mask that is True for prior samples in the batch in order to train
                them with prior loss.
            proposal: Proposal distribution.

        Returns:
            Importance-weighted log-probability of the proposal posterior.
        """

        # Evaluate prior
        # we accept prior log prob to be -Inf at theta
        # meaning that theta is out of the prior range (the weight is thus 0)
        utils.assert_not_nan_or_plus_inf(
            self._prior.log_prob(theta), "prior log probs of proposal samples"
        )
        prior = torch.exp(self._prior.log_prob(theta))

        # Evaluate proposal
        # (as theta comes from prior and proposal from previous rounds,
        # the last proposal is actually a mixture of the prior
        # and of all the previous proposals with coefficients representing
        # the proportion of the new theta added at each round)
        prop = torch.zeros(self._round + 1, device=theta.device)
        nb_samples = 0  # total number of theta from all the rounds

        for k in range(self._round + 1):
            nb_samples += self._theta_roundwise[k].size(0)
            # the number of new theta sampled in the round k
            prop[k] = self._theta_roundwise[k].size(0)

        prop /= nb_samples
        log_prop = torch.log(prop).repeat(theta.size(0), 1)

        log_previous_proposals = torch.zeros(
            (theta.size(0), self._round + 1), device=theta.device
        )
        for k, density in enumerate(self._proposal_roundwise):
            # we accept the k th proposal log prob to be -Inf at theta
            # meaning that theta is out of the k th proposal range
            log_previous_proposals[:, k] = density.log_prob(theta)
            utils.assert_not_nan_or_plus_inf(
                log_previous_proposals[:, k], "proposal log probs of proposal samples"
            )

        log_proposal = torch.logsumexp(log_prop + log_previous_proposals, dim=1)
        proposal = torch.exp(log_proposal)

        # Construct the importance weights and normalize them
        importance_weights = prior / proposal
        importance_weights /= importance_weights.sum()

        theta = reshape_to_sample_batch_event(theta, theta.shape[1:])
        # Reshape the density estimator log probs
        # from (sample_shape, batch_shape) to (batch_shape)
        posterior_log_probs = self._neural_net.log_prob(theta, x).squeeze(dim=0)

        return importance_weights * posterior_log_probs

__init__(prior=None, density_estimator='maf', device='cpu', logging_level='WARNING', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize NPE-B.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them.	None
density_estimator	Union[Literal['nsf', 'maf', 'mdn', 'made'], ConditionalEstimatorBuilder[ConditionalDensityEstimator]]	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 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() and must return a ConditionalDensityEstimator.	'maf'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'WARNING'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during training.	True

Source code in sbi/inference/trainers/npe/npe_b.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    density_estimator: Union[
        Literal["nsf", "maf", "mdn", "made"],
        ConditionalEstimatorBuilder[ConditionalDensityEstimator],
    ] = "maf",
    device: str = "cpu",
    logging_level: Union[int, str] = "WARNING",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize NPE-B.

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them.
        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
            `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()` 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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during training.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

NPE_C

Bases: PosteriorEstimatorTrainer

Neural Posterior Estimation algorithm (NPE-C) as in Greenberg et al. (2019) [1].

NPE-C (also known as APT - Automatic Posterior Transformation, aka SNPE-C) trains a neural network over multiple rounds to directly approximate the posterior for a specific observation x_o. In the first round, NPE-C is equivalent to other NPE methods and is fully amortized (direct inference for any new observation). After the first round, NPE-C automatically selects between two loss variants depending on the chosen density estimator: the non-atomic loss (for Mixture of Gaussians) which is stable and avoids leakage, or the atomic loss (for flows) which is more flexible but may suffer from leakage issues.

For single-round inference, NPE-A, NPE-B, and NPE-C are equivalent and use plain NLL loss.

[1] Automatic Posterior Transformation for Likelihood-free Inference, Greenberg et al., ICML 2019, https://arxiv.org/abs/1905.07488.

Example:

::

import torch
from sbi.inference import NPE_C
from sbi.utils import BoxUniform

# 1. Setup simulator, prior, and observation
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
x_o = torch.randn(1, 3)  # Observed data

def simulator(theta):
    return theta + torch.randn_like(theta) * 0.1

# 2. Multi-round inference
inference = NPE_C(prior=prior)
proposal = prior

for round_idx in range(5):
    theta = proposal.sample((100,))
    x = simulator(theta)
    density_estimator = inference.append_simulations(theta, x).train()
    posterior = inference.build_posterior(density_estimator)
    proposal = posterior.set_default_x(x_o)

# 3. Sample from final posterior
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/npe/npe_c.py

class NPE_C(PosteriorEstimatorTrainer):
    r"""Neural Posterior Estimation algorithm (NPE-C) as in Greenberg et al. (2019) [1].

    NPE-C (also known as APT - Automatic Posterior Transformation, aka SNPE-C) trains
    a neural network over multiple rounds to directly approximate the posterior for a
    specific observation x_o. In the first round, NPE-C is equivalent to other NPE
    methods and is fully amortized (direct inference for any new observation). After
    the first round, NPE-C automatically selects between two loss variants depending
    on the chosen density estimator: the non-atomic loss (for Mixture of Gaussians)
    which is stable and avoids leakage, or the atomic loss (for flows) which is more
    flexible but may suffer from leakage issues.

    For single-round inference, NPE-A, NPE-B, and NPE-C are equivalent and use
    plain NLL loss.

    [1] *Automatic Posterior Transformation for Likelihood-free Inference*,
        Greenberg et al., ICML 2019, https://arxiv.org/abs/1905.07488.

    Example:
    --------

    ::

        import torch
        from sbi.inference import NPE_C
        from sbi.utils import BoxUniform

        # 1. Setup simulator, prior, and observation
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        x_o = torch.randn(1, 3)  # Observed data

        def simulator(theta):
            return theta + torch.randn_like(theta) * 0.1

        # 2. Multi-round inference
        inference = NPE_C(prior=prior)
        proposal = prior

        for round_idx in range(5):
            theta = proposal.sample((100,))
            x = simulator(theta)
            density_estimator = inference.append_simulations(theta, x).train()
            posterior = inference.build_posterior(density_estimator)
            proposal = posterior.set_default_x(x_o)

        # 3. Sample from final posterior
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        density_estimator: Union[
            Literal["nsf", "maf", "mdn", "made"],
            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
        ] = "maf",
        device: str = "cpu",
        logging_level: Union[int, str] = "WARNING",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize NPE-C.

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them.
            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
                `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()` 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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during training.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

    def train(
        self,
        num_atoms: int = 10,
        training_batch_size: int = 200,
        learning_rate: float = 5e-4,
        validation_fraction: float = 0.1,
        stop_after_epochs: int = 20,
        max_num_epochs: int = 2**31 - 1,
        clip_max_norm: Optional[float] = 5.0,
        calibration_kernel: Optional[Callable] = None,
        resume_training: bool = False,
        force_first_round_loss: bool = False,
        discard_prior_samples: bool = False,
        use_combined_loss: bool = False,
        retrain_from_scratch: bool = False,
        show_train_summary: bool = False,
        dataloader_kwargs: Optional[Dict] = None,
    ) -> ConditionalDensityEstimator:
        r"""Return density estimator that approximates the distribution $p(\theta|x)$.

        Args:
            num_atoms: Number of atoms to use for classification.
            training_batch_size: Training batch size.
            learning_rate: Learning rate for Adam optimizer.
            validation_fraction: The fraction of data to use for validation.
            stop_after_epochs: The number of epochs to wait for improvement on the
                validation set before terminating training.
            max_num_epochs: Maximum number of epochs to run. If reached, we stop
                training even when the validation loss is still decreasing. Otherwise,
                we train until validation loss increases (see also
                ``stop_after_epochs``).
            clip_max_norm: Value at which to clip the total gradient norm in order to
                prevent exploding gradients. Use None for no clipping.
            calibration_kernel: A function to calibrate the loss with respect to the
                simulations ``x``. See Lueckmann, Gonçalves et al., NeurIPS 2017.
            resume_training: Can be used in case training time is limited, e.g. on a
                cluster. If ``True``, the split between train and validation set, the
                optimizer, the number of epochs, and the best validation log-prob will
                be restored from the last time ``.train()`` was called.
            force_first_round_loss: If ``True``, train with maximum likelihood,
                i.e., potentially ignoring the correction for using a proposal
                distribution different from the prior.
            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                from the prior. Training may be sped up by ignoring such less targeted
                samples.
            use_combined_loss: Whether to train the neural net also on prior samples
                using maximum likelihood in addition to training it on all samples using
                atomic loss. The extra MLE loss helps prevent density leaking with
                bounded priors.
            retrain_from_scratch: Whether to retrain the conditional density
                estimator for the posterior from scratch each round.
            show_train_summary: Whether to print the number of epochs and validation
                loss and leakage after the training.
            dataloader_kwargs: Additional or updated kwargs to be passed to the training
                and validation dataloaders (like, e.g., a collate_fn)

        Returns:
            Density estimator that approximates the distribution $p(\theta|x)$.
        """

        if len(self._data_round_index) == 0:
            raise RuntimeError(
                "No simulations found. You must call .append_simulations() "
                "before calling .train()."
            )
        # WARNING: sneaky trick ahead. We proxy the parent's `train` here,
        # requiring the signature to have `num_atoms`, save it for use below, and
        # continue. It's sneaky because we are using the object (self) as a namespace
        # to pass arguments between functions, and that's implicit state management.
        self._num_atoms = num_atoms
        self._use_combined_loss = use_combined_loss
        kwargs = del_entries(
            locals(),
            entries=("self", "__class__", "num_atoms", "use_combined_loss"),
        )

        self._round = max(self._data_round_index)

        if self._round > 0:
            # Set the proposal to the last proposal that was passed by the user. For
            # atomic SNPE, it does not matter what the proposal is. For non-atomic
            # SNPE, we only use the latest data that was passed, i.e. the one from the
            # last proposal.
            proposal = self._proposal_roundwise[-1]
            self.use_non_atomic_loss = (
                isinstance(proposal, DirectPosterior)
                and isinstance(proposal.posterior_estimator, MixtureDensityEstimator)
                and isinstance(self._neural_net, MixtureDensityEstimator)
                and check_dist_class(
                    self._prior, class_to_check=(Uniform, MultivariateNormal)
                )[0]
            )

            algorithm = "non-atomic" if self.use_non_atomic_loss else "atomic"
            print(f"Using SNPE-C with {algorithm} loss")

            if self.use_non_atomic_loss:
                # Take care of z-scoring, pre-compute and store prior terms.
                self._set_state_for_mog_proposal()

        return super().train(**kwargs)

    def _set_state_for_mog_proposal(self) -> None:
        """Set state variables that are used at each training step of non-atomic SNPE-C.

        Three things are computed:
        1) Check if z-scoring was requested. We check if the MixtureDensityEstimator
            has an input transform enabled via the has_input_transform property.
        2) Define a (potentially standardized) prior. It's standardized if z-scoring
            had been requested.
        3) Compute (Precision * mean) for the prior. This quantity is used at every
            training step if the prior is Gaussian.
        """
        # Check if z-scoring is enabled on the MixtureDensityEstimator
        assert isinstance(self._neural_net, MixtureDensityEstimator)
        self.z_score_theta = self._neural_net.has_input_transform

        self._set_maybe_z_scored_prior()

        if isinstance(self._maybe_z_scored_prior, MultivariateNormal):
            self.prec_m_prod_prior = torch.mv(
                self._maybe_z_scored_prior.precision_matrix,  # type: ignore
                self._maybe_z_scored_prior.loc,  # type: ignore
            )

    def _set_maybe_z_scored_prior(self) -> None:
        r"""Compute and store potentially standardized prior (if z-scoring was done).

        The proposal posterior is:
        $pp(\theta|x) = 1/Z * q(\theta|x) * prop(\theta) / p(\theta)$

        Let's denote z-scored theta by `a`: a = (theta - mean) / std
        Then pp'(a|x) = 1/Z_2 * q'(a|x) * prop'(a) / p'(a)$

        The ' indicates that the evaluation occurs in standardized space. The constant
        scaling factor has been absorbed into Z_2.
        From the above equation, we see that we need to evaluate the prior **in
        standardized space**. We build the standardized prior in this function.

        The standardize transform that is applied to the samples theta does not use
        the exact prior mean and std (due to implementation issues). Hence, the z-scored
        prior will not be exactly have mean=0 and std=1.
        """

        if self.z_score_theta:
            # Get z-score parameters from the MixtureDensityEstimator
            # The transform is: z = (theta - shift) / scale
            # where shift = mean (estimated from samples) and scale = std (estimated)
            assert isinstance(self._neural_net, MixtureDensityEstimator)
            shift = self._neural_net._transform_shift
            scale = self._neural_net._transform_scale

            # The MixtureDensityEstimator uses: z = (theta - shift) / scale
            # where shift = mean and scale = std (estimated from training data)
            estim_prior_mean = shift
            estim_prior_std = scale

            # Compute the discrepancy of the true prior mean and std and the mean and
            # std that was empirically estimated from samples.
            # N(theta|m,s) = N((theta-m_e)/s_e|(m-m_e)/s_e, s/s_e)
            # Above: m,s are true prior mean and std. m_e,s_e are estimated prior mean
            # and std (estimated from samples and used to build standardize transform).
            almost_zero_mean = (self._prior.mean - estim_prior_mean) / estim_prior_std
            almost_one_std = torch.sqrt(self._prior.variance) / estim_prior_std

            if isinstance(self._prior, MultivariateNormal):
                self._maybe_z_scored_prior = MultivariateNormal(
                    almost_zero_mean, torch.diag(almost_one_std)
                )
            else:
                range_ = torch.sqrt(almost_one_std * 3.0)
                self._maybe_z_scored_prior = BoxUniform(
                    almost_zero_mean - range_, almost_zero_mean + range_
                )
        else:
            self._maybe_z_scored_prior = self._prior

    def _log_prob_proposal_posterior(
        self,
        theta: Tensor,
        x: Tensor,
        masks: Tensor,
        proposal: DirectPosterior,
    ) -> Tensor:
        """Return the log-probability of the proposal posterior.

        If the proposal is a MoG, the density estimator is a MoG, and the prior is
        either Gaussian or uniform, we use non-atomic loss. Else, use atomic loss (which
        suffers from leakage).

        Args:
            theta: Batch of parameters θ.
            x: Batch of data.
            masks: Mask that is True for prior samples in the batch in order to train
                them with prior loss.
            proposal: Proposal distribution.

        Returns: Log-probability of the proposal posterior.
        """

        if self.use_non_atomic_loss:
            if not isinstance(self._neural_net, MixtureDensityEstimator):
                raise ValueError(
                    "The density estimator must be a MixtureDensityEstimator "
                    "for non-atomic loss."
                )

            return self._log_prob_proposal_posterior_mog(theta, x, proposal)
        else:
            if not hasattr(self._neural_net, "log_prob"):
                raise ValueError(
                    "The neural estimator must have a log_prob method, for\
                                 atomic loss. It should at best follow the \
                                 sbi.neural_nets 'DensityEstiamtor' interface."
                )
            return self._log_prob_proposal_posterior_atomic(theta, x, masks)

    def _log_prob_proposal_posterior_atomic(
        self, theta: Tensor, x: Tensor, masks: Tensor
    ):
        """Return log probability of the proposal posterior for atomic proposals.

        We have two main options when evaluating the proposal posterior.
            (1) Generate atoms from the proposal prior.
            (2) Generate atoms from a more targeted distribution, such as the most
                recent posterior.
        If we choose the latter, it is likely beneficial not to do this in the first
        round, since we would be sampling from a randomly-initialized neural density
        estimator.

        Args:
            theta: Batch of parameters θ.
            x: Batch of data.
            masks: Mask that is True for prior samples in the batch in order to train
                them with prior loss.

        Returns:
            Log-probability of the proposal posterior.
        """
        batch_size = theta.shape[0]

        num_atoms = int(
            clamp_and_warn("num_atoms", self._num_atoms, min_val=2, max_val=batch_size)
        )

        # Each set of parameter atoms is evaluated using the same x,
        # so we repeat rows of the data x, e.g. [1, 2] -> [1, 1, 2, 2]
        repeated_x = repeat_rows(x, num_atoms)

        # To generate the full set of atoms for a given item in the batch,
        # we sample without replacement num_atoms - 1 times from the rest
        # of the theta in the batch.
        probs = ones(batch_size, batch_size) * (1 - eye(batch_size)) / (batch_size - 1)

        choices = torch.multinomial(probs, num_samples=num_atoms - 1, replacement=False)
        contrasting_theta = theta[choices]

        # We can now create our sets of atoms from the contrasting parameter sets
        # we have generated.
        atomic_theta = torch.cat((theta[:, None, :], contrasting_theta), dim=1).reshape(
            batch_size * num_atoms, -1
        )

        # Get (batch_size * num_atoms) log prob prior evals.
        log_prob_prior = self._prior.log_prob(atomic_theta)
        log_prob_prior = log_prob_prior.reshape(batch_size, num_atoms)
        assert_all_finite(log_prob_prior, "prior eval")

        # Evaluate large batch giving (batch_size * num_atoms) log prob posterior evals.
        atomic_theta = reshape_to_sample_batch_event(
            atomic_theta, atomic_theta.shape[1:]
        )
        repeated_x = reshape_to_batch_event(
            repeated_x, self._neural_net.condition_shape
        )
        log_prob_posterior = self._neural_net.log_prob(atomic_theta, repeated_x)
        assert_all_finite(log_prob_posterior, "posterior eval")
        log_prob_posterior = log_prob_posterior.reshape(batch_size, num_atoms)

        # Compute unnormalized proposal posterior.
        unnormalized_log_prob = log_prob_posterior - log_prob_prior

        # Normalize proposal posterior across discrete set of atoms.
        log_prob_proposal_posterior = unnormalized_log_prob[:, 0] - torch.logsumexp(
            unnormalized_log_prob, dim=-1
        )
        assert_all_finite(log_prob_proposal_posterior, "proposal posterior eval")

        # XXX This evaluates the posterior on _all_ prior samples
        if self._use_combined_loss:
            theta = reshape_to_sample_batch_event(theta, self._neural_net.input_shape)
            x = reshape_to_batch_event(x, self._neural_net.condition_shape)
            log_prob_posterior_non_atomic = self._neural_net.log_prob(theta, x)
            # squeeze to remove sample dimension, which is always one during the loss
            # evaluation of `SNPE_C` (because we have one theta vector per x vector).
            log_prob_posterior_non_atomic = log_prob_posterior_non_atomic.squeeze(dim=0)
            masks = masks.reshape(-1)
            log_prob_proposal_posterior = (
                masks * log_prob_posterior_non_atomic + log_prob_proposal_posterior
            )

        return log_prob_proposal_posterior

    def _log_prob_proposal_posterior_mog(
        self, theta: Tensor, x: Tensor, proposal: DirectPosterior
    ) -> Tensor:
        """Return log-probability of the proposal posterior for MoG proposal.

        For MoG proposals and MoG density estimators, this can be done in closed form
        and does not require atomic loss (i.e. there will be no leakage issues).

        Notation:

        m are mean vectors.
        prec are precision matrices.
        cov are covariance matrices.

        _p at the end indicates that it is the proposal.
        _d indicates that it is the density estimator.
        _pp indicates the proposal posterior.

        All tensors will have shapes (batch_dim, num_components, ...)

        Args:
            theta: Batch of parameters θ.
            x: Batch of data.
            proposal: Proposal distribution.

        Returns:
            Log-probability of the proposal posterior.
        """
        # Get the proposal MoG at the default_x
        assert isinstance(proposal.posterior_estimator, MixtureDensityEstimator)
        assert proposal.default_x is not None, "Proposal must have default_x set"
        mog_p = proposal.posterior_estimator.get_uncorrected_mog(proposal.default_x)
        norm_logits_p = mog_p.log_weights  # Already normalized
        m_p = mog_p.means
        prec_p = mog_p.precisions

        # Get the density estimator MoG at the training data x
        assert isinstance(self._neural_net, MixtureDensityEstimator)
        mog_d = self._neural_net.get_uncorrected_mog(x)
        norm_logits_d = mog_d.log_weights  # Already normalized
        m_d = mog_d.means
        prec_d = mog_d.precisions

        # z-score theta if z-scoring was requested.
        theta = self._maybe_z_score_theta(theta)

        # Compute the MoG parameters of the proposal posterior.
        (
            logits_pp,
            m_pp,
            prec_pp,
            cov_pp,
        ) = self._automatic_posterior_transformation(
            norm_logits_p, m_p, prec_p, norm_logits_d, m_d, prec_d
        )

        # Create MoG for proposal posterior and compute log_prob
        # We need precision_factors for MoG, compute via Cholesky
        precf_pp = torch.linalg.cholesky(prec_pp, upper=True)
        mog_pp = MoG(
            logits=logits_pp,
            means=m_pp,
            precisions=prec_pp,
            precision_factors=precf_pp,
        )

        # Compute the log_prob of theta under the product.
        log_prob_proposal_posterior = mog_pp.log_prob(theta)
        assert_all_finite(
            log_prob_proposal_posterior,
            """the evaluation of the MoG proposal posterior. This is likely due to a
            numerical instability in the training procedure. Please create an issue on
            Github.""",
        )

        return log_prob_proposal_posterior

    def _automatic_posterior_transformation(
        self,
        logits_p: Tensor,
        means_p: Tensor,
        precisions_p: Tensor,
        logits_d: Tensor,
        means_d: Tensor,
        precisions_d: Tensor,
    ):
        r"""Returns the MoG parameters of the proposal posterior.

        The proposal posterior is:
        $pp(\theta|x) = 1/Z * q(\theta|x) * prop(\theta) / p(\theta)$
        In words: proposal posterior = posterior estimate * proposal / prior.

        If the posterior estimate and the proposal are MoG and the prior is either
        Gaussian or uniform, we can solve this in closed-form. The is implemented in
        this function.

        This function implements Appendix A1 from Greenberg et al. 2019.

        We have to build L*K components. How do we do this?
        Example: proposal has two components, density estimator has three components.
        Let's call the two components of the proposal i,j and the three components
        of the density estimator x,y,z. We have to multiply every component of the
        proposal with every component of the density estimator. So, what we do is:
        1) for the proposal, build: i,i,i,j,j,j. Done with torch.repeat_interleave()
        2) for the density estimator, build: x,y,z,x,y,z. Done with torch.repeat()
        3) Multiply them with simple matrix operations.

        Args:
            logits_p: Component weight of each Gaussian of the proposal.
            means_p: Mean of each Gaussian of the proposal.
            precisions_p: Precision matrix of each Gaussian of the proposal.
            logits_d: Component weight for each Gaussian of the density estimator.
            means_d: Mean of each Gaussian of the density estimator.
            precisions_d: Precision matrix of each Gaussian of the density estimator.

        Returns: (Component weight, mean, precision matrix, covariance matrix) of each
            Gaussian of the proposal posterior. Has L*K terms (proposal has L terms,
            density estimator has K terms).
        """

        precisions_pp, covariances_pp = self._precisions_proposal_posterior(
            precisions_p, precisions_d
        )

        means_pp = self._means_proposal_posterior(
            covariances_pp, means_p, precisions_p, means_d, precisions_d
        )

        logits_pp = self._logits_proposal_posterior(
            means_pp,
            precisions_pp,
            covariances_pp,
            logits_p,
            means_p,
            precisions_p,
            logits_d,
            means_d,
            precisions_d,
        )

        return logits_pp, means_pp, precisions_pp, covariances_pp

    def _precisions_proposal_posterior(
        self, precisions_p: Tensor, precisions_d: Tensor
    ):
        """Return the precisions and covariances of the proposal posterior.

        Args:
            precisions_p: Precision matrices of the proposal distribution.
            precisions_d: Precision matrices of the density estimator.

        Returns: (Precisions, Covariances) of the proposal posterior. L*K terms.
        """

        num_comps_p = precisions_p.shape[1]
        num_comps_d = precisions_d.shape[1]

        precisions_p_rep = precisions_p.repeat_interleave(num_comps_d, dim=1)
        precisions_d_rep = precisions_d.repeat(1, num_comps_p, 1, 1)

        precisions_pp = precisions_p_rep + precisions_d_rep
        if isinstance(self._maybe_z_scored_prior, MultivariateNormal):
            precisions_pp -= self._maybe_z_scored_prior.precision_matrix

        covariances_pp = torch.inverse(precisions_pp)

        return precisions_pp, covariances_pp

    def _means_proposal_posterior(
        self,
        covariances_pp: Tensor,
        means_p: Tensor,
        precisions_p: Tensor,
        means_d: Tensor,
        precisions_d: Tensor,
    ):
        """Return the means of the proposal posterior.

        means_pp = C_ix * (P_i * m_i + P_x * m_x - P_o * m_o).

        Args:
            covariances_pp: Covariance matrices of the proposal posterior.
            means_p: Means of the proposal distribution.
            precisions_p: Precision matrices of the proposal distribution.
            means_d: Means of the density estimator.
            precisions_d: Precision matrices of the density estimator.

        Returns: Means of the proposal posterior. L*K terms.
        """

        num_comps_p = precisions_p.shape[1]
        num_comps_d = precisions_d.shape[1]

        # First, compute the product P_i * m_i and P_j * m_j
        prec_m_prod_p = batched_mixture_mv(precisions_p, means_p)
        prec_m_prod_d = batched_mixture_mv(precisions_d, means_d)

        # Repeat them to allow for matrix operations: same trick as for the precisions.
        prec_m_prod_p_rep = prec_m_prod_p.repeat_interleave(num_comps_d, dim=1)
        prec_m_prod_d_rep = prec_m_prod_d.repeat(1, num_comps_p, 1)

        # Means = C_ij * (P_i * m_i + P_x * m_x - P_o * m_o).
        summed_cov_m_prod_rep = prec_m_prod_p_rep + prec_m_prod_d_rep
        if isinstance(self._maybe_z_scored_prior, MultivariateNormal):
            summed_cov_m_prod_rep -= self.prec_m_prod_prior

        means_pp = batched_mixture_mv(covariances_pp, summed_cov_m_prod_rep)

        return means_pp

    def _maybe_z_score_theta(self, theta: Tensor) -> Tensor:
        """Return potentially standardized theta if z-scoring was requested."""

        if self.z_score_theta:
            assert isinstance(self._neural_net, MixtureDensityEstimator)
            theta = self._neural_net._transform_input(theta)

        return theta

    @staticmethod
    def _logits_proposal_posterior(
        means_pp: Tensor,
        precisions_pp: Tensor,
        covariances_pp: Tensor,
        logits_p: Tensor,
        means_p: Tensor,
        precisions_p: Tensor,
        logits_d: Tensor,
        means_d: Tensor,
        precisions_d: Tensor,
    ):
        """Return the component weights (i.e. logits) of the proposal posterior.

        Args:
            means_pp: Means of the proposal posterior.
            precisions_pp: Precision matrices of the proposal posterior.
            covariances_pp: Covariance matrices of the proposal posterior.
            logits_p: Component weights (i.e. logits) of the proposal distribution.
            means_p: Means of the proposal distribution.
            precisions_p: Precision matrices of the proposal distribution.
            logits_d: Component weights (i.e. logits) of the density estimator.
            means_d: Means of the density estimator.
            precisions_d: Precision matrices of the density estimator.

        Returns: Component weights of the proposal posterior. L*K terms.
        """

        num_comps_p = precisions_p.shape[1]
        num_comps_d = precisions_d.shape[1]

        # Compute log(alpha_i * beta_j)
        logits_p_rep = logits_p.repeat_interleave(num_comps_d, dim=1)
        logits_d_rep = logits_d.repeat(1, num_comps_p)
        logit_factors = logits_p_rep + logits_d_rep

        # Compute sqrt(det()/(det()*det()))
        logdet_covariances_pp = torch.logdet(covariances_pp)
        logdet_covariances_p = -torch.logdet(precisions_p)
        logdet_covariances_d = -torch.logdet(precisions_d)

        # Repeat the proposal and density estimator terms such that there are LK terms.
        # Same trick as has been used above.
        logdet_covariances_p_rep = logdet_covariances_p.repeat_interleave(
            num_comps_d, dim=1
        )
        logdet_covariances_d_rep = logdet_covariances_d.repeat(1, num_comps_p)

        log_sqrt_det_ratio = 0.5 * (
            logdet_covariances_pp
            - (logdet_covariances_p_rep + logdet_covariances_d_rep)
        )

        # Compute for proposal, density estimator, and proposal posterior:
        # mu_i.T * P_i * mu_i
        exponent_p = batched_mixture_vmv(precisions_p, means_p)
        exponent_d = batched_mixture_vmv(precisions_d, means_d)
        exponent_pp = batched_mixture_vmv(precisions_pp, means_pp)

        # Extend proposal and density estimator exponents to get LK terms.
        exponent_p_rep = exponent_p.repeat_interleave(num_comps_d, dim=1)
        exponent_d_rep = exponent_d.repeat(1, num_comps_p)
        exponent = -0.5 * (exponent_p_rep + exponent_d_rep - exponent_pp)

        logits_pp = logit_factors + log_sqrt_det_ratio + exponent

        return logits_pp

__init__(prior=None, density_estimator='maf', device='cpu', logging_level='WARNING', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize NPE-C.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them.	None
density_estimator	Union[Literal['nsf', 'maf', 'mdn', 'made'], ConditionalEstimatorBuilder[ConditionalDensityEstimator]]	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 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() and must return a ConditionalDensityEstimator.	'maf'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'WARNING'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during training.	True

Source code in sbi/inference/trainers/npe/npe_c.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    density_estimator: Union[
        Literal["nsf", "maf", "mdn", "made"],
        ConditionalEstimatorBuilder[ConditionalDensityEstimator],
    ] = "maf",
    device: str = "cpu",
    logging_level: Union[int, str] = "WARNING",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize NPE-C.

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them.
        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
            `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()` 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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during training.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

train(num_atoms=10, training_batch_size=200, learning_rate=0.0005, validation_fraction=0.1, stop_after_epochs=20, max_num_epochs=2 ** 31 - 1, clip_max_norm=5.0, calibration_kernel=None, resume_training=False, force_first_round_loss=False, discard_prior_samples=False, use_combined_loss=False, retrain_from_scratch=False, show_train_summary=False, dataloader_kwargs=None)

Return density estimator that approximates the distribution \(p(\theta|x)\).

Parameters:

Name	Type	Description	Default
num_atoms	int	Number of atoms to use for classification.	10
training_batch_size	int	Training batch size.	200
learning_rate	float	Learning rate for Adam optimizer.	0.0005
validation_fraction	float	The fraction of data to use for validation.	0.1
stop_after_epochs	int	The number of epochs to wait for improvement on the validation set before terminating training.	20
max_num_epochs	int	Maximum number of epochs to run. If reached, we stop training even when the validation loss is still decreasing. Otherwise, we train until validation loss increases (see also stop_after_epochs).	2 ** 31 - 1
clip_max_norm	Optional[float]	Value at which to clip the total gradient norm in order to prevent exploding gradients. Use None for no clipping.	5.0
calibration_kernel	Optional[Callable]	A function to calibrate the loss with respect to the simulations x. See Lueckmann, Gonçalves et al., NeurIPS 2017.	None
resume_training	bool	Can be used in case training time is limited, e.g. on a cluster. If True, the split between train and validation set, the optimizer, the number of epochs, and the best validation log-prob will be restored from the last time .train() was called.	False
force_first_round_loss	bool	If True, train with maximum likelihood, i.e., potentially ignoring the correction for using a proposal distribution different from the prior.	False
discard_prior_samples	bool	Whether to discard samples simulated in round 1, i.e. from the prior. Training may be sped up by ignoring such less targeted samples.	False
use_combined_loss	bool	Whether to train the neural net also on prior samples using maximum likelihood in addition to training it on all samples using atomic loss. The extra MLE loss helps prevent density leaking with bounded priors.	False
retrain_from_scratch	bool	Whether to retrain the conditional density estimator for the posterior from scratch each round.	False
show_train_summary	bool	Whether to print the number of epochs and validation loss and leakage after the training.	False
dataloader_kwargs	Optional[Dict]	Additional or updated kwargs to be passed to the training and validation dataloaders (like, e.g., a collate_fn)	None

Returns:

Type	Description
ConditionalDensityEstimator	Density estimator that approximates the distribution \(p(\theta|x)\).

Source code in sbi/inference/trainers/npe/npe_c.py

def train(
    self,
    num_atoms: int = 10,
    training_batch_size: int = 200,
    learning_rate: float = 5e-4,
    validation_fraction: float = 0.1,
    stop_after_epochs: int = 20,
    max_num_epochs: int = 2**31 - 1,
    clip_max_norm: Optional[float] = 5.0,
    calibration_kernel: Optional[Callable] = None,
    resume_training: bool = False,
    force_first_round_loss: bool = False,
    discard_prior_samples: bool = False,
    use_combined_loss: bool = False,
    retrain_from_scratch: bool = False,
    show_train_summary: bool = False,
    dataloader_kwargs: Optional[Dict] = None,
) -> ConditionalDensityEstimator:
    r"""Return density estimator that approximates the distribution $p(\theta|x)$.

    Args:
        num_atoms: Number of atoms to use for classification.
        training_batch_size: Training batch size.
        learning_rate: Learning rate for Adam optimizer.
        validation_fraction: The fraction of data to use for validation.
        stop_after_epochs: The number of epochs to wait for improvement on the
            validation set before terminating training.
        max_num_epochs: Maximum number of epochs to run. If reached, we stop
            training even when the validation loss is still decreasing. Otherwise,
            we train until validation loss increases (see also
            ``stop_after_epochs``).
        clip_max_norm: Value at which to clip the total gradient norm in order to
            prevent exploding gradients. Use None for no clipping.
        calibration_kernel: A function to calibrate the loss with respect to the
            simulations ``x``. See Lueckmann, Gonçalves et al., NeurIPS 2017.
        resume_training: Can be used in case training time is limited, e.g. on a
            cluster. If ``True``, the split between train and validation set, the
            optimizer, the number of epochs, and the best validation log-prob will
            be restored from the last time ``.train()`` was called.
        force_first_round_loss: If ``True``, train with maximum likelihood,
            i.e., potentially ignoring the correction for using a proposal
            distribution different from the prior.
        discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
            from the prior. Training may be sped up by ignoring such less targeted
            samples.
        use_combined_loss: Whether to train the neural net also on prior samples
            using maximum likelihood in addition to training it on all samples using
            atomic loss. The extra MLE loss helps prevent density leaking with
            bounded priors.
        retrain_from_scratch: Whether to retrain the conditional density
            estimator for the posterior from scratch each round.
        show_train_summary: Whether to print the number of epochs and validation
            loss and leakage after the training.
        dataloader_kwargs: Additional or updated kwargs to be passed to the training
            and validation dataloaders (like, e.g., a collate_fn)

    Returns:
        Density estimator that approximates the distribution $p(\theta|x)$.
    """

    if len(self._data_round_index) == 0:
        raise RuntimeError(
            "No simulations found. You must call .append_simulations() "
            "before calling .train()."
        )
    # WARNING: sneaky trick ahead. We proxy the parent's `train` here,
    # requiring the signature to have `num_atoms`, save it for use below, and
    # continue. It's sneaky because we are using the object (self) as a namespace
    # to pass arguments between functions, and that's implicit state management.
    self._num_atoms = num_atoms
    self._use_combined_loss = use_combined_loss
    kwargs = del_entries(
        locals(),
        entries=("self", "__class__", "num_atoms", "use_combined_loss"),
    )

    self._round = max(self._data_round_index)

    if self._round > 0:
        # Set the proposal to the last proposal that was passed by the user. For
        # atomic SNPE, it does not matter what the proposal is. For non-atomic
        # SNPE, we only use the latest data that was passed, i.e. the one from the
        # last proposal.
        proposal = self._proposal_roundwise[-1]
        self.use_non_atomic_loss = (
            isinstance(proposal, DirectPosterior)
            and isinstance(proposal.posterior_estimator, MixtureDensityEstimator)
            and isinstance(self._neural_net, MixtureDensityEstimator)
            and check_dist_class(
                self._prior, class_to_check=(Uniform, MultivariateNormal)
            )[0]
        )

        algorithm = "non-atomic" if self.use_non_atomic_loss else "atomic"
        print(f"Using SNPE-C with {algorithm} loss")

        if self.use_non_atomic_loss:
            # Take care of z-scoring, pre-compute and store prior terms.
            self._set_state_for_mog_proposal()

    return super().train(**kwargs)

FMPE

Bases: VectorFieldTrainer

Flow Matching Posterior Estimation (FMPE) [1].

FMPE trains a continuous normalizing flow (CNF) to transform samples from the prior distribution to the posterior distribution using flow matching. Instead of maximum likelihood, it trains a vector field to match the marginal vector field of a conditional flow that interpolates between the prior and posterior. The neural network architecture for the vector field is not constrained like for flows and can be any expressive network. Sampling is performed by solving an ODE, which can be slower than flow-based NPE, but log_prob evaluation can also be slower.

NOTE: FMPE does not support multi-round inference with flexible proposals yet. You can try multi-round with truncated proposals, but this is not tested.

[1] Flow Matching for Generative Modeling, Lipman et al., ICLR 2023, https://arxiv.org/abs/2210.02747

Example:

::

import torch
from sbi.inference import FMPE
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train flow matching estimator
inference = FMPE(prior=prior)
flow_estimator = inference.append_simulations(theta, x).train()

# 3. Build posterior (uses ODE solver for sampling)
posterior = inference.build_posterior(flow_estimator)

# 4. Sample from posterior
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/vfpe/fmpe.py

class FMPE(VectorFieldTrainer):
    r"""Flow Matching Posterior Estimation (FMPE) [1].

    FMPE trains a continuous normalizing flow (CNF) to transform samples from the
    prior distribution to the posterior distribution using flow matching. Instead of
    maximum likelihood, it trains a vector field to match the marginal vector field
    of a conditional flow that interpolates between the prior and posterior. The
    neural network architecture for the vector field is not constrained like for
    flows and can be any expressive network. Sampling is performed by solving an ODE,
    which can be slower than flow-based NPE, but log_prob evaluation can also be slower.

    NOTE: FMPE does not support multi-round inference with flexible proposals yet.
    You can try multi-round with truncated proposals, but this is not tested.

    [1] Flow Matching for Generative Modeling, Lipman et al., ICLR 2023,
        https://arxiv.org/abs/2210.02747

    Example:
    --------

    ::

        import torch
        from sbi.inference import FMPE
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train flow matching estimator
        inference = FMPE(prior=prior)
        flow_estimator = inference.append_simulations(theta, x).train()

        # 3. Build posterior (uses ODE solver for sampling)
        posterior = inference.build_posterior(flow_estimator)

        # 4. Sample from posterior
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        vf_estimator: Union[
            Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
            ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator],
        ] = "mlp",
        density_estimator: Optional[
            ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]
        ] = None,
        device: str = "cpu",
        logging_level: Union[int, str] = "WARNING",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ) -> None:
        """Initialization method for the FMPE class.

        Args:
            prior: Prior distribution.
            vf_estimator: Neural network architecture used to learn the
                vector field estimator. Can be a string (e.g. 'mlp', 'ada_mlp',
                'transformer' or 'transformer_cross_attn') or a callable that implements
                the `ConditionalEstimatorBuilder` protocol with `__call__` that receives
                `theta` and `x` and returns a `ConditionalVectorFieldEstimator`.
                To configure estimator-level options, use `posterior_flow_nn` to
                build a custom callable and pass it here.
            density_estimator: Deprecated. Use `vf_estimator` instead.
            device: Device to use for training.
            logging_level: Logging level.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show progress bars.
        """

        if density_estimator is not None:
            warnings.warn(
                "`density_estimator` is deprecated and will be removed in a future "
                "release. Use `vf_estimator` instead.",
                FutureWarning,
                stacklevel=2,
            )
            vf_estimator = density_estimator

        super().__init__(
            prior=prior,
            device=device,
            logging_level=logging_level,
            summary_writer=summary_writer,
            tracker=tracker,
            show_progress_bars=show_progress_bars,
            vector_field_estimator_builder=vf_estimator,
        )

        # When vf_estimator is a string, build the default neural net.
        if isinstance(vf_estimator, str):
            self._build_neural_net = self._build_default_nn_fn(model=vf_estimator)

    def build_posterior(
        self,
        vector_field_estimator: Optional[ConditionalVectorFieldEstimator] = None,
        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.

        Note that this is the same as the NPSE posterior, but the sample_with method
        is set to "ode" by default.

        For FMPE, the posterior distribution that is returned here implements
        the following functionality over the raw neural density estimator:
            - correct the calculation of the log probability such that
              samples outside of the prior bounds have log probability -inf.
            - reject samples that lie outside of the prior bounds.

        Args:
            vector_field_estimator: The flow matching estimator that
                the posterior is based on. If `None`, use the latest neural
                flow matching estimator that was trained.
            prior: Prior distribution.
            sample_with: Method to use for sampling from the posterior.
                Can be one of 'sde' (default) or 'ode'. The 'sde' method uses
                the score to do a Langevin diffusion step, while the 'ode' method
                uses the score to define a probabilistic ODE and solves it with
                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,
        model: Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
    ) -> ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]:
        return posterior_flow_nn(model=model)

__init__(prior=None, vf_estimator='mlp', density_estimator=None, device='cpu', logging_level='WARNING', summary_writer=None, tracker=None, show_progress_bars=True)

Initialization method for the FMPE class.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	Prior distribution.	None
vf_estimator	Union[Literal['mlp', 'ada_mlp', 'transformer', 'transformer_cross_attn'], ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]]	Neural network architecture used to learn the vector field estimator. Can be a string (e.g. ‘mlp’, ‘ada_mlp’, ‘transformer’ or ‘transformer_cross_attn’) or a callable that implements the ConditionalEstimatorBuilder protocol with __call__ that receives theta and x and returns a ConditionalVectorFieldEstimator. To configure estimator-level options, use posterior_flow_nn to build a custom callable and pass it here.	'mlp'
density_estimator	Optional[ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]]	Deprecated. Use vf_estimator instead.	None
device	str	Device to use for training.	'cpu'
logging_level	Union[int, str]	Logging level.	'WARNING'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show progress bars.	True

Source code in sbi/inference/trainers/vfpe/fmpe.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    vf_estimator: Union[
        Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
        ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator],
    ] = "mlp",
    density_estimator: Optional[
        ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]
    ] = None,
    device: str = "cpu",
    logging_level: Union[int, str] = "WARNING",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
) -> None:
    """Initialization method for the FMPE class.

    Args:
        prior: Prior distribution.
        vf_estimator: Neural network architecture used to learn the
            vector field estimator. Can be a string (e.g. 'mlp', 'ada_mlp',
            'transformer' or 'transformer_cross_attn') or a callable that implements
            the `ConditionalEstimatorBuilder` protocol with `__call__` that receives
            `theta` and `x` and returns a `ConditionalVectorFieldEstimator`.
            To configure estimator-level options, use `posterior_flow_nn` to
            build a custom callable and pass it here.
        density_estimator: Deprecated. Use `vf_estimator` instead.
        device: Device to use for training.
        logging_level: Logging level.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show progress bars.
    """

    if density_estimator is not None:
        warnings.warn(
            "`density_estimator` is deprecated and will be removed in a future "
            "release. Use `vf_estimator` instead.",
            FutureWarning,
            stacklevel=2,
        )
        vf_estimator = density_estimator

    super().__init__(
        prior=prior,
        device=device,
        logging_level=logging_level,
        summary_writer=summary_writer,
        tracker=tracker,
        show_progress_bars=show_progress_bars,
        vector_field_estimator_builder=vf_estimator,
    )

    # When vf_estimator is a string, build the default neural net.
    if isinstance(vf_estimator, str):
        self._build_neural_net = self._build_default_nn_fn(model=vf_estimator)

build_posterior(vector_field_estimator=None, prior=None, sample_with='ode', vectorfield_sampling_parameters=None, posterior_parameters=None)

Build posterior from the flow matching estimator.

Note that this is the same as the NPSE posterior, but the sample_with method is set to “ode” by default.

For FMPE, the posterior distribution that is returned here implements the following functionality over the raw neural density estimator: - correct the calculation of the log probability such that samples outside of the prior bounds have log probability -inf. - reject samples that lie outside of the prior bounds.

Parameters:

Name	Type	Description	Default
vector_field_estimator	Optional[ConditionalVectorFieldEstimator]	The flow matching estimator that the posterior is based on. If None, use the latest neural flow matching estimator that was trained.	None
prior	Optional[Distribution]	Prior distribution.	None
sample_with	Literal['ode', 'sde']	Method to use for sampling from the posterior. Can be one of ‘sde’ (default) or ‘ode’. The ‘sde’ method uses the score to do a Langevin diffusion step, while the ‘ode’ method uses the score to define a probabilistic ODE and solves it with a numerical ODE solver.	'ode'
vectorfield_sampling_parameters	Optional[Dict[str, Any]]	Additional keyword arguments passed to VectorFieldPosterior.	None
posterior_parameters	Optional[VectorFieldPosteriorParameters]	Configuration passed to the init method for VectorFieldPosterior.	None

Returns:

Type	Description
NeuralPosterior	Posterior \(p(\theta|x)\) with .sample() and .log_prob() methods.

Source code in sbi/inference/trainers/vfpe/fmpe.py

def build_posterior(
    self,
    vector_field_estimator: Optional[ConditionalVectorFieldEstimator] = None,
    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.

    Note that this is the same as the NPSE posterior, but the sample_with method
    is set to "ode" by default.

    For FMPE, the posterior distribution that is returned here implements
    the following functionality over the raw neural density estimator:
        - correct the calculation of the log probability such that
          samples outside of the prior bounds have log probability -inf.
        - reject samples that lie outside of the prior bounds.

    Args:
        vector_field_estimator: The flow matching estimator that
            the posterior is based on. If `None`, use the latest neural
            flow matching estimator that was trained.
        prior: Prior distribution.
        sample_with: Method to use for sampling from the posterior.
            Can be one of 'sde' (default) or 'ode'. The 'sde' method uses
            the score to do a Langevin diffusion step, while the 'ode' method
            uses the score to define a probabilistic ODE and solves it with
            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,
    )

NPSE

Bases: VectorFieldTrainer

Neural Posterior Score Estimation (NPSE) [1, 2].

NPSE trains a neural network to estimate the score function (gradient of the log posterior) \(\nabla_\theta \log p(\theta|x)\) using denoising score matching. NPSE learns the score of a diffusion process that transforms the prior into the posterior. The neural network can be any expressive architecture. Sampling is performed using SDE solvers (e.g., Langevin dynamics) or ODE solvers, which can be slower than flow-based NPE, but expressiveness can be higher.

NOTE: NPSE does not support multi-round inference with flexible proposals yet. You can try multi-round with truncated proposals, but this is not tested.

[1] Score modeling for simulation-based inference, Geffner et al., ICML 2023. [2] Sequential neural score estimation: Likelihood-free inference with conditional score based diffusion models, Sharrock et al., ICML 2024.

Example:

::

import torch
from sbi.inference import NPSE
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train score estimator
inference = NPSE(prior=prior, sde_type="ve")
score_estimator = inference.append_simulations(theta, x).train()

# 3. Build posterior (uses SDE solver by default)
posterior = inference.build_posterior(score_estimator)

# 4. Sample from posterior using Langevin dynamics
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/vfpe/npse.py

class NPSE(VectorFieldTrainer):
    r"""Neural Posterior Score Estimation (NPSE) [1, 2].

    NPSE trains a neural network to estimate the score function (gradient of the log
    posterior) $\nabla_\theta \log p(\theta|x)$ using denoising score matching. NPSE
    learns the score of a diffusion process that transforms the prior into the
    posterior. The neural network can be any expressive architecture. Sampling is
    performed using SDE solvers (e.g., Langevin dynamics) or ODE solvers, which can
    be slower than flow-based NPE, but expressiveness can be higher.

    NOTE: NPSE does not support multi-round inference with flexible proposals yet.
    You can try multi-round with truncated proposals, but this is not tested.

    [1] Score modeling for simulation-based inference, Geffner et al., ICML 2023.
    [2] Sequential neural score estimation: Likelihood-free inference with conditional
        score based diffusion models, Sharrock et al., ICML 2024.

    Example:
    --------

    ::

        import torch
        from sbi.inference import NPSE
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train score estimator
        inference = NPSE(prior=prior, sde_type="ve")
        score_estimator = inference.append_simulations(theta, x).train()

        # 3. Build posterior (uses SDE solver by default)
        posterior = inference.build_posterior(score_estimator)

        # 4. Sample from posterior using Langevin dynamics
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        vf_estimator: Union[
            Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
            ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator],
        ] = "mlp",
        score_estimator: Optional[
            Union[
                Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
                ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator],
            ]
        ] = None,
        density_estimator: Optional[
            ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]
        ] = None,
        sde_type: Literal["vp", "ve", "subvp"] = "ve",
        device: str = "cpu",
        logging_level: Union[int, str] = "WARNING",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize Neural Posterior Score Estimation.

        Args:
            prior: Prior distribution.
            vf_estimator: Neural network architecture for the
                vector field estimator aiming to estimate the marginal scores of the
                target diffusion process. Can be a string (e.g. 'mlp', 'ada_mlp',
                'transformer' or 'transformer_cross_attn') or a callable that implements
                the `ConditionalEstimatorBuilder` protocol with `__call__` that receives
                `theta` and `x` and returns a `ConditionalVectorFieldEstimator`.
                To configure estimator-level options (e.g. noise schedules for VE),
                use `posterior_score_nn` to build a custom callable and pass it here.
            score_estimator: Deprecated, use `vf_estimator` instead.
            density_estimator: Deprecated, use `vf_estimator` instead.
            sde_type: Type of SDE to use. Must be one of ['vp', 've', 'subvp'].
                Only used when `vf_estimator` is a string (i.e. when using the
                default builder). Ignored when a custom callable is passed.
            device: Device to run the training on.
            logging_level: Logging level for the training. Can be an integer or a
                string.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show progress bars during training.

        References:
            - Geffner, Tomas, George Papamakarios, and Andriy Mnih. "Score modeling for
                simulation-based inference." ICML 2023.
            - Sharrock, Louis, et al. "Sequential neural score estimation: Likelihood-
                free inference with conditional score based diffusion models." ICML 2024
        """
        if score_estimator is not None:
            warnings.warn(
                "`score_estimator` is deprecated and will be removed in a future "
                "release. Use `vf_estimator` instead.",
                FutureWarning,
                stacklevel=2,
            )
            vf_estimator = score_estimator

        if density_estimator is not None:
            warnings.warn(
                "`density_estimator` is deprecated and will be removed in a future "
                "release. Use `vf_estimator` instead.",
                FutureWarning,
                stacklevel=2,
            )
            vf_estimator = density_estimator

        super().__init__(
            prior=prior,
            vector_field_estimator_builder=vf_estimator,
            device=device,
            logging_level=logging_level,
            summary_writer=summary_writer,
            tracker=tracker,
            show_progress_bars=show_progress_bars,
        )

        # When vf_estimator is a string, build the default neural net using
        # the NPSE-specific builder which requires sde_type.
        if isinstance(vf_estimator, str):
            self._build_neural_net = self._build_default_nn_fn(
                model=vf_estimator, sde_type=sde_type
            )

    def build_posterior(
        self,
        vector_field_estimator: Optional[ConditionalVectorFieldEstimator] = None,
        prior: Optional[Distribution] = None,
        sample_with: Literal["ode", "sde"] = "sde",
        vectorfield_sampling_parameters: Optional[Dict[str, Any]] = None,
        posterior_parameters: Optional[VectorFieldPosteriorParameters] = None,
    ) -> NeuralPosterior:
        r"""Build posterior from the vector field estimator.

        Note that this is the same as the FMPE posterior, but the sample_with
        method is set to "sde" by default.

        For NPSE, the posterior distribution that is returned here implements
        the following functionality over the raw neural density estimator:
            - correct the calculation of the log probability such that
              samples outside of the prior bounds have log probability -inf.
            - reject samples that lie outside of the prior bounds.

        Args:
            vector_field_estimator: The vector field estimator that the posterior is
                based on. If `None`, use the latest vector field estimator that was
                trained.
            prior: Prior distribution.
            sample_with: Method to use for sampling from the posterior. Can be one of
                'sde' (default) or 'ode'. The 'sde' method uses the score to
                do a Langevin diffusion step, while the 'ode' method solves a
                probabilistic ODE with 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,
        model: Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
        sde_type: Literal["vp", "ve", "subvp"] = "ve",
    ) -> ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]:
        return posterior_score_nn(model=model, sde_type=sde_type)

__init__(prior=None, vf_estimator='mlp', score_estimator=None, density_estimator=None, sde_type='ve', device='cpu', logging_level='WARNING', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize Neural Posterior Score Estimation.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	Prior distribution.	None
vf_estimator	Union[Literal['mlp', 'ada_mlp', 'transformer', 'transformer_cross_attn'], ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]]	Neural network architecture for the vector field estimator aiming to estimate the marginal scores of the target diffusion process. Can be a string (e.g. ‘mlp’, ‘ada_mlp’, ‘transformer’ or ‘transformer_cross_attn’) or a callable that implements the ConditionalEstimatorBuilder protocol with __call__ that receives theta and x and returns a ConditionalVectorFieldEstimator. To configure estimator-level options (e.g. noise schedules for VE), use posterior_score_nn to build a custom callable and pass it here.	'mlp'
score_estimator	Optional[Union[Literal['mlp', 'ada_mlp', 'transformer', 'transformer_cross_attn'], ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]]]	Deprecated, use vf_estimator instead.	None
density_estimator	Optional[ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]]	Deprecated, use vf_estimator instead.	None
sde_type	Literal['vp', 've', 'subvp']	Type of SDE to use. Must be one of [‘vp’, ‘ve’, ‘subvp’]. Only used when vf_estimator is a string (i.e. when using the default builder). Ignored when a custom callable is passed.	've'
device	str	Device to run the training on.	'cpu'
logging_level	Union[int, str]	Logging level for the training. Can be an integer or a string.	'WARNING'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show progress bars during training.	True

References

• Geffner, Tomas, George Papamakarios, and Andriy Mnih. “Score modeling for simulation-based inference.” ICML 2023.
• Sharrock, Louis, et al. “Sequential neural score estimation: Likelihood- free inference with conditional score based diffusion models.” ICML 2024

Source code in sbi/inference/trainers/vfpe/npse.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    vf_estimator: Union[
        Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
        ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator],
    ] = "mlp",
    score_estimator: Optional[
        Union[
            Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
            ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator],
        ]
    ] = None,
    density_estimator: Optional[
        ConditionalEstimatorBuilder[ConditionalVectorFieldEstimator]
    ] = None,
    sde_type: Literal["vp", "ve", "subvp"] = "ve",
    device: str = "cpu",
    logging_level: Union[int, str] = "WARNING",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize Neural Posterior Score Estimation.

    Args:
        prior: Prior distribution.
        vf_estimator: Neural network architecture for the
            vector field estimator aiming to estimate the marginal scores of the
            target diffusion process. Can be a string (e.g. 'mlp', 'ada_mlp',
            'transformer' or 'transformer_cross_attn') or a callable that implements
            the `ConditionalEstimatorBuilder` protocol with `__call__` that receives
            `theta` and `x` and returns a `ConditionalVectorFieldEstimator`.
            To configure estimator-level options (e.g. noise schedules for VE),
            use `posterior_score_nn` to build a custom callable and pass it here.
        score_estimator: Deprecated, use `vf_estimator` instead.
        density_estimator: Deprecated, use `vf_estimator` instead.
        sde_type: Type of SDE to use. Must be one of ['vp', 've', 'subvp'].
            Only used when `vf_estimator` is a string (i.e. when using the
            default builder). Ignored when a custom callable is passed.
        device: Device to run the training on.
        logging_level: Logging level for the training. Can be an integer or a
            string.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show progress bars during training.

    References:
        - Geffner, Tomas, George Papamakarios, and Andriy Mnih. "Score modeling for
            simulation-based inference." ICML 2023.
        - Sharrock, Louis, et al. "Sequential neural score estimation: Likelihood-
            free inference with conditional score based diffusion models." ICML 2024
    """
    if score_estimator is not None:
        warnings.warn(
            "`score_estimator` is deprecated and will be removed in a future "
            "release. Use `vf_estimator` instead.",
            FutureWarning,
            stacklevel=2,
        )
        vf_estimator = score_estimator

    if density_estimator is not None:
        warnings.warn(
            "`density_estimator` is deprecated and will be removed in a future "
            "release. Use `vf_estimator` instead.",
            FutureWarning,
            stacklevel=2,
        )
        vf_estimator = density_estimator

    super().__init__(
        prior=prior,
        vector_field_estimator_builder=vf_estimator,
        device=device,
        logging_level=logging_level,
        summary_writer=summary_writer,
        tracker=tracker,
        show_progress_bars=show_progress_bars,
    )

    # When vf_estimator is a string, build the default neural net using
    # the NPSE-specific builder which requires sde_type.
    if isinstance(vf_estimator, str):
        self._build_neural_net = self._build_default_nn_fn(
            model=vf_estimator, sde_type=sde_type
        )

build_posterior(vector_field_estimator=None, prior=None, sample_with='sde', vectorfield_sampling_parameters=None, posterior_parameters=None)

Build posterior from the vector field estimator.

Note that this is the same as the FMPE posterior, but the sample_with method is set to “sde” by default.

For NPSE, the posterior distribution that is returned here implements the following functionality over the raw neural density estimator: - correct the calculation of the log probability such that samples outside of the prior bounds have log probability -inf. - reject samples that lie outside of the prior bounds.

Parameters:

Name	Type	Description	Default
vector_field_estimator	Optional[ConditionalVectorFieldEstimator]	The vector field estimator that the posterior is based on. If None, use the latest vector field estimator that was trained.	None
prior	Optional[Distribution]	Prior distribution.	None
sample_with	Literal['ode', 'sde']	Method to use for sampling from the posterior. Can be one of ‘sde’ (default) or ‘ode’. The ‘sde’ method uses the score to do a Langevin diffusion step, while the ‘ode’ method solves a probabilistic ODE with a numerical ODE solver.	'sde'
vectorfield_sampling_parameters	Optional[Dict[str, Any]]	Additional keyword arguments passed to VectorFieldPosterior.	None
posterior_parameters	Optional[VectorFieldPosteriorParameters]	Configuration passed to the init method for VectorFieldPosterior.	None

Returns:

Type	Description
NeuralPosterior	Posterior \(p(\theta|x)\) with .sample() and .log_prob() methods.

Source code in sbi/inference/trainers/vfpe/npse.py

def build_posterior(
    self,
    vector_field_estimator: Optional[ConditionalVectorFieldEstimator] = None,
    prior: Optional[Distribution] = None,
    sample_with: Literal["ode", "sde"] = "sde",
    vectorfield_sampling_parameters: Optional[Dict[str, Any]] = None,
    posterior_parameters: Optional[VectorFieldPosteriorParameters] = None,
) -> NeuralPosterior:
    r"""Build posterior from the vector field estimator.

    Note that this is the same as the FMPE posterior, but the sample_with
    method is set to "sde" by default.

    For NPSE, the posterior distribution that is returned here implements
    the following functionality over the raw neural density estimator:
        - correct the calculation of the log probability such that
          samples outside of the prior bounds have log probability -inf.
        - reject samples that lie outside of the prior bounds.

    Args:
        vector_field_estimator: The vector field estimator that the posterior is
            based on. If `None`, use the latest vector field estimator that was
            trained.
        prior: Prior distribution.
        sample_with: Method to use for sampling from the posterior. Can be one of
            'sde' (default) or 'ode'. The 'sde' method uses the score to
            do a Langevin diffusion step, while the 'ode' method solves a
            probabilistic ODE with 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,
    )

NLE_A

Bases: LikelihoodEstimatorTrainer

Neural Likelihood Estimation (NLE) as in Papamakarios et al. (2019) [1].

NLE trains a neural network to approximate the likelihood \(p(x|\theta)\) using a conditional density estimator (normalizing flow). Unlike NPE methods, which directly estimate the posterior, NLE estimates the likelihood.

NLE can be run multi-round without need for correction, but requires running potentially expensive posterior sampling in each round.

[1] Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows, Papamakarios et al., AISTATS 2019, https://arxiv.org/abs/1805.07226

Example:

::

import torch
from sbi.inference import NLE_A
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train likelihood estimator
inference = NLE_A(prior=prior)
likelihood_estimator = inference.append_simulations(theta, x).train()

# 3. Build posterior (uses MCMC for sampling)
posterior = inference.build_posterior(likelihood_estimator)

# 4. Sample from posterior
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/nle/nle_a.py

class NLE_A(LikelihoodEstimatorTrainer):
    r"""Neural Likelihood Estimation (NLE) as in Papamakarios et al. (2019) [1].

    NLE trains a neural network to approximate the likelihood $p(x|\theta)$ using a
    conditional density estimator (normalizing flow). Unlike NPE methods, which directly
    estimate the posterior, NLE estimates the likelihood.

    NLE can be run multi-round without need for correction, but requires running
    potentially expensive posterior sampling in each round.

    [1] Sequential Neural Likelihood: Fast Likelihood-free Inference with
        Autoregressive Flows, Papamakarios et al., AISTATS 2019,
        https://arxiv.org/abs/1805.07226

    Example:
    --------

    ::

        import torch
        from sbi.inference import NLE_A
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train likelihood estimator
        inference = NLE_A(prior=prior)
        likelihood_estimator = inference.append_simulations(theta, x).train()

        # 3. Build posterior (uses MCMC for sampling)
        posterior = inference.build_posterior(likelihood_estimator)

        # 4. Sample from posterior
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        density_estimator: Union[
            Literal["nsf", "maf", "mdn", "made"],
            ConditionalEstimatorBuilder[ConditionalDensityEstimator],
        ] = "maf",
        device: str = "cpu",
        logging_level: Union[int, str] = "WARNING",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize Neural Likelihood Estimation.

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. If `None`, the
                prior must be passed to `.build_posterior()`.
            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
                `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()` 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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during simulation and
                sampling.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

__init__(prior=None, density_estimator='maf', device='cpu', logging_level='WARNING', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize Neural Likelihood Estimation.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. If None, the prior must be passed to .build_posterior().	None
density_estimator	Union[Literal['nsf', 'maf', 'mdn', 'made'], ConditionalEstimatorBuilder[ConditionalDensityEstimator]]	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 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() and must return a ConditionalDensityEstimator.	'maf'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'WARNING'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during simulation and sampling.	True

Source code in sbi/inference/trainers/nle/nle_a.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    density_estimator: Union[
        Literal["nsf", "maf", "mdn", "made"],
        ConditionalEstimatorBuilder[ConditionalDensityEstimator],
    ] = "maf",
    device: str = "cpu",
    logging_level: Union[int, str] = "WARNING",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize Neural Likelihood Estimation.

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. If `None`, the
            prior must be passed to `.build_posterior()`.
        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
            `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()` 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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during simulation and
            sampling.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

NRE_A

Bases: RatioEstimatorTrainer

Neural Ratio Estimation (NRE-A / AALR) as in Hermans et al. (2020) [1].

NRE-A trains a neural classifier to estimate the likelihood-to-evidence ratio \(r(\theta, x) = p(x|\theta) / p(x)\) by distinguishing between samples from the joint distribution \(p(\theta, x)\) and samples from the marginals \(p(\theta)p(x)\). Posterior sampling is then performed via MCMC, rejection sampling, or variational inference using the estimated ratio.

NRE can be run multi-round without need for correction, but requires running potentially expensive posterior sampling in each round.

[1] Likelihood-free MCMC with Amortized Approximate Likelihood Ratios, Hermans et al., ICML 2020, https://arxiv.org/abs/1903.04057

Example:

::

import torch
from sbi.inference import NRE_A
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train ratio estimator
inference = NRE_A(prior=prior)
ratio_estimator = inference.append_simulations(theta, x).train()

# 3. Build posterior (uses MCMC or rejection sampling)
posterior = inference.build_posterior(ratio_estimator)

# 4. Sample from posterior
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/nre/nre_a.py

class NRE_A(RatioEstimatorTrainer):
    r"""Neural Ratio Estimation (NRE-A / AALR) as in Hermans et al. (2020) [1].

    NRE-A trains a neural classifier to estimate the likelihood-to-evidence ratio
    $r(\theta, x) = p(x|\theta) / p(x)$ by distinguishing between samples from the
    joint distribution $p(\theta, x)$ and samples from the marginals $p(\theta)p(x)$.
    Posterior sampling is then performed via MCMC, rejection sampling, or variational
    inference using the estimated ratio.

    NRE can be run multi-round without need for correction, but requires running
    potentially expensive posterior sampling in each round.

    [1] *Likelihood-free MCMC with Amortized Approximate Likelihood Ratios*, Hermans
        et al., ICML 2020, https://arxiv.org/abs/1903.04057

    Example:
    --------

    ::

        import torch
        from sbi.inference import NRE_A
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train ratio estimator
        inference = NRE_A(prior=prior)
        ratio_estimator = inference.append_simulations(theta, x).train()

        # 3. Build posterior (uses MCMC or rejection sampling)
        posterior = inference.build_posterior(ratio_estimator)

        # 4. Sample from posterior
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
        device: str = "cpu",
        logging_level: Union[int, str] = "warning",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize NRE_A.

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. If `None`, the
                prior must be passed to `.build_posterior()`.
            classifier: Classifier trained to approximate likelihood ratios. If it is
                a string, use a pre-configured network of the provided type (one of
                linear, mlp, resnet), or a callable that implements the
                `ConditionalEstimatorBuilder` protocol. The callable will
                be called with the first batch of simulations (theta, x), which can
                thus be used for shape inference and potentially for z-scoring. It
                returns a `RatioEstimator`.
            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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during simulation and
                sampling.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

    def train(
        self,
        training_batch_size: int = 200,
        learning_rate: float = 5e-4,
        validation_fraction: float = 0.1,
        stop_after_epochs: int = 20,
        max_num_epochs: int = 2**31 - 1,
        clip_max_norm: Optional[float] = 5.0,
        resume_training: bool = False,
        discard_prior_samples: bool = False,
        retrain_from_scratch: bool = False,
        show_train_summary: bool = False,
        dataloader_kwargs: Optional[Dict] = None,
        loss_kwargs: Optional[LossArgsNRE_A] = None,
    ) -> RatioEstimator:
        r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.

        Args:
            training_batch_size: Training batch size.
            learning_rate: Learning rate for Adam optimizer.
            validation_fraction: The fraction of data to use for validation.
            stop_after_epochs: The number of epochs to wait for improvement on the
                validation set before terminating training.
            max_num_epochs: Maximum number of epochs to run. If reached, we stop
                training even when the validation loss is still decreasing. Otherwise,
                we train until validation loss increases (see also `stop_after_epochs`).
            clip_max_norm: Value at which to clip the total gradient norm in order to
                prevent exploding gradients. Use None for no clipping.
            resume_training: Can be used in case training time is limited, e.g. on a
                cluster. If `True`, the split between train and validation set, the
                optimizer, the number of epochs, and the best validation log-prob will
                be restored from the last time `.train()` was called.
            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                from the prior. Training may be sped up by ignoring such less targeted
                samples.
            retrain_from_scratch: Whether to retrain the conditional density
                estimator for the posterior from scratch each round.
            show_train_summary: Whether to print the number of epochs and validation
                loss and leakage after the training.
            dataloader_kwargs: Additional or updated kwargs to be passed to the training
                and validation dataloaders (like, e.g., a collate_fn)
            loss_kwargs: Additional or updated kwargs to be passed to the self._loss fn.

        Returns:
            Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))

        if loss_kwargs is None:
            kwargs["loss_kwargs"] = LossArgsNRE_A()
        elif not issubclass(type(loss_kwargs), LossArgsNRE_A):
            raise TypeError(
                "Expected loss_kwargs to be a subclass of LossArgsNRE_A,"
                f" but got {type(loss_kwargs)}"
            )

        return super().train(**kwargs)

    def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:
        """Returns the binary cross-entropy loss for the trained classifier.

        The classifier takes as input a $(\theta,x)$ pair. It is trained to predict 1
        if the pair was sampled from the joint $p(\theta,x)$, and to predict 0 if the
        pair was sampled from the marginals $p(\theta)p(x)$.
        """

        assert theta.shape[0] == x.shape[0], "Batch sizes for theta and x must match."
        batch_size = theta.shape[0]

        logits = self._classifier_logits(theta, x, num_atoms)
        likelihood = torch.sigmoid(logits).squeeze()

        # Alternating pairs where there is one sampled from the joint and one
        # sampled from the marginals. The first element is sampled from the
        # joint p(theta, x) and is labelled 1. The second element is sampled
        # from the marginals p(theta)p(x) and is labelled 0. And so on.
        labels = ones(2 * batch_size, device=self._device)  # two atoms
        labels[1::2] = 0.0

        # Binary cross entropy to learn the likelihood (AALR-specific)
        loss = nn.BCELoss()(likelihood, labels)
        assert_all_finite(loss, "NRE-A loss")
        return loss

__init__(prior=None, classifier='resnet', device='cpu', logging_level='warning', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize NRE_A.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. If None, the prior must be passed to .build_posterior().	None
classifier	Union[str, ConditionalEstimatorBuilder[RatioEstimator]]	Classifier trained to approximate likelihood ratios. If it is a string, use a pre-configured network of the provided type (one of linear, mlp, resnet), or a callable that implements the ConditionalEstimatorBuilder protocol. The callable will be called with the first batch of simulations (theta, x), which can thus be used for shape inference and potentially for z-scoring. It returns a RatioEstimator.	'resnet'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'warning'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during simulation and sampling.	True

Source code in sbi/inference/trainers/nre/nre_a.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
    device: str = "cpu",
    logging_level: Union[int, str] = "warning",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize NRE_A.

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. If `None`, the
            prior must be passed to `.build_posterior()`.
        classifier: Classifier trained to approximate likelihood ratios. If it is
            a string, use a pre-configured network of the provided type (one of
            linear, mlp, resnet), or a callable that implements the
            `ConditionalEstimatorBuilder` protocol. The callable will
            be called with the first batch of simulations (theta, x), which can
            thus be used for shape inference and potentially for z-scoring. It
            returns a `RatioEstimator`.
        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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during simulation and
            sampling.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

train(training_batch_size=200, learning_rate=0.0005, validation_fraction=0.1, stop_after_epochs=20, max_num_epochs=2 ** 31 - 1, clip_max_norm=5.0, resume_training=False, discard_prior_samples=False, retrain_from_scratch=False, show_train_summary=False, dataloader_kwargs=None, loss_kwargs=None)

Return classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Parameters:

Name	Type	Description	Default
training_batch_size	int	Training batch size.	200
learning_rate	float	Learning rate for Adam optimizer.	0.0005
validation_fraction	float	The fraction of data to use for validation.	0.1
stop_after_epochs	int	The number of epochs to wait for improvement on the validation set before terminating training.	20
max_num_epochs	int	Maximum number of epochs to run. If reached, we stop training even when the validation loss is still decreasing. Otherwise, we train until validation loss increases (see also stop_after_epochs).	2 ** 31 - 1
clip_max_norm	Optional[float]	Value at which to clip the total gradient norm in order to prevent exploding gradients. Use None for no clipping.	5.0
resume_training	bool	Can be used in case training time is limited, e.g. on a cluster. If True, the split between train and validation set, the optimizer, the number of epochs, and the best validation log-prob will be restored from the last time .train() was called.	False
discard_prior_samples	bool	Whether to discard samples simulated in round 1, i.e. from the prior. Training may be sped up by ignoring such less targeted samples.	False
retrain_from_scratch	bool	Whether to retrain the conditional density estimator for the posterior from scratch each round.	False
show_train_summary	bool	Whether to print the number of epochs and validation loss and leakage after the training.	False
dataloader_kwargs	Optional[Dict]	Additional or updated kwargs to be passed to the training and validation dataloaders (like, e.g., a collate_fn)	None
loss_kwargs	Optional[LossArgsNRE_A]	Additional or updated kwargs to be passed to the self._loss fn.	None

Returns:

Type	Description
RatioEstimator	Classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Source code in sbi/inference/trainers/nre/nre_a.py

def train(
    self,
    training_batch_size: int = 200,
    learning_rate: float = 5e-4,
    validation_fraction: float = 0.1,
    stop_after_epochs: int = 20,
    max_num_epochs: int = 2**31 - 1,
    clip_max_norm: Optional[float] = 5.0,
    resume_training: bool = False,
    discard_prior_samples: bool = False,
    retrain_from_scratch: bool = False,
    show_train_summary: bool = False,
    dataloader_kwargs: Optional[Dict] = None,
    loss_kwargs: Optional[LossArgsNRE_A] = None,
) -> RatioEstimator:
    r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.

    Args:
        training_batch_size: Training batch size.
        learning_rate: Learning rate for Adam optimizer.
        validation_fraction: The fraction of data to use for validation.
        stop_after_epochs: The number of epochs to wait for improvement on the
            validation set before terminating training.
        max_num_epochs: Maximum number of epochs to run. If reached, we stop
            training even when the validation loss is still decreasing. Otherwise,
            we train until validation loss increases (see also `stop_after_epochs`).
        clip_max_norm: Value at which to clip the total gradient norm in order to
            prevent exploding gradients. Use None for no clipping.
        resume_training: Can be used in case training time is limited, e.g. on a
            cluster. If `True`, the split between train and validation set, the
            optimizer, the number of epochs, and the best validation log-prob will
            be restored from the last time `.train()` was called.
        discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
            from the prior. Training may be sped up by ignoring such less targeted
            samples.
        retrain_from_scratch: Whether to retrain the conditional density
            estimator for the posterior from scratch each round.
        show_train_summary: Whether to print the number of epochs and validation
            loss and leakage after the training.
        dataloader_kwargs: Additional or updated kwargs to be passed to the training
            and validation dataloaders (like, e.g., a collate_fn)
        loss_kwargs: Additional or updated kwargs to be passed to the self._loss fn.

    Returns:
        Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))

    if loss_kwargs is None:
        kwargs["loss_kwargs"] = LossArgsNRE_A()
    elif not issubclass(type(loss_kwargs), LossArgsNRE_A):
        raise TypeError(
            "Expected loss_kwargs to be a subclass of LossArgsNRE_A,"
            f" but got {type(loss_kwargs)}"
        )

    return super().train(**kwargs)

NRE_B

Bases: RatioEstimatorTrainer

Neural Ratio Estimation (NRE-B / SRE) as in Durkan et al. (2020) [1].

NRE-B is an extension of NRE-A that trains a neural classifier using a contrastive (1-out-of-K) loss to estimate the likelihood-to-evidence ratio. Instead of binary classification, it contrasts one sample from the joint \(p(\theta, x)\) against \(K-1\) samples from the marginals \(p(\theta)p(x)\). This multi-class formulation improves training stability compared to NRE-A.

NRE can be run multi-round without need for correction, but requires running potentially expensive posterior sampling in each round.

[1] On Contrastive Learning for Likelihood-free Inference, Durkan et al., ICML 2020, https://arxiv.org/pdf/2002.03712

Example:

::

import torch
from sbi.inference import NRE_B
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train ratio estimator with contrastive loss
inference = NRE_B(prior=prior)
ratio_estimator = inference.append_simulations(theta, x).train(num_atoms=10)

# 3. Build posterior
posterior = inference.build_posterior(ratio_estimator)

# 4. Sample from posterior
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/nre/nre_b.py

class NRE_B(RatioEstimatorTrainer):
    r"""Neural Ratio Estimation (NRE-B / SRE) as in Durkan et al. (2020) [1].

    NRE-B is an extension of NRE-A that trains a neural classifier using a contrastive
    (1-out-of-K) loss to estimate the likelihood-to-evidence ratio. Instead of binary
    classification, it contrasts one sample from the joint $p(\theta, x)$ against $K-1$
    samples from the marginals $p(\theta)p(x)$. This multi-class formulation improves
    training stability compared to NRE-A.

    NRE can be run multi-round without need for correction, but requires running
    potentially expensive posterior sampling in each round.

    [1] *On Contrastive Learning for Likelihood-free Inference*, Durkan et al.,
        ICML 2020, https://arxiv.org/pdf/2002.03712

    Example:
    --------

    ::

        import torch
        from sbi.inference import NRE_B
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train ratio estimator with contrastive loss
        inference = NRE_B(prior=prior)
        ratio_estimator = inference.append_simulations(theta, x).train(num_atoms=10)

        # 3. Build posterior
        posterior = inference.build_posterior(ratio_estimator)

        # 4. Sample from posterior
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
        device: str = "cpu",
        logging_level: Union[int, str] = "warning",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize NRE_B.

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. If `None`, the
                prior must be passed to `.build_posterior()`.
            classifier: Classifier trained to approximate likelihood ratios. If it is
                a string, use a pre-configured network of the provided type (one of
                linear, mlp, resnet), or a callable that implements the
                `ConditionalEstimatorBuilder` protocol. The callable will
                be called with the first batch of simulations (theta, x), which can thus
                be used for shape inference and potentially for z-scoring. It returns a
                `RatioEstimator`.
            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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during simulation and
                sampling.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

    def train(
        self,
        num_atoms: int = 10,
        training_batch_size: int = 200,
        learning_rate: float = 5e-4,
        validation_fraction: float = 0.1,
        stop_after_epochs: int = 20,
        max_num_epochs: int = 2**31 - 1,
        clip_max_norm: Optional[float] = 5.0,
        resume_training: bool = False,
        discard_prior_samples: bool = False,
        retrain_from_scratch: bool = False,
        show_train_summary: bool = False,
        dataloader_kwargs: Optional[Dict] = None,
    ) -> RatioEstimator:
        r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.

        Args:
            num_atoms: Number of atoms to use for classification.
            training_batch_size: Training batch size.
            learning_rate: Learning rate for Adam optimizer.
            validation_fraction: The fraction of data to use for validation.
            stop_after_epochs: The number of epochs to wait for improvement on the
                validation set before terminating training.
            max_num_epochs: Maximum number of epochs to run. If reached, we stop
                training even when the validation loss is still decreasing. Otherwise,
                we train until validation loss increases (see also `stop_after_epochs`).
            clip_max_norm: Value at which to clip the total gradient norm in order to
                prevent exploding gradients. Use None for no clipping.
            resume_training: Can be used in case training time is limited, e.g. on a
                cluster. If `True`, the split between train and validation set, the
                optimizer, the number of epochs, and the best validation log-prob will
                be restored from the last time `.train()` was called.
            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                from the prior. Training may be sped up by ignoring such less targeted
                samples.
            retrain_from_scratch: Whether to retrain the conditional density
                estimator for the posterior from scratch each round.
            show_train_summary: Whether to print the number of epochs and validation
                loss and leakage after the training.
            dataloader_kwargs: Additional or updated kwargs to be passed to the training
                and validation dataloaders (like, e.g., a collate_fn)

        Returns:
            Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
        """
        kwargs = del_entries(locals(), entries=("self", "__class__"))
        kwargs["loss_kwargs"] = LossArgsNRE(num_atoms=kwargs.pop("num_atoms"))

        return super().train(**kwargs)

    def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:
        r"""Return cross-entropy (via softmax activation) loss for 1-out-of-`num_atoms`
        classification.

        The classifier takes as input `num_atoms` $(\theta,x)$ pairs. Out of these
        pairs, one pair was sampled from the joint $p(\theta,x)$ and all others from the
        marginals $p(\theta)p(x)$. The classifier is trained to predict which of the
        pairs was sampled from the joint $p(\theta,x)$.
        """

        assert theta.shape[0] == x.shape[0], "Batch sizes for theta and x must match."
        batch_size = theta.shape[0]
        logits = self._classifier_logits(theta, x, num_atoms)

        # For 1-out-of-`num_atoms` classification each datapoint consists
        # of `num_atoms` points, with one of them being the correct one.
        # We have a batch of `batch_size` such datapoints.
        logits = logits.reshape(batch_size, num_atoms)

        # Index 0 is the theta-x-pair sampled from the joint p(theta,x) and hence the
        # "correct" one for the 1-out-of-N classification.
        log_prob = logits[:, 0] - torch.logsumexp(logits, dim=-1)

        loss = -torch.mean(log_prob)
        assert_all_finite(loss, "NRE-B loss")
        return loss

__init__(prior=None, classifier='resnet', device='cpu', logging_level='warning', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize NRE_B.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. If None, the prior must be passed to .build_posterior().	None
classifier	Union[str, ConditionalEstimatorBuilder[RatioEstimator]]	Classifier trained to approximate likelihood ratios. If it is a string, use a pre-configured network of the provided type (one of linear, mlp, resnet), or a callable that implements the ConditionalEstimatorBuilder protocol. The callable will be called with the first batch of simulations (theta, x), which can thus be used for shape inference and potentially for z-scoring. It returns a RatioEstimator.	'resnet'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'warning'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during simulation and sampling.	True

Source code in sbi/inference/trainers/nre/nre_b.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
    device: str = "cpu",
    logging_level: Union[int, str] = "warning",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize NRE_B.

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. If `None`, the
            prior must be passed to `.build_posterior()`.
        classifier: Classifier trained to approximate likelihood ratios. If it is
            a string, use a pre-configured network of the provided type (one of
            linear, mlp, resnet), or a callable that implements the
            `ConditionalEstimatorBuilder` protocol. The callable will
            be called with the first batch of simulations (theta, x), which can thus
            be used for shape inference and potentially for z-scoring. It returns a
            `RatioEstimator`.
        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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during simulation and
            sampling.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

train(num_atoms=10, training_batch_size=200, learning_rate=0.0005, validation_fraction=0.1, stop_after_epochs=20, max_num_epochs=2 ** 31 - 1, clip_max_norm=5.0, resume_training=False, discard_prior_samples=False, retrain_from_scratch=False, show_train_summary=False, dataloader_kwargs=None)

Return classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Parameters:

Name	Type	Description	Default
num_atoms	int	Number of atoms to use for classification.	10
training_batch_size	int	Training batch size.	200
learning_rate	float	Learning rate for Adam optimizer.	0.0005
validation_fraction	float	The fraction of data to use for validation.	0.1
stop_after_epochs	int	The number of epochs to wait for improvement on the validation set before terminating training.	20
max_num_epochs	int	Maximum number of epochs to run. If reached, we stop training even when the validation loss is still decreasing. Otherwise, we train until validation loss increases (see also stop_after_epochs).	2 ** 31 - 1
clip_max_norm	Optional[float]	Value at which to clip the total gradient norm in order to prevent exploding gradients. Use None for no clipping.	5.0
resume_training	bool	Can be used in case training time is limited, e.g. on a cluster. If True, the split between train and validation set, the optimizer, the number of epochs, and the best validation log-prob will be restored from the last time .train() was called.	False
discard_prior_samples	bool	Whether to discard samples simulated in round 1, i.e. from the prior. Training may be sped up by ignoring such less targeted samples.	False
retrain_from_scratch	bool	Whether to retrain the conditional density estimator for the posterior from scratch each round.	False
show_train_summary	bool	Whether to print the number of epochs and validation loss and leakage after the training.	False
dataloader_kwargs	Optional[Dict]	Additional or updated kwargs to be passed to the training and validation dataloaders (like, e.g., a collate_fn)	None

Returns:

Type	Description
RatioEstimator	Classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Source code in sbi/inference/trainers/nre/nre_b.py

def train(
    self,
    num_atoms: int = 10,
    training_batch_size: int = 200,
    learning_rate: float = 5e-4,
    validation_fraction: float = 0.1,
    stop_after_epochs: int = 20,
    max_num_epochs: int = 2**31 - 1,
    clip_max_norm: Optional[float] = 5.0,
    resume_training: bool = False,
    discard_prior_samples: bool = False,
    retrain_from_scratch: bool = False,
    show_train_summary: bool = False,
    dataloader_kwargs: Optional[Dict] = None,
) -> RatioEstimator:
    r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.

    Args:
        num_atoms: Number of atoms to use for classification.
        training_batch_size: Training batch size.
        learning_rate: Learning rate for Adam optimizer.
        validation_fraction: The fraction of data to use for validation.
        stop_after_epochs: The number of epochs to wait for improvement on the
            validation set before terminating training.
        max_num_epochs: Maximum number of epochs to run. If reached, we stop
            training even when the validation loss is still decreasing. Otherwise,
            we train until validation loss increases (see also `stop_after_epochs`).
        clip_max_norm: Value at which to clip the total gradient norm in order to
            prevent exploding gradients. Use None for no clipping.
        resume_training: Can be used in case training time is limited, e.g. on a
            cluster. If `True`, the split between train and validation set, the
            optimizer, the number of epochs, and the best validation log-prob will
            be restored from the last time `.train()` was called.
        discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
            from the prior. Training may be sped up by ignoring such less targeted
            samples.
        retrain_from_scratch: Whether to retrain the conditional density
            estimator for the posterior from scratch each round.
        show_train_summary: Whether to print the number of epochs and validation
            loss and leakage after the training.
        dataloader_kwargs: Additional or updated kwargs to be passed to the training
            and validation dataloaders (like, e.g., a collate_fn)

    Returns:
        Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
    """
    kwargs = del_entries(locals(), entries=("self", "__class__"))
    kwargs["loss_kwargs"] = LossArgsNRE(num_atoms=kwargs.pop("num_atoms"))

    return super().train(**kwargs)

NRE_C

Bases: RatioEstimatorTrainer

Neural Ratio Estimation (NRE-C / CNRE) as in Miller et al. (2022) [1].

NRE-C generalizes NRE-A and NRE-B using a “multi-class sigmoid” loss that ensures the estimated ratio \(p(\theta,x)/p(\theta)p(x)\) is exact at optimum in the first round. This addresses the issue that NRE-B’s ratio is only defined up to an arbitrary function of \(x\). NRE-C provides more accurate ratio estimates while maintaining the benefits of contrastive learning.

[1] Contrastive Neural Ratio Estimation, Benjamin Kurt Miller, et al., NeurIPS 2022, https://arxiv.org/abs/2210.06170

Example:

::

import torch
from sbi.inference import NRE_C
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train ratio estimator
inference = NRE_C(prior=prior)
ratio_estimator = inference.append_simulations(theta, x).train(num_classes=5)

# 3. Build posterior
posterior = inference.build_posterior(ratio_estimator)

# 4. Sample from posterior
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/nre/nre_c.py

class NRE_C(RatioEstimatorTrainer):
    r"""Neural Ratio Estimation (NRE-C / CNRE) as in Miller et al. (2022) [1].

    NRE-C generalizes NRE-A and NRE-B using a "multi-class sigmoid" loss that
    ensures the estimated ratio $p(\theta,x)/p(\theta)p(x)$ is exact at optimum
    in the first round. This addresses the issue that NRE-B's ratio is only defined
    up to an arbitrary function of $x$. NRE-C provides more accurate ratio estimates
    while maintaining the benefits of contrastive learning.

    [1] *Contrastive Neural Ratio Estimation*, Benjamin Kurt Miller, et al.,
        NeurIPS 2022, https://arxiv.org/abs/2210.06170

    Example:
    --------

    ::

        import torch
        from sbi.inference import NRE_C
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train ratio estimator
        inference = NRE_C(prior=prior)
        ratio_estimator = inference.append_simulations(theta, x).train(num_classes=5)

        # 3. Build posterior
        posterior = inference.build_posterior(ratio_estimator)

        # 4. Sample from posterior
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
        device: str = "cpu",
        logging_level: Union[int, str] = "warning",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Initialize NRE-C.

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. If `None`, the
                prior must be passed to `.build_posterior()`.
            classifier: Classifier trained to approximate likelihood ratios. If it is
                a string, use a pre-configured network of the provided type (one of
                linear, mlp, resnet), or a callable that implements the
                `ConditionalEstimatorBuilder` protocol. The callable will
                be called with the first batch of simulations (theta, x), which can thus
                be used for shape inference and potentially for z-scoring. It returns a
                `RatioEstimator`.
            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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during simulation and
                sampling.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

    def train(
        self,
        num_classes: int = 5,
        gamma: float = 1.0,
        training_batch_size: int = 200,
        learning_rate: float = 5e-4,
        validation_fraction: float = 0.1,
        stop_after_epochs: int = 20,
        max_num_epochs: int = 2**31 - 1,
        clip_max_norm: Optional[float] = 5.0,
        resume_training: bool = False,
        discard_prior_samples: bool = False,
        retrain_from_scratch: bool = False,
        show_train_summary: bool = False,
        dataloader_kwargs: Optional[Dict] = None,
    ) -> RatioEstimator:
        r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.

        Args:
            num_classes: Number of theta to classify against, corresponds to $K$ in
                _Contrastive Neural Ratio Estimation_. Minimum value is 1. Similar to
                `num_atoms` for SNRE_B except SNRE_C has an additional independently
                drawn sample. The total number of alternative parameters `NRE-C` "sees"
                is $2K-1$ or `2 * num_classes - 1` divided between two loss terms.
            gamma: Determines the relative weight of the sum of all $K$ dependently
                drawn classes against the marginally drawn one. Specifically,
                $p(y=k) :=p_K$, $p(y=0) := p_0$, $p_0 = 1 - K p_K$, and finally
                $\gamma := K p_K / p_0$.
            training_batch_size: Training batch size.
            learning_rate: Learning rate for Adam optimizer.
            validation_fraction: The fraction of data to use for validation.
            stop_after_epochs: The number of epochs to wait for improvement on the
                validation set before terminating training.
            max_num_epochs: Maximum number of epochs to run. If reached, we stop
                training even when the validation loss is still decreasing. Otherwise,
                we train until validation loss increases (see also `stop_after_epochs`).
            clip_max_norm: Value at which to clip the total gradient norm in order to
                prevent exploding gradients. Use None for no clipping.
            exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
                during training. Expect errors, silent or explicit, when `False`.
            resume_training: Can be used in case training time is limited, e.g. on a
                cluster. If `True`, the split between train and validation set, the
                optimizer, the number of epochs, and the best validation log-prob will
                be restored from the last time `.train()` was called.
            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                from the prior. Training may be sped up by ignoring such less targeted
                samples.
            retrain_from_scratch: Whether to retrain the conditional density
                estimator for the posterior from scratch each round.
            show_train_summary: Whether to print the number of epochs and validation
                loss and leakage after the training.
            dataloader_kwargs: Additional or updated kwargs to be passed to the training
                and validation dataloaders (like, e.g., a collate_fn)

        Returns:
            Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        kwargs["loss_kwargs"] = LossArgsNRE_C(
            num_atoms=kwargs.pop("num_classes") + 1, gamma=kwargs.pop("gamma")
        )

        return super().train(**kwargs)

    def _loss(
        self, theta: Tensor, x: Tensor, num_atoms: int, gamma: float
    ) -> torch.Tensor:
        r"""Return cross-entropy loss (via ''multi-class sigmoid'' activation) for
        1-out-of-`K + 1` classification.

        At optimum, this loss function returns the exact likelihood-to-evidence ratio
        in the first round.
        Details of loss computation are described in Contrastive Neural Ratio
        Estimation[1]. The paper does not discuss the sequential case.

        [1] _Contrastive Neural Ratio Estimation_, Benajmin Kurt Miller, et. al.,
            NeurIPS 2022, https://arxiv.org/abs/2210.06170
        """

        # Reminder: K = num_classes
        # The algorithm is written with K, so we convert back to K format rather than
        # reasoning in num_atoms.
        num_classes = num_atoms - 1
        assert num_classes >= 1, f"num_classes = {num_classes} must be greater than 1."

        assert theta.shape[0] == x.shape[0], "Batch sizes for theta and x must match."
        batch_size = theta.shape[0]

        # We append a contrastive theta to the marginal case because we will remove
        # the jointly drawn
        # sample in the logits_marginal[:, 0] position. That makes the remaining sample
        # marginally drawn.
        # We have a batch of `batch_size` datapoints.
        logits_marginal = self._classifier_logits(theta, x, num_classes + 1).reshape(
            batch_size, num_classes + 1
        )
        logits_joint = self._classifier_logits(theta, x, num_classes).reshape(
            batch_size, num_classes
        )

        dtype = logits_marginal.dtype
        device = logits_marginal.device

        # Index 0 is the theta-x-pair sampled from the joint p(theta,x) and hence
        # we remove the jointly drawn sample from the logits_marginal
        logits_marginal = logits_marginal[:, 1:]
        # ... and retain it in the logits_joint. Now we have two arrays with K choices.

        # To use logsumexp, we extend the denominator logits with loggamma
        loggamma = torch.tensor(gamma, dtype=dtype, device=device).log()
        logK = torch.tensor(num_classes, dtype=dtype, device=device).log()
        denominator_marginal = torch.concat(
            [loggamma + logits_marginal, logK.expand((batch_size, 1))],
            dim=-1,
        )
        denominator_joint = torch.concat(
            [loggamma + logits_joint, logK.expand((batch_size, 1))],
            dim=-1,
        )

        # Compute the contributions to the loss from each term in the classification.
        log_prob_marginal = logK - torch.logsumexp(denominator_marginal, dim=-1)
        log_prob_joint = (
            loggamma + logits_joint[:, 0] - torch.logsumexp(denominator_joint, dim=-1)
        )

        # relative weights. p_marginal := p_0, and p_joint := p_K * K from the notation.
        p_marginal, p_joint = self._get_prior_probs_marginal_and_joint(gamma)

        loss = -torch.mean(p_marginal * log_prob_marginal + p_joint * log_prob_joint)
        assert_all_finite(loss, "NRE-C loss")
        return loss

    @staticmethod
    def _get_prior_probs_marginal_and_joint(gamma: float) -> Tuple[float, float]:
        r"""Return a tuple (p_marginal, p_joint) where `p_marginal := `$p_0$,
        `p_joint := `$p_K \cdot K$.

        We let the joint (dependently drawn) class to be equally likely across K
        options. The marginal class is therefore restricted to get the remaining
        probability.
        """
        p_joint = gamma / (1 + gamma)
        p_marginal = 1 / (1 + gamma)
        return p_marginal, p_joint

    def _get_losses(self, batch: Sequence[Tensor], loss_args: LossArgs) -> Tensor:
        """Overrides the parent class method to check the type of loss_args."""

        if not isinstance(loss_args, LossArgsNRE_C):
            raise TypeError(
                "Expected type of loss_args to be LossArgsNRE_C,"
                f" but got {type(loss_args)}"
            )

        return super()._get_losses(batch=batch, loss_args=loss_args)

__init__(prior=None, classifier='resnet', device='cpu', logging_level='warning', summary_writer=None, tracker=None, show_progress_bars=True)

Initialize NRE-C.

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. If None, the prior must be passed to .build_posterior().	None
classifier	Union[str, ConditionalEstimatorBuilder[RatioEstimator]]	Classifier trained to approximate likelihood ratios. If it is a string, use a pre-configured network of the provided type (one of linear, mlp, resnet), or a callable that implements the ConditionalEstimatorBuilder protocol. The callable will be called with the first batch of simulations (theta, x), which can thus be used for shape inference and potentially for z-scoring. It returns a RatioEstimator.	'resnet'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'warning'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during simulation and sampling.	True

Source code in sbi/inference/trainers/nre/nre_c.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
    device: str = "cpu",
    logging_level: Union[int, str] = "warning",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Initialize NRE-C.

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. If `None`, the
            prior must be passed to `.build_posterior()`.
        classifier: Classifier trained to approximate likelihood ratios. If it is
            a string, use a pre-configured network of the provided type (one of
            linear, mlp, resnet), or a callable that implements the
            `ConditionalEstimatorBuilder` protocol. The callable will
            be called with the first batch of simulations (theta, x), which can thus
            be used for shape inference and potentially for z-scoring. It returns a
            `RatioEstimator`.
        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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during simulation and
            sampling.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

train(num_classes=5, gamma=1.0, training_batch_size=200, learning_rate=0.0005, validation_fraction=0.1, stop_after_epochs=20, max_num_epochs=2 ** 31 - 1, clip_max_norm=5.0, resume_training=False, discard_prior_samples=False, retrain_from_scratch=False, show_train_summary=False, dataloader_kwargs=None)

Return classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Parameters:

Name	Type	Description	Default
num_classes	int	Number of theta to classify against, corresponds to \(K\) in Contrastive Neural Ratio Estimation. Minimum value is 1. Similar to num_atoms for SNRE_B except SNRE_C has an additional independently drawn sample. The total number of alternative parameters NRE-C “sees” is \(2K-1\) or 2 * num_classes - 1 divided between two loss terms.	5
gamma	float	Determines the relative weight of the sum of all \(K\) dependently drawn classes against the marginally drawn one. Specifically, \(p(y=k) :=p_K\), \(p(y=0) := p_0\), \(p_0 = 1 - K p_K\), and finally \(\gamma := K p_K / p_0\).	1.0
training_batch_size	int	Training batch size.	200
learning_rate	float	Learning rate for Adam optimizer.	0.0005
validation_fraction	float	The fraction of data to use for validation.	0.1
stop_after_epochs	int	The number of epochs to wait for improvement on the validation set before terminating training.	20
max_num_epochs	int	Maximum number of epochs to run. If reached, we stop training even when the validation loss is still decreasing. Otherwise, we train until validation loss increases (see also stop_after_epochs).	2 ** 31 - 1
clip_max_norm	Optional[float]	Value at which to clip the total gradient norm in order to prevent exploding gradients. Use None for no clipping.	5.0
exclude_invalid_x		Whether to exclude simulation outputs x=NaN or x=±∞ during training. Expect errors, silent or explicit, when False.	required
resume_training	bool	Can be used in case training time is limited, e.g. on a cluster. If True, the split between train and validation set, the optimizer, the number of epochs, and the best validation log-prob will be restored from the last time .train() was called.	False
discard_prior_samples	bool	Whether to discard samples simulated in round 1, i.e. from the prior. Training may be sped up by ignoring such less targeted samples.	False
retrain_from_scratch	bool	Whether to retrain the conditional density estimator for the posterior from scratch each round.	False
show_train_summary	bool	Whether to print the number of epochs and validation loss and leakage after the training.	False
dataloader_kwargs	Optional[Dict]	Additional or updated kwargs to be passed to the training and validation dataloaders (like, e.g., a collate_fn)	None

Returns:

Type	Description
RatioEstimator	Classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Source code in sbi/inference/trainers/nre/nre_c.py

def train(
    self,
    num_classes: int = 5,
    gamma: float = 1.0,
    training_batch_size: int = 200,
    learning_rate: float = 5e-4,
    validation_fraction: float = 0.1,
    stop_after_epochs: int = 20,
    max_num_epochs: int = 2**31 - 1,
    clip_max_norm: Optional[float] = 5.0,
    resume_training: bool = False,
    discard_prior_samples: bool = False,
    retrain_from_scratch: bool = False,
    show_train_summary: bool = False,
    dataloader_kwargs: Optional[Dict] = None,
) -> RatioEstimator:
    r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.

    Args:
        num_classes: Number of theta to classify against, corresponds to $K$ in
            _Contrastive Neural Ratio Estimation_. Minimum value is 1. Similar to
            `num_atoms` for SNRE_B except SNRE_C has an additional independently
            drawn sample. The total number of alternative parameters `NRE-C` "sees"
            is $2K-1$ or `2 * num_classes - 1` divided between two loss terms.
        gamma: Determines the relative weight of the sum of all $K$ dependently
            drawn classes against the marginally drawn one. Specifically,
            $p(y=k) :=p_K$, $p(y=0) := p_0$, $p_0 = 1 - K p_K$, and finally
            $\gamma := K p_K / p_0$.
        training_batch_size: Training batch size.
        learning_rate: Learning rate for Adam optimizer.
        validation_fraction: The fraction of data to use for validation.
        stop_after_epochs: The number of epochs to wait for improvement on the
            validation set before terminating training.
        max_num_epochs: Maximum number of epochs to run. If reached, we stop
            training even when the validation loss is still decreasing. Otherwise,
            we train until validation loss increases (see also `stop_after_epochs`).
        clip_max_norm: Value at which to clip the total gradient norm in order to
            prevent exploding gradients. Use None for no clipping.
        exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
            during training. Expect errors, silent or explicit, when `False`.
        resume_training: Can be used in case training time is limited, e.g. on a
            cluster. If `True`, the split between train and validation set, the
            optimizer, the number of epochs, and the best validation log-prob will
            be restored from the last time `.train()` was called.
        discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
            from the prior. Training may be sped up by ignoring such less targeted
            samples.
        retrain_from_scratch: Whether to retrain the conditional density
            estimator for the posterior from scratch each round.
        show_train_summary: Whether to print the number of epochs and validation
            loss and leakage after the training.
        dataloader_kwargs: Additional or updated kwargs to be passed to the training
            and validation dataloaders (like, e.g., a collate_fn)

    Returns:
        Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    kwargs["loss_kwargs"] = LossArgsNRE_C(
        num_atoms=kwargs.pop("num_classes") + 1, gamma=kwargs.pop("gamma")
    )

    return super().train(**kwargs)

BNRE

Bases: NRE_A

Balanced Neural Ratio Estimation (BNRE) as in Delaunoy et al. (2022) [1].

BNRE is a variation of NRE-A that adds a balancing regularizer to the binary cross-entropy loss. This regularizer encourages the classifier to predict equal probabilities for joint and marginal samples on average, which can lead to more conservative and reliable posterior approximations. BNRE is particularly useful when robustness is prioritized over tightness of the posterior.

NRE can be run multi-round without need for correction, but requires running potentially expensive posterior sampling in each round.

[1] Towards Reliable Simulation-Based Inference with Balanced Neural Ratio Estimation, Delaunoy, A., Hermans, J., Rozet, F., Wehenkel, A., & Louppe, G., NeurIPS 2022. https://arxiv.org/abs/2208.13624

Example:

::

import torch
from sbi.inference import BNRE
from sbi.utils import BoxUniform

# 1. Setup prior and simulate data
prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
theta = prior.sample((100,))
x = theta + torch.randn_like(theta) * 0.1

# 2. Train balanced ratio estimator
inference = BNRE(prior=prior)
# Note: regularization_strength needs to be tuned carefully for your problem
ratio_estimator = inference.append_simulations(theta, x).train(
    regularization_strength=100.0
)

# 3. Build posterior
posterior = inference.build_posterior(ratio_estimator)

# 4. Sample from posterior
x_o = torch.randn(1, 3)
samples = posterior.sample((1000,), x=x_o)

Source code in sbi/inference/trainers/nre/bnre.py

class BNRE(NRE_A):
    r"""Balanced Neural Ratio Estimation (BNRE) as in Delaunoy et al. (2022) [1].

    BNRE is a variation of NRE-A that adds a balancing regularizer to the binary
    cross-entropy loss. This regularizer encourages the classifier to predict equal
    probabilities for joint and marginal samples on average, which can lead to more
    conservative and reliable posterior approximations. BNRE is particularly useful
    when robustness is prioritized over tightness of the posterior.

    NRE can be run multi-round without need for correction, but requires running
    potentially expensive posterior sampling in each round.

    [1] Towards Reliable Simulation-Based Inference with Balanced Neural Ratio
        Estimation, Delaunoy, A., Hermans, J., Rozet, F., Wehenkel, A., & Louppe, G.,
        NeurIPS 2022. https://arxiv.org/abs/2208.13624

    Example:
    --------

    ::

        import torch
        from sbi.inference import BNRE
        from sbi.utils import BoxUniform

        # 1. Setup prior and simulate data
        prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3))
        theta = prior.sample((100,))
        x = theta + torch.randn_like(theta) * 0.1

        # 2. Train balanced ratio estimator
        inference = BNRE(prior=prior)
        # Note: regularization_strength needs to be tuned carefully for your problem
        ratio_estimator = inference.append_simulations(theta, x).train(
            regularization_strength=100.0
        )

        # 3. Build posterior
        posterior = inference.build_posterior(ratio_estimator)

        # 4. Sample from posterior
        x_o = torch.randn(1, 3)
        samples = posterior.sample((1000,), x=x_o)
    """

    def __init__(
        self,
        prior: Optional[Distribution] = None,
        classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
        device: str = "cpu",
        logging_level: Union[int, str] = "warning",
        summary_writer: Optional[SummaryWriter] = None,
        tracker: Optional[Tracker] = None,
        show_progress_bars: bool = True,
    ):
        r"""Balanced neural ratio estimation (BNRE).

        Args:
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. If `None`, the
                prior must be passed to `.build_posterior()`.
            classifier: Classifier trained to approximate likelihood ratios. If it is
                a string, use a pre-configured network of the provided type (one of
                linear, mlp, resnet), or a callable that implements the
                `ConditionalEstimatorBuilder` protocol. The callable will
                be called with the first batch of simulations (theta, x), which can thus
                be used for shape inference and potentially for z-scoring. It returns a
                `RatioEstimator`.
            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.
            summary_writer: Deprecated alias for the TensorBoard summary writer.
                Use ``tracker`` instead.
            tracker: Tracking adapter used to log training metrics. If None, a
                TensorBoard tracker is used with a default log directory.
            show_progress_bars: Whether to show a progressbar during simulation and
                sampling.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))
        super().__init__(**kwargs)

    def train(
        self,
        regularization_strength: float = 100.0,
        training_batch_size: int = 200,
        learning_rate: float = 5e-4,
        validation_fraction: float = 0.1,
        stop_after_epochs: int = 20,
        max_num_epochs: int = 2**31 - 1,
        clip_max_norm: Optional[float] = 5.0,
        resume_training: bool = False,
        discard_prior_samples: bool = False,
        retrain_from_scratch: bool = False,
        show_train_summary: bool = False,
        dataloader_kwargs: Optional[Dict] = None,
    ) -> RatioEstimator:
        r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
        Args:

            regularization_strength: The multiplicative coefficient applied to the
                balancing regularizer ($\lambda$).
            training_batch_size: Training batch size.
            learning_rate: Learning rate for Adam optimizer.
            validation_fraction: The fraction of data to use for validation.
            stop_after_epochs: The number of epochs to wait for improvement on the
                validation set before terminating training.
            max_num_epochs: Maximum number of epochs to run. If reached, we stop
                training even when the validation loss is still decreasing. Otherwise,
                we train until validation loss increases (see also `stop_after_epochs`).
            clip_max_norm: Value at which to clip the total gradient norm in order to
                prevent exploding gradients. Use None for no clipping.
            exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
                during training. Expect errors, silent or explicit, when `False`.
            resume_training: Can be used in case training time is limited, e.g. on a
                cluster. If `True`, the split between train and validation set, the
                optimizer, the number of epochs, and the best validation log-prob will
                be restored from the last time `.train()` was called.
            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                from the prior. Training may be sped up by ignoring such less targeted
                samples.
            retrain_from_scratch: Whether to retrain the conditional density
                estimator for the posterior from scratch each round.
            show_train_summary: Whether to print the number of epochs and validation
                loss and leakage after the training.
            dataloader_kwargs: Additional or updated kwargs to be passed to the training
                and validation dataloaders (like, e.g., a collate_fn)
        Returns:
            Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
        """

        kwargs = del_entries(locals(), entries=("self", "__class__"))

        # Configure _loss function parameters by initializing LossArgsBNRE
        # with the given regularization strength.
        kwargs["loss_kwargs"] = LossArgsBNRE(
            regularization_strength=kwargs.pop("regularization_strength"),
        )

        return super().train(**kwargs)

    def _loss(
        self, theta: Tensor, x: Tensor, num_atoms: int, regularization_strength: float
    ) -> Tensor:
        """Returns the binary cross-entropy loss for the trained classifier.

        The classifier takes as input a $(\theta,x)$ pair. It is trained to predict 1
        if the pair was sampled from the joint $p(\theta,x)$, and to predict 0 if the
        pair was sampled from the marginals $p(\theta)p(x)$.
        """

        assert theta.shape[0] == x.shape[0], "Batch sizes for theta and x must match."
        batch_size = theta.shape[0]

        logits = self._classifier_logits(theta, x, num_atoms)
        likelihood = torch.sigmoid(logits).squeeze()

        # Alternating pairs where there is one sampled from the joint and one
        # sampled from the marginals. The first element is sampled from the
        # joint p(theta, x) and is labelled 1. The second element is sampled
        # from the marginals p(theta)p(x) and is labelled 0. And so on.
        labels = ones(2 * batch_size, device=self._device)  # two atoms
        labels[1::2] = 0.0

        # Binary cross entropy to learn the likelihood (AALR-specific)
        bce = nn.BCELoss()(likelihood, labels)

        # Balancing regularizer
        regularizer = (
            (torch.sigmoid(logits[0::2]) + torch.sigmoid(logits[1::2]) - 1)
            .mean()
            .square()
        )

        loss = bce + regularization_strength * regularizer
        assert_all_finite(loss, "BNRE loss")
        return loss

    def _get_losses(self, batch: Sequence[Tensor], loss_args: LossArgs) -> Tensor:
        """Overrides the parent class method to check the type of loss_args."""

        if not isinstance(loss_args, LossArgsBNRE):
            raise TypeError(
                "Expected type of loss_args to be LossArgsBNRE,"
                f" but got {type(loss_args)}"
            )

        return super()._get_losses(batch=batch, loss_args=loss_args)

__init__(prior=None, classifier='resnet', device='cpu', logging_level='warning', summary_writer=None, tracker=None, show_progress_bars=True)

Balanced neural ratio estimation (BNRE).

Parameters:

Name	Type	Description	Default
prior	Optional[Distribution]	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. If None, the prior must be passed to .build_posterior().	None
classifier	Union[str, ConditionalEstimatorBuilder[RatioEstimator]]	Classifier trained to approximate likelihood ratios. If it is a string, use a pre-configured network of the provided type (one of linear, mlp, resnet), or a callable that implements the ConditionalEstimatorBuilder protocol. The callable will be called with the first batch of simulations (theta, x), which can thus be used for shape inference and potentially for z-scoring. It returns a RatioEstimator.	'resnet'
device	str	Training device, e.g., “cpu”, “cuda” or “cuda:{0, 1, …}”.	'cpu'
logging_level	Union[int, str]	Minimum severity of messages to log. One of the strings INFO, WARNING, DEBUG, ERROR and CRITICAL.	'warning'
summary_writer	Optional[SummaryWriter]	Deprecated alias for the TensorBoard summary writer. Use tracker instead.	None
tracker	Optional[Tracker]	Tracking adapter used to log training metrics. If None, a TensorBoard tracker is used with a default log directory.	None
show_progress_bars	bool	Whether to show a progressbar during simulation and sampling.	True

Source code in sbi/inference/trainers/nre/bnre.py

def __init__(
    self,
    prior: Optional[Distribution] = None,
    classifier: Union[str, ConditionalEstimatorBuilder[RatioEstimator]] = "resnet",
    device: str = "cpu",
    logging_level: Union[int, str] = "warning",
    summary_writer: Optional[SummaryWriter] = None,
    tracker: Optional[Tracker] = None,
    show_progress_bars: bool = True,
):
    r"""Balanced neural ratio estimation (BNRE).

    Args:
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. If `None`, the
            prior must be passed to `.build_posterior()`.
        classifier: Classifier trained to approximate likelihood ratios. If it is
            a string, use a pre-configured network of the provided type (one of
            linear, mlp, resnet), or a callable that implements the
            `ConditionalEstimatorBuilder` protocol. The callable will
            be called with the first batch of simulations (theta, x), which can thus
            be used for shape inference and potentially for z-scoring. It returns a
            `RatioEstimator`.
        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.
        summary_writer: Deprecated alias for the TensorBoard summary writer.
            Use ``tracker`` instead.
        tracker: Tracking adapter used to log training metrics. If None, a
            TensorBoard tracker is used with a default log directory.
        show_progress_bars: Whether to show a progressbar during simulation and
            sampling.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))
    super().__init__(**kwargs)

train(regularization_strength=100.0, training_batch_size=200, learning_rate=0.0005, validation_fraction=0.1, stop_after_epochs=20, max_num_epochs=2 ** 31 - 1, clip_max_norm=5.0, resume_training=False, discard_prior_samples=False, retrain_from_scratch=False, show_train_summary=False, dataloader_kwargs=None)

Return classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\). Args:

regularization_strength: The multiplicative coefficient applied to the
    balancing regularizer ($\lambda$).
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
    validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
    training even when the validation loss is still decreasing. Otherwise,
    we train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
    prevent exploding gradients. Use None for no clipping.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
    during training. Expect errors, silent or explicit, when `False`.
resume_training: Can be used in case training time is limited, e.g. on a
    cluster. If `True`, the split between train and validation set, the
    optimizer, the number of epochs, and the best validation log-prob will
    be restored from the last time `.train()` was called.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
    from the prior. Training may be sped up by ignoring such less targeted
    samples.
retrain_from_scratch: Whether to retrain the conditional density
    estimator for the posterior from scratch each round.
show_train_summary: Whether to print the number of epochs and validation
    loss and leakage after the training.
dataloader_kwargs: Additional or updated kwargs to be passed to the training
    and validation dataloaders (like, e.g., a collate_fn)

Returns: Classifier that approximates the ratio \(p(\theta,x)/p(\theta)p(x)\).

Source code in sbi/inference/trainers/nre/bnre.py

def train(
    self,
    regularization_strength: float = 100.0,
    training_batch_size: int = 200,
    learning_rate: float = 5e-4,
    validation_fraction: float = 0.1,
    stop_after_epochs: int = 20,
    max_num_epochs: int = 2**31 - 1,
    clip_max_norm: Optional[float] = 5.0,
    resume_training: bool = False,
    discard_prior_samples: bool = False,
    retrain_from_scratch: bool = False,
    show_train_summary: bool = False,
    dataloader_kwargs: Optional[Dict] = None,
) -> RatioEstimator:
    r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
    Args:

        regularization_strength: The multiplicative coefficient applied to the
            balancing regularizer ($\lambda$).
        training_batch_size: Training batch size.
        learning_rate: Learning rate for Adam optimizer.
        validation_fraction: The fraction of data to use for validation.
        stop_after_epochs: The number of epochs to wait for improvement on the
            validation set before terminating training.
        max_num_epochs: Maximum number of epochs to run. If reached, we stop
            training even when the validation loss is still decreasing. Otherwise,
            we train until validation loss increases (see also `stop_after_epochs`).
        clip_max_norm: Value at which to clip the total gradient norm in order to
            prevent exploding gradients. Use None for no clipping.
        exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
            during training. Expect errors, silent or explicit, when `False`.
        resume_training: Can be used in case training time is limited, e.g. on a
            cluster. If `True`, the split between train and validation set, the
            optimizer, the number of epochs, and the best validation log-prob will
            be restored from the last time `.train()` was called.
        discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
            from the prior. Training may be sped up by ignoring such less targeted
            samples.
        retrain_from_scratch: Whether to retrain the conditional density
            estimator for the posterior from scratch each round.
        show_train_summary: Whether to print the number of epochs and validation
            loss and leakage after the training.
        dataloader_kwargs: Additional or updated kwargs to be passed to the training
            and validation dataloaders (like, e.g., a collate_fn)
    Returns:
        Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
    """

    kwargs = del_entries(locals(), entries=("self", "__class__"))

    # Configure _loss function parameters by initializing LossArgsBNRE
    # with the given regularization strength.
    kwargs["loss_kwargs"] = LossArgsBNRE(
        regularization_strength=kwargs.pop("regularization_strength"),
    )

    return super().train(**kwargs)

MCABC

Bases: ABCBASE

Monte-Carlo Approximate Bayesian Computation (Rejection ABC).

Source code in sbi/inference/abc/mcabc.py

class MCABC(ABCBASE):
    """Monte-Carlo Approximate Bayesian Computation (Rejection ABC)."""

    def __init__(
        self,
        simulator: Callable,
        prior: Distribution,
        distance: Union[str, Callable] = "l2",
        requires_iid_data: Optional[bool] = None,
        distance_kwargs: Optional[Dict] = None,
        num_workers: int = 1,
        simulation_batch_size: int = 1,
        distance_batch_size: int = -1,
        show_progress_bars: bool = True,
    ):
        r"""Monte-Carlo Approximate Bayesian Computation (Rejection ABC) [1].

        [1] Pritchard, J. K., Seielstad, M. T., Perez-Lezaun, A., & Feldman, M. W.
        (1999). Population growth of human Y chromosomes: a study of Y chromosome
        microsatellites. Molecular biology and evolution, 16(12), 1791-1798.

        Args:
            simulator: A function that takes parameters $\theta$ and maps them to
                simulations, or observations, `x`, $\mathrm{sim}(\theta)\to x$. Any
                regular Python callable (i.e. function or class with `__call__` method)
                can be used.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Any
                object with `.log_prob()`and `.sample()` (for example, a PyTorch
                distribution) can be used.
            distance: Distance function to compare observed and simulated data. Can be
                a custom callable function or one of `l1`, `l2`, `mse`,
                `mmd`, `wasserstein`.
            requires_iid_data: Whether to allow conditioning on iid sampled data or not.
                Typically, this information is inferred by the choice of the distance,
                but in case a custom distance is used, this information is pivotal.
            distance_kwargs: Configurations parameters for the distances. In particular
                useful for the MMD and Wasserstein distance.
            num_workers: Number of parallel workers to use for simulations.
            simulation_batch_size: Number of parameter sets that the simulator
                maps to data x at once. If None, we simulate all parameter sets at the
                same time. If >= 1, the simulator has to process data of shape
                (simulation_batch_size, parameter_dimension).
            distance_batch_size: Number of simulations that the distance function
                evaluates against the reference observations at once. If -1, we evaluate
                all simulations at the same time.
        """

        super().__init__(
            simulator=simulator,
            prior=prior,
            distance=distance,
            requires_iid_data=requires_iid_data,
            distance_kwargs=distance_kwargs,
            num_workers=num_workers,
            simulation_batch_size=simulation_batch_size,
            distance_batch_size=distance_batch_size,
            show_progress_bars=show_progress_bars,
        )

    def __call__(
        self,
        x_o: Union[Tensor, ndarray],
        num_simulations: int,
        eps: Optional[float] = None,
        quantile: Optional[float] = None,
        lra: bool = False,
        sass: bool = False,
        sass_fraction: float = 0.25,
        sass_expansion_degree: int = 1,
        kde: bool = False,
        kde_kwargs: Optional[Dict[str, Any]] = None,
        return_summary: bool = False,
        num_iid_samples: int = 1,
    ) -> Union[Tuple[Tensor, dict], Tuple[KDEWrapper, dict], Tensor, KDEWrapper]:
        r"""Run MCABC and return accepted parameters or KDE object fitted on them.

        Args:
            x_o: Observed data.
            num_simulations: Number of simulations to run.
            eps: Acceptance threshold $\epsilon$ for distance between observed and
                simulated data.
            quantile: Upper quantile of smallest distances for which the corresponding
                parameters are returned, e.g, q=0.01 will return the top 1%. Exactly
                one of quantile or `eps` have to be passed.
            lra: Whether to run linear regression adjustment as in Beaumont et al. 2002
            sass: Whether to determine semi-automatic summary statistics as in
                Fearnhead & Prangle 2012.
            sass_fraction: Fraction of simulation budget used for the initial sass run.
            sass_expansion_degree: Degree of the polynomial feature expansion for the
                sass regression, default 1 - no expansion.
            kde: Whether to run KDE on the accepted parameters to return a KDE
                object from which one can sample.
            kde_kwargs: kwargs for performing KDE:
                'bandwidth='; either a float, or a string naming a bandwidth
                heuristics, e.g., 'cv' (cross validation), 'silvermann' or 'scott',
                default 'cv'.
                'transform': transform applied to the parameters before doing KDE.
                'sample_weights': weights associated with samples. See 'get_kde' for
                more details
            return_summary: Whether to return the distances and data corresponding to
                the accepted parameters.
            num_iid_samples: Number of simulations per parameter. Choose
                `num_iid_samples>1`, if you have chosen a statistical distance that
                evaluates sets of simulations against a set of reference observations
                instead of a single data-point comparison.

        Returns:
            theta (if kde False): accepted parameters
            kde (if kde True): KDE object based on accepted parameters from which one
                can .sample() and .log_prob().
            summary (if summary True): dictionary containing the accepted paramters (if
                kde True), distances and simulated data x.
        """

        # Exactly one of eps or quantile need to be passed.
        assert (eps is not None) ^ (quantile is not None), (
            "Eps or quantile must be passed, but not both."
        )
        if kde_kwargs is None:
            kde_kwargs = {}

        # Run SASS and change the simulator and x_o accordingly.
        if sass:
            num_pilot_simulations = int(sass_fraction * num_simulations)
            self.logger.info(
                "Running SASS with %s pilot samples.", num_pilot_simulations
            )
            num_simulations -= num_pilot_simulations

            pilot_theta = self.prior.sample((num_pilot_simulations,))
            pilot_x = self._batched_simulator(pilot_theta)

            sass_transform = super()._get_sass_transform(
                pilot_theta, pilot_x, sass_expansion_degree
            )

            # Add sass transform to simulator and x_o.
            def simulator(theta):
                return sass_transform(self._batched_simulator(theta))

            x_o = sass_transform(process_x(x_o))
        else:
            simulator = self._batched_simulator

        # Simulate and calculate distances.
        theta = self.prior.sample((num_simulations,))
        theta_repeat = theta.repeat_interleave(num_iid_samples, dim=0)
        x = simulator(theta_repeat)
        x = x.reshape((
            num_simulations,
            num_iid_samples,
            -1,
        ))  # Dim(num_initial_pop, num_iid_samples, -1)

        # Infer x shape to test and set x_o.
        if not self.distance.requires_iid_data:
            x = x.squeeze(1)
            self.x_shape = x[0].shape
        else:
            self.x_shape = x[0, 0].shape

        self.x_o = process_x(x_o, self.x_shape)
        distances = self.distance(self.x_o, x)

        # Select based on acceptance threshold epsilon.
        if eps is not None:
            is_accepted = distances < eps
            num_accepted = is_accepted.sum().item()
            assert num_accepted > 0, f"No parameters accepted, eps={eps} too small"

            theta_accepted = theta[is_accepted]
            distances_accepted = distances[is_accepted]
            x_accepted = x[is_accepted]

        # Select based on quantile on sorted distances.
        elif quantile is not None:
            num_top_samples = int(num_simulations * quantile)
            sort_idx = torch.argsort(distances)
            theta_accepted = theta[sort_idx][:num_top_samples]
            distances_accepted = distances[sort_idx][:num_top_samples]
            x_accepted = x[sort_idx][:num_top_samples]

        else:
            raise ValueError("One of epsilon or quantile has to be passed.")

        # Maybe adjust theta with LRA.
        if lra:
            self.logger.info("Running Linear regression adjustment.")
            final_theta = super()._run_lra(
                theta_accepted, x_accepted, observation=self.x_o
            )
        else:
            final_theta = theta_accepted

        if kde:
            self.logger.info(
                "KDE on %s samples with bandwidth option "
                f"{kde_kwargs.get('bandwidth', 'cv')}. "
                "Beware that KDE can give unreliable results when used with too few"
                " samples and in high dimensions.",
                final_theta.shape[0],
            )

            kde_dist = get_kde(final_theta, **kde_kwargs)

            if return_summary:
                return (
                    kde_dist,
                    dict(theta=final_theta, distances=distances_accepted, x=x_accepted),
                )
            else:
                return kde_dist
        elif return_summary:
            return final_theta, dict(distances=distances_accepted, x=x_accepted)
        else:
            return final_theta

__call__(x_o, num_simulations, eps=None, quantile=None, lra=False, sass=False, sass_fraction=0.25, sass_expansion_degree=1, kde=False, kde_kwargs=None, return_summary=False, num_iid_samples=1)

Run MCABC and return accepted parameters or KDE object fitted on them.

Parameters:

Name	Type	Description	Default
x_o	Union[Tensor, ndarray]	Observed data.	required
num_simulations	int	Number of simulations to run.	required
eps	Optional[float]	Acceptance threshold \(\epsilon\) for distance between observed and simulated data.	None
quantile	Optional[float]	Upper quantile of smallest distances for which the corresponding parameters are returned, e.g, q=0.01 will return the top 1%. Exactly one of quantile or eps have to be passed.	None
lra	bool	Whether to run linear regression adjustment as in Beaumont et al. 2002	False
sass	bool	Whether to determine semi-automatic summary statistics as in Fearnhead & Prangle 2012.	False
sass_fraction	float	Fraction of simulation budget used for the initial sass run.	0.25
sass_expansion_degree	int	Degree of the polynomial feature expansion for the sass regression, default 1 - no expansion.	1
kde	bool	Whether to run KDE on the accepted parameters to return a KDE object from which one can sample.	False
kde_kwargs	Optional[Dict[str, Any]]	kwargs for performing KDE: ‘bandwidth=’; either a float, or a string naming a bandwidth heuristics, e.g., ‘cv’ (cross validation), ‘silvermann’ or ‘scott’, default ‘cv’. ‘transform’: transform applied to the parameters before doing KDE. ‘sample_weights’: weights associated with samples. See ‘get_kde’ for more details	None
return_summary	bool	Whether to return the distances and data corresponding to the accepted parameters.	False
num_iid_samples	int	Number of simulations per parameter. Choose num_iid_samples>1, if you have chosen a statistical distance that evaluates sets of simulations against a set of reference observations instead of a single data-point comparison.	1

Returns:

Name	Type	Description
theta	if kde False	accepted parameters
kde	if kde True	KDE object based on accepted parameters from which one can .sample() and .log_prob().
summary	if summary True	dictionary containing the accepted paramters (if kde True), distances and simulated data x.

Source code in sbi/inference/abc/mcabc.py

def __call__(
    self,
    x_o: Union[Tensor, ndarray],
    num_simulations: int,
    eps: Optional[float] = None,
    quantile: Optional[float] = None,
    lra: bool = False,
    sass: bool = False,
    sass_fraction: float = 0.25,
    sass_expansion_degree: int = 1,
    kde: bool = False,
    kde_kwargs: Optional[Dict[str, Any]] = None,
    return_summary: bool = False,
    num_iid_samples: int = 1,
) -> Union[Tuple[Tensor, dict], Tuple[KDEWrapper, dict], Tensor, KDEWrapper]:
    r"""Run MCABC and return accepted parameters or KDE object fitted on them.

    Args:
        x_o: Observed data.
        num_simulations: Number of simulations to run.
        eps: Acceptance threshold $\epsilon$ for distance between observed and
            simulated data.
        quantile: Upper quantile of smallest distances for which the corresponding
            parameters are returned, e.g, q=0.01 will return the top 1%. Exactly
            one of quantile or `eps` have to be passed.
        lra: Whether to run linear regression adjustment as in Beaumont et al. 2002
        sass: Whether to determine semi-automatic summary statistics as in
            Fearnhead & Prangle 2012.
        sass_fraction: Fraction of simulation budget used for the initial sass run.
        sass_expansion_degree: Degree of the polynomial feature expansion for the
            sass regression, default 1 - no expansion.
        kde: Whether to run KDE on the accepted parameters to return a KDE
            object from which one can sample.
        kde_kwargs: kwargs for performing KDE:
            'bandwidth='; either a float, or a string naming a bandwidth
            heuristics, e.g., 'cv' (cross validation), 'silvermann' or 'scott',
            default 'cv'.
            'transform': transform applied to the parameters before doing KDE.
            'sample_weights': weights associated with samples. See 'get_kde' for
            more details
        return_summary: Whether to return the distances and data corresponding to
            the accepted parameters.
        num_iid_samples: Number of simulations per parameter. Choose
            `num_iid_samples>1`, if you have chosen a statistical distance that
            evaluates sets of simulations against a set of reference observations
            instead of a single data-point comparison.

    Returns:
        theta (if kde False): accepted parameters
        kde (if kde True): KDE object based on accepted parameters from which one
            can .sample() and .log_prob().
        summary (if summary True): dictionary containing the accepted paramters (if
            kde True), distances and simulated data x.
    """

    # Exactly one of eps or quantile need to be passed.
    assert (eps is not None) ^ (quantile is not None), (
        "Eps or quantile must be passed, but not both."
    )
    if kde_kwargs is None:
        kde_kwargs = {}

    # Run SASS and change the simulator and x_o accordingly.
    if sass:
        num_pilot_simulations = int(sass_fraction * num_simulations)
        self.logger.info(
            "Running SASS with %s pilot samples.", num_pilot_simulations
        )
        num_simulations -= num_pilot_simulations

        pilot_theta = self.prior.sample((num_pilot_simulations,))
        pilot_x = self._batched_simulator(pilot_theta)

        sass_transform = super()._get_sass_transform(
            pilot_theta, pilot_x, sass_expansion_degree
        )

        # Add sass transform to simulator and x_o.
        def simulator(theta):
            return sass_transform(self._batched_simulator(theta))

        x_o = sass_transform(process_x(x_o))
    else:
        simulator = self._batched_simulator

    # Simulate and calculate distances.
    theta = self.prior.sample((num_simulations,))
    theta_repeat = theta.repeat_interleave(num_iid_samples, dim=0)
    x = simulator(theta_repeat)
    x = x.reshape((
        num_simulations,
        num_iid_samples,
        -1,
    ))  # Dim(num_initial_pop, num_iid_samples, -1)

    # Infer x shape to test and set x_o.
    if not self.distance.requires_iid_data:
        x = x.squeeze(1)
        self.x_shape = x[0].shape
    else:
        self.x_shape = x[0, 0].shape

    self.x_o = process_x(x_o, self.x_shape)
    distances = self.distance(self.x_o, x)

    # Select based on acceptance threshold epsilon.
    if eps is not None:
        is_accepted = distances < eps
        num_accepted = is_accepted.sum().item()
        assert num_accepted > 0, f"No parameters accepted, eps={eps} too small"

        theta_accepted = theta[is_accepted]
        distances_accepted = distances[is_accepted]
        x_accepted = x[is_accepted]

    # Select based on quantile on sorted distances.
    elif quantile is not None:
        num_top_samples = int(num_simulations * quantile)
        sort_idx = torch.argsort(distances)
        theta_accepted = theta[sort_idx][:num_top_samples]
        distances_accepted = distances[sort_idx][:num_top_samples]
        x_accepted = x[sort_idx][:num_top_samples]

    else:
        raise ValueError("One of epsilon or quantile has to be passed.")

    # Maybe adjust theta with LRA.
    if lra:
        self.logger.info("Running Linear regression adjustment.")
        final_theta = super()._run_lra(
            theta_accepted, x_accepted, observation=self.x_o
        )
    else:
        final_theta = theta_accepted

    if kde:
        self.logger.info(
            "KDE on %s samples with bandwidth option "
            f"{kde_kwargs.get('bandwidth', 'cv')}. "
            "Beware that KDE can give unreliable results when used with too few"
            " samples and in high dimensions.",
            final_theta.shape[0],
        )

        kde_dist = get_kde(final_theta, **kde_kwargs)

        if return_summary:
            return (
                kde_dist,
                dict(theta=final_theta, distances=distances_accepted, x=x_accepted),
            )
        else:
            return kde_dist
    elif return_summary:
        return final_theta, dict(distances=distances_accepted, x=x_accepted)
    else:
        return final_theta

__init__(simulator, prior, distance='l2', requires_iid_data=None, distance_kwargs=None, num_workers=1, simulation_batch_size=1, distance_batch_size=-1, show_progress_bars=True)

Monte-Carlo Approximate Bayesian Computation (Rejection ABC) [1].

[1] Pritchard, J. K., Seielstad, M. T., Perez-Lezaun, A., & Feldman, M. W. (1999). Population growth of human Y chromosomes: a study of Y chromosome microsatellites. Molecular biology and evolution, 16(12), 1791-1798.

Parameters:

Name	Type	Description	Default
simulator	Callable	A function that takes parameters \(\theta\) and maps them to simulations, or observations, x, \(\mathrm{sim}(\theta)\to x\). Any regular Python callable (i.e. function or class with __call__ method) can be used.	required
prior	Distribution	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. Any object with .log_prob()and .sample() (for example, a PyTorch distribution) can be used.	required
distance	Union[str, Callable]	Distance function to compare observed and simulated data. Can be a custom callable function or one of l1, l2, mse, mmd, wasserstein.	'l2'
requires_iid_data	Optional[bool]	Whether to allow conditioning on iid sampled data or not. Typically, this information is inferred by the choice of the distance, but in case a custom distance is used, this information is pivotal.	None
distance_kwargs	Optional[Dict]	Configurations parameters for the distances. In particular useful for the MMD and Wasserstein distance.	None
num_workers	int	Number of parallel workers to use for simulations.	1
simulation_batch_size	int	Number of parameter sets that the simulator maps to data x at once. If None, we simulate all parameter sets at the same time. If >= 1, the simulator has to process data of shape (simulation_batch_size, parameter_dimension).	1
distance_batch_size	int	Number of simulations that the distance function evaluates against the reference observations at once. If -1, we evaluate all simulations at the same time.	-1

Source code in sbi/inference/abc/mcabc.py

def __init__(
    self,
    simulator: Callable,
    prior: Distribution,
    distance: Union[str, Callable] = "l2",
    requires_iid_data: Optional[bool] = None,
    distance_kwargs: Optional[Dict] = None,
    num_workers: int = 1,
    simulation_batch_size: int = 1,
    distance_batch_size: int = -1,
    show_progress_bars: bool = True,
):
    r"""Monte-Carlo Approximate Bayesian Computation (Rejection ABC) [1].

    [1] Pritchard, J. K., Seielstad, M. T., Perez-Lezaun, A., & Feldman, M. W.
    (1999). Population growth of human Y chromosomes: a study of Y chromosome
    microsatellites. Molecular biology and evolution, 16(12), 1791-1798.

    Args:
        simulator: A function that takes parameters $\theta$ and maps them to
            simulations, or observations, `x`, $\mathrm{sim}(\theta)\to x$. Any
            regular Python callable (i.e. function or class with `__call__` method)
            can be used.
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. Any
            object with `.log_prob()`and `.sample()` (for example, a PyTorch
            distribution) can be used.
        distance: Distance function to compare observed and simulated data. Can be
            a custom callable function or one of `l1`, `l2`, `mse`,
            `mmd`, `wasserstein`.
        requires_iid_data: Whether to allow conditioning on iid sampled data or not.
            Typically, this information is inferred by the choice of the distance,
            but in case a custom distance is used, this information is pivotal.
        distance_kwargs: Configurations parameters for the distances. In particular
            useful for the MMD and Wasserstein distance.
        num_workers: Number of parallel workers to use for simulations.
        simulation_batch_size: Number of parameter sets that the simulator
            maps to data x at once. If None, we simulate all parameter sets at the
            same time. If >= 1, the simulator has to process data of shape
            (simulation_batch_size, parameter_dimension).
        distance_batch_size: Number of simulations that the distance function
            evaluates against the reference observations at once. If -1, we evaluate
            all simulations at the same time.
    """

    super().__init__(
        simulator=simulator,
        prior=prior,
        distance=distance,
        requires_iid_data=requires_iid_data,
        distance_kwargs=distance_kwargs,
        num_workers=num_workers,
        simulation_batch_size=simulation_batch_size,
        distance_batch_size=distance_batch_size,
        show_progress_bars=show_progress_bars,
    )

SMCABC

Bases: ABCBASE

Sequential Monte Carlo Approximate Bayesian Computation.

Source code in sbi/inference/abc/smcabc.py

class SMCABC(ABCBASE):
    """Sequential Monte Carlo Approximate Bayesian Computation."""

    def __init__(
        self,
        simulator: Callable,
        prior: Distribution,
        distance: Union[str, Callable] = "l2",
        requires_iid_data: Optional[bool] = None,
        distance_kwargs: Optional[Dict] = None,
        num_workers: int = 1,
        simulation_batch_size: int = 1,
        distance_batch_size: int = -1,
        show_progress_bars: bool = True,
        kernel: Optional[str] = "gaussian",
        algorithm_variant: str = "C",
    ):
        r"""Sequential Monte Carlo Approximate Bayesian Computation.

        We distinguish between three different SMC methods here:
            - A: Toni et al. 2010 (Phd Thesis)
            - B: Sisson et al. 2007 (with correction from 2009)
            - C: Beaumont et al. 2009

        In Toni et al. 2010 we find an overview of the differences on page 34:
            - B: same as A except for resampling of weights if the effective sampling
                size is too small.
            - C: same as A except for calculation of the covariance of the perturbation
                kernel: the kernel covariance is a scaled version of the covariance of
                the previous population.

        Args:
            simulator: A function that takes parameters $\theta$ and maps them to
                simulations, or observations, `x`, $\mathrm{sim}(\theta)\to x$. Any
                regular Python callable (i.e. function or class with `__call__` method)
                can be used.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Any
                object with `.log_prob()`and `.sample()` (for example, a PyTorch
                distribution) can be used.
            distance: Distance function to compare observed and simulated data. Can be
                a custom callable function or one of `l1`, `l2`, `mse`,
                `mmd`, `wasserstein`.
            requires_iid_data: Whether to allow conditioning on iid sampled data or not.
                Typically, this information is inferred by the choice of the distance,
                but in case a custom distance is used, this information is pivotal.
            distance_kwargs: Configurations parameters for the distances. In particular
                useful for the MMD and Wasserstein distance.
            num_workers: Number of parallel workers to use for simulations.
            simulation_batch_size: Number of parameter sets that the simulator
                maps to data x at once. If None, we simulate all parameter sets at the
                same time. If >= 1, the simulator has to process data of shape
                (simulation_batch_size, parameter_dimension).
            distance_batch_size: Number of simulations that the distance function
                evaluates against the reference observations at once. If -1, we evaluate
                all simulations at the same time.
            show_progress_bars: Whether to show a progressbar during simulation and
                sampling.
            kernel: Perturbation kernel.
            algorithm_variant: Indicating the choice of algorithm variant, A, B, or C.
        """

        super().__init__(
            simulator=simulator,
            prior=prior,
            distance=distance,
            requires_iid_data=requires_iid_data,
            distance_kwargs=distance_kwargs,
            num_workers=num_workers,
            simulation_batch_size=simulation_batch_size,
            distance_batch_size=distance_batch_size,
            show_progress_bars=show_progress_bars,
        )

        kernels = ("gaussian", "uniform")
        assert kernel in kernels, (
            f"Kernel '{kernel}' not supported. Choose one from {kernels}."
        )
        self.kernel = kernel

        algorithm_variants = ("A", "B", "C")
        assert algorithm_variant in algorithm_variants, (
            f"SMCABC variant '{algorithm_variant}' not supported, choose one from"
            " {algorithm_variants}."
        )
        self.algorithm_variant = algorithm_variant
        self.distance_to_x0 = None
        self.simulation_counter = 0
        self.num_simulations = 0
        self.kernel_variance = None

        # Define simulator that keeps track of budget.
        def simulate_with_budget(theta):
            self.simulation_counter += theta.shape[0]
            return self._batched_simulator(theta)

        self._simulate_with_budget = simulate_with_budget

    def __call__(
        self,
        x_o: Union[Tensor, ndarray],
        num_particles: int,
        num_initial_pop: int,
        num_simulations: int,
        epsilon_decay: float,
        distance_based_decay: bool = False,
        ess_min: Optional[float] = None,
        kernel_variance_scale: float = 1.0,
        use_last_pop_samples: bool = True,
        return_summary: bool = False,
        kde: bool = False,
        kde_kwargs: Optional[Dict[str, Any]] = None,
        kde_sample_weights: bool = False,
        lra: bool = False,
        lra_with_weights: bool = False,
        sass: bool = False,
        sass_fraction: float = 0.25,
        sass_expansion_degree: int = 1,
        num_iid_samples: int = 1,
    ) -> Union[Tensor, KDEWrapper, Tuple[Tensor, dict], Tuple[KDEWrapper, dict]]:
        r"""Run SMCABC and return accepted parameters or KDE object fitted on them.

        Args:
            x_o: Observed data.
            num_particles: Number of particles in each population.
            num_initial_pop: Number of simulations used for initial population.
            num_simulations: Total number of possible simulations.
            epsilon_decay: Factor with which the acceptance threshold $\epsilon$ decays.
            distance_based_decay: Whether the $\epsilon$ decay is constant over
                populations or calculated from the previous populations distribution of
                distances.
            ess_min: Threshold of effective sampling size for resampling weights. Not
                used when None (default).
            kernel_variance_scale: Factor for scaling the perturbation kernel variance.
            use_last_pop_samples: Whether to fill up the current population with
                samples from the previous population when the budget is used up. If
                False, the current population is discarded and the previous population
                is returned.
            lra: Whether to run linear regression adjustment as in Beaumont et al. 2002
            lra_with_weights: Whether to run lra as weighted linear regression with SMC
                weights
            sass: Whether to determine semi-automatic summary statistics (sass) as in
                Fearnhead & Prangle 2012.
            sass_fraction: Fraction of simulation budget used for the initial sass run.
            sass_expansion_degree: Degree of the polynomial feature expansion for the
                sass regression, default 1 - no expansion.
            kde: Whether to run KDE on the accepted parameters to return a KDE
                object from which one can sample.
            kde_kwargs: kwargs for performing KDE:
                'bandwidth='; either a float, or a string naming a bandwidth
                heuristics, e.g., 'cv' (cross validation), 'silvermann' or 'scott',
                default 'cv'.
                'transform': transform applied to the parameters before doing KDE.
                'sample_weights': weights associated with samples. See 'get_kde' for
                more details
            kde_sample_weights: Whether perform weighted KDE with SMC weights or on raw
                particles.
            return_summary: Whether to return a dictionary with all accepted particles,
                weights, etc. at the end.
            num_iid_samples: Number of simulations per parameter. Choose
                `num_iid_samples>1`, if you have chosen a statistical distance that
                evaluates sets of simulations against a set of reference observations
                instead of a single data-point comparison.

        Returns:
            theta (if kde False): accepted parameters of the last population.
            kde (if kde True): KDE object fitted on accepted parameters, from which one
                can .sample() and .log_prob().
            summary (if return_summary True): dictionary containing the accepted
                paramters (if kde True), distances and simulated data x of all
                populations.
        """

        pop_idx = 0
        self.num_simulations = num_simulations * num_iid_samples
        if kde_kwargs is None:
            kde_kwargs = {}
        assert isinstance(epsilon_decay, float) and epsilon_decay > 0.0
        assert not (self.distance.requires_iid_data and lra), (
            "Currently there is no support to run inference "
        )
        "on multiple observations together with lra."
        assert not (self.distance.requires_iid_data and sass), (
            "Currently there is no support to run inference "
        )
        "on multiple observations together with sass."

        # Pilot run for SASS.
        if sass:
            num_pilot_simulations = int(sass_fraction * num_simulations)
            self.logger.info(
                "Running SASS with %s pilot samples.", num_pilot_simulations
            )
            sass_transform = self._run_sass_set_xo(
                num_particles,
                num_pilot_simulations,
                x_o,
                num_iid_samples,
                lra,
                sass_expansion_degree,
            )
            # Udpate simulator and xo
            x_o = sass_transform(self.x_o)

            def sass_simulator(theta):
                self.simulation_counter += theta.shape[0]
                return sass_transform(self._batched_simulator(theta))

            self._simulate_with_budget = sass_simulator

        # run initial population
        particles, epsilon, distances, x = self._set_xo_and_sample_initial_population(
            x_o, num_particles, num_initial_pop, num_iid_samples
        )
        log_weights = torch.log(1 / num_particles * torch.ones(num_particles))

        self.logger.info((
            "population=%s, eps=%s, ess=%s, num_sims=%s",
            pop_idx,
            epsilon,
            1.0,
            num_initial_pop,
        ))

        all_particles = [particles]
        all_log_weights = [log_weights]
        all_distances = [distances]
        all_epsilons = [epsilon]
        all_x = [x]

        while self.simulation_counter < self.num_simulations:
            pop_idx += 1
            # Decay based on quantile of distances from previous pop.
            if distance_based_decay:
                epsilon = self._get_next_epsilon(
                    all_distances[pop_idx - 1], epsilon_decay
                )
            # Constant decay.
            else:
                epsilon *= epsilon_decay

            # Get kernel variance from previous pop.
            self.kernel_variance = self._get_kernel_variance(
                all_particles[pop_idx - 1],
                torch.exp(all_log_weights[pop_idx - 1]),
                samples_per_dim=500,
                kernel_variance_scale=kernel_variance_scale,
            )
            particles, log_weights, distances, x = self._sample_next_population(
                particles=all_particles[pop_idx - 1],
                log_weights=all_log_weights[pop_idx - 1],
                distances=all_distances[pop_idx - 1],
                epsilon=epsilon,
                x=all_x[pop_idx - 1],
                num_iid_samples=num_iid_samples,
                use_last_pop_samples=use_last_pop_samples,
            )

            # Resample population if effective sampling size is too small.
            if ess_min is not None:
                particles, log_weights = self._resample_if_ess_too_small(
                    particles, log_weights, ess_min, pop_idx
                )

            self.logger.info((
                "population=%s done: eps={epsilon:.6f}, num_sims=%s.",
                pop_idx,
                epsilon,
                self.simulation_counter,
            ))

            # collect results
            all_particles.append(particles)
            all_log_weights.append(log_weights)
            all_distances.append(distances)
            all_epsilons.append(epsilon)
            all_x.append(x)

        # Maybe run LRA and adjust weights.
        if lra:
            self.logger.info("Running Linear regression adjustment.")
            adjusted_particles, _ = self._run_lra_update_weights(
                particles=all_particles[-1],
                xs=all_x[-1],
                observation=process_x(x_o),
                log_weights=all_log_weights[-1],
                lra_with_weights=lra_with_weights,
            )
            final_particles = adjusted_particles
        else:
            final_particles = all_particles[-1]

        if kde:
            self.logger.info(
                """KDE on %s samples with bandwidth option %s. Beware that KDE can give
                unreliable results when used with too few samples and in high
                dimensions.""",
                final_particles.shape[0],
                kde_kwargs.get("bandwidth", "cv"),
            )
            # Maybe get particles weights from last population for weighted KDE.
            if kde_sample_weights:
                kde_kwargs["sample_weights"] = all_log_weights[-1].exp()

            kde_dist = get_kde(final_particles, **kde_kwargs)

            if return_summary:
                return (
                    kde_dist,
                    dict(
                        particles=all_particles,
                        weights=all_log_weights,
                        epsilons=all_epsilons,
                        distances=all_distances,
                        xs=all_x,
                    ),
                )
            else:
                return kde_dist

        if return_summary:
            return (
                final_particles,
                dict(
                    particles=all_particles,
                    weights=all_log_weights,
                    epsilons=all_epsilons,
                    distances=all_distances,
                    xs=all_x,
                ),
            )
        else:
            return final_particles

    def _set_xo_and_sample_initial_population(
        self,
        x_o: Array,
        num_particles: int,
        num_initial_pop: int,
        num_iid_samples: int,
    ) -> Tuple[Tensor, float, Tensor, Tensor]:
        """Return particles, epsilon and distances of initial population."""

        assert num_particles <= num_initial_pop, (
            "number of initial round simulations must be greater than population size"
        )

        assert (x_o.shape[0] == 1) or self.distance.requires_iid_data, (
            "Your data contain iid data-points, but the choice of "
            "your distance does not allow multiple conditioning "
            "observations."
        )

        theta = self.prior.sample((num_initial_pop,))

        theta_repeat = theta.repeat_interleave(num_iid_samples, dim=0)
        x = self._simulate_with_budget(theta_repeat)
        x = x.reshape((
            num_initial_pop,
            num_iid_samples,
            -1,
        ))  # Dim(num_initial_pop, num_iid_samples, -1)

        # Infer x shape to test and set x_o.
        if not self.distance.requires_iid_data:
            x = x.squeeze(1)
            self.x_shape = x[0].shape
        else:
            self.x_shape = x[0, 0].shape
        self.x_o = process_x(x_o, self.x_shape)

        distances = self.distance(self.x_o, x)
        sortidx = torch.argsort(distances)
        particles = theta[sortidx][:num_particles]
        # Take last accepted distance as epsilon.
        initial_epsilon = distances[sortidx][num_particles - 1].item()

        if not math.isfinite(initial_epsilon):
            initial_epsilon = 1e8

        return (
            particles,
            initial_epsilon,
            distances[sortidx][:num_particles],
            x[sortidx][:num_particles],
        )

    def _sample_next_population(
        self,
        particles: Tensor,
        log_weights: Tensor,
        distances: Tensor,
        epsilon: float,
        x: Tensor,
        num_iid_samples: int,
        use_last_pop_samples: bool = True,
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """Return particles, weights and distances of new population."""

        new_particles = []
        new_log_weights = []
        new_distances = []
        new_x = []

        num_accepted_particles = 0
        num_particles = particles.shape[0]

        while num_accepted_particles < num_particles:
            # Upperbound for batch size to not exceed simulation budget.
            num_batch = min(
                num_particles - num_accepted_particles,
                self.num_simulations - self.simulation_counter,
            )

            # Sample from previous population and perturb.
            particle_candidates = self._sample_and_perturb(
                particles, torch.exp(log_weights), num_samples=num_batch
            )
            # Simulate and select based on distance.
            candidates_repeated = particle_candidates.repeat_interleave(
                num_iid_samples, dim=0
            )
            x_candidates = self._simulate_with_budget(candidates_repeated)
            x_candidates = x_candidates.reshape((
                num_batch,
                num_iid_samples,
                -1,
            ))  # Dim(num_initial_pop, num_iid_samples, -1)
            if not self.distance.requires_iid_data:
                x_candidates = x_candidates.squeeze(1)

            dists = self.distance(self.x_o, x_candidates)
            is_accepted = dists <= epsilon
            num_accepted_batch = int(is_accepted.sum().item())

            if num_accepted_batch > 0:
                new_particles.append(particle_candidates[is_accepted])
                new_log_weights.append(
                    self._calculate_new_log_weights(
                        particle_candidates[is_accepted],
                        particles,
                        log_weights,
                    )
                )
                new_distances.append(dists[is_accepted])
                new_x.append(x_candidates[is_accepted])
                num_accepted_particles += num_accepted_batch

            # If simulation budget was exceeded and we still need particles, take
            # previous population or fill up with previous population.
            if (
                self.simulation_counter >= self.num_simulations
                and num_accepted_particles < num_particles
            ):
                if use_last_pop_samples:
                    num_remaining = num_particles - num_accepted_particles
                    self.logger.info(
                        """Simulation Budget exceeded, filling up with %s
                        samples from last population.""",
                        num_remaining,
                    )
                    # Some new particles have been accepted already, therefore
                    # fill up the remaining once with old particles and weights.
                    new_particles.append(particles[:num_remaining, :])
                    # Recalculate weights with new particles.
                    new_log_weights = [
                        self._calculate_new_log_weights(
                            torch.cat(new_particles),
                            particles,
                            log_weights,
                        )
                    ]
                    new_distances.append(distances[:num_remaining])
                    new_x.append(x[:num_remaining])
                else:
                    self.logger.info(
                        "Simulation Budget exceeded, returning previous population."
                    )
                    new_particles = [particles]
                    new_log_weights = [log_weights]
                    new_distances = [distances]
                    new_x = [x]

                break

        # collect lists of tensors into tensors
        new_particles = torch.cat(new_particles)
        new_log_weights = torch.cat(new_log_weights)
        new_distances = torch.cat(new_distances)
        new_x = torch.cat(new_x)

        # normalize the new weights
        new_log_weights -= torch.logsumexp(new_log_weights, dim=0)

        # Return sorted wrt distances.
        sort_idx = torch.argsort(new_distances)

        return (
            new_particles[sort_idx],
            new_log_weights[sort_idx],
            new_distances[sort_idx],
            new_x[sort_idx],
        )

    def _get_next_epsilon(self, distances: Tensor, quantile: float) -> float:
        """Return epsilon for next round based on quantile of this round's distances.

        Note: distances are made unique to avoid repeated distances from simulations
        that result in the same observation.

        Args:
            distances: The distances accepted in this round.
            quantile: Quantile in the distance distribution to determine new epsilon.

        Returns:
            epsilon: Epsilon for the next population.
        """
        # Take unique distances to skip same distances simulations (return is sorted).
        distances = torch.unique(distances)
        # Cumsum as cdf proxy.
        distances_cdf = torch.cumsum(distances, dim=0) / distances.sum()
        # Take the q quantile of distances.
        try:
            qidx = torch.where(distances_cdf >= quantile)[0][0]
        except IndexError:
            self.logger.warning((
                """Accepted unique distances=%s don't match quantile=%s. Selecting
                    last distance.""",
                distances,
                quantile,
            ))
            qidx = -1

        # The new epsilon is given by that distance.
        return distances[qidx].item()

    def _calculate_new_log_weights(
        self,
        new_particles: Tensor,
        old_particles: Tensor,
        old_log_weights: Tensor,
    ) -> Tensor:
        """Return new log weights following formulas in publications A,B anc C."""

        # Prior can be batched across new particles.
        prior_log_probs = self.prior.log_prob(new_particles)

        # Contstruct function to get kernel log prob for given old particle.
        # The kernel is centered on each old particle as in all three variants (A,B,C).
        def kernel_log_prob(new_particle):
            return self._get_new_kernel(old_particles).log_prob(new_particle)

        # We still have to loop over particles here because
        # the kernel log probs are already batched across old particles.
        log_weighted_sum = torch.tensor(
            [
                torch.logsumexp(old_log_weights + kernel_log_prob(new_particle), dim=0)
                for new_particle in new_particles
            ],
            dtype=torch.float32,
        )
        # new weights are prior probs over weighted sum:
        return prior_log_probs - log_weighted_sum

    @staticmethod
    def sample_from_population_with_weights(
        particles: Tensor, weights: Tensor, num_samples: int = 1
    ) -> Tensor:
        """Return samples from particles sampled with weights."""

        # define multinomial with weights as probs
        multi = Multinomial(probs=weights)
        # sample num samples, with replacement
        samples = multi.sample(sample_shape=torch.Size((num_samples,)))
        # get indices of success trials
        indices = torch.where(samples)[1]
        # return those indices from trace
        return particles[indices]

    def _sample_and_perturb(
        self, particles: Tensor, weights: Tensor, num_samples: int = 1
    ) -> Tensor:
        """Sample and perturb batch of new parameters from trace.

        Reject sampled and perturbed parameters outside of prior.
        """

        num_accepted = 0
        parameters = []
        while num_accepted < num_samples:
            parms = self.sample_from_population_with_weights(
                particles, weights, num_samples=num_samples - num_accepted
            )

            # Create kernel on params and perturb.
            parms_perturbed = self._get_new_kernel(parms).sample()

            is_within_prior = within_support(self.prior, parms_perturbed)
            num_accepted += int(is_within_prior.sum().item())

            if num_accepted > 0:
                parameters.append(parms_perturbed[is_within_prior])

        return torch.cat(parameters)

    def _get_kernel_variance(
        self,
        particles: Tensor,
        weights: Tensor,
        samples_per_dim: int = 100,
        kernel_variance_scale: float = 1.0,
    ) -> Tensor:
        """Return kernel variance for a given population of particles and weights."""
        if self.kernel == "gaussian":
            # For variant C, Beaumont et al. 2009, the kernel variance comes from the
            # previous population.
            if self.algorithm_variant == "C":
                # Calculate weighted covariance of particles.
                population_cov = torch.tensor(
                    np.atleast_2d(np.cov(particles, rowvar=False, aweights=weights)),
                    dtype=torch.float32,
                )
                # Make sure variance is nonsingular.
                try:
                    torch.linalg.cholesky(kernel_variance_scale * population_cov)
                except RuntimeError:
                    self.logger.warning(
                        """"Singular particle covariance, using unit covariance."""
                    )
                    population_cov = torch.eye(particles.shape[1])
                return kernel_variance_scale * population_cov
            # While for Toni et al. and Sisson et al. it comes from the parameter
            # ranges.
            elif self.algorithm_variant in ("A", "B"):
                particle_ranges = self._get_particle_ranges(
                    particles, weights, samples_per_dim=samples_per_dim
                )
                return kernel_variance_scale * torch.diag(particle_ranges)
            else:
                raise ValueError(f"Variant, '{self.algorithm_variant}' not supported.")
        elif self.kernel == "uniform":
            # Variance spans the range of parameters for every dimension.
            return kernel_variance_scale * self._get_particle_ranges(
                particles, weights, samples_per_dim=samples_per_dim
            )
        else:
            raise ValueError(f"Kernel, '{self.kernel}' not supported.")

    def _get_new_kernel(self, thetas: Tensor) -> Distribution:
        """Return new kernel distribution for a given set of paramters."""

        if self.kernel == "gaussian":
            assert self.kernel_variance is not None, "get kernel variance first."
            assert self.kernel_variance.ndim == 2
            return MultivariateNormal(
                loc=thetas, covariance_matrix=self.kernel_variance
            )

        elif self.kernel == "uniform":
            low = thetas - self.kernel_variance
            high = thetas + self.kernel_variance
            # Move batch shape to event shape to get Uniform that is multivariate in
            # parameter dimension.
            return BoxUniform(low=low, high=high)
        else:
            raise ValueError(f"Kernel, '{self.kernel}' not supported.")

    def _resample_if_ess_too_small(
        self,
        particles: Tensor,
        log_weights: Tensor,
        ess_min: float,
        pop_idx: int,
    ) -> Tuple[Tensor, Tensor]:
        """Return resampled particles and uniform weights if effectice sampling size is
        too small.
        """

        num_particles = particles.shape[0]
        # Calculate relative ESS (ESS/N) which ranges from 0 to 1
        # This is scale-invariant and commonly used in SMC literature
        weights = torch.exp(log_weights)
        ess = ((torch.sum(weights) ** 2) / torch.sum(weights**2)) / num_particles
        # Resampling of weights for low ESS only for Sisson et al. 2007.
        if ess < ess_min:
            self.logger.info("ESS=%s too low, resampling pop %s...", ess, pop_idx)
            # First resample, then set to uniform weights as in Sisson et al. 2007.
            particles = self.sample_from_population_with_weights(
                particles, torch.exp(log_weights), num_samples=num_particles
            )
            log_weights = torch.log(1 / num_particles * torch.ones(num_particles))

        return particles, log_weights

    def _run_lra_update_weights(
        self,
        particles: Tensor,
        xs: Tensor,
        observation: Tensor,
        log_weights: Tensor,
        lra_with_weights: bool,
    ) -> Tuple[Tensor, Tensor]:
        """Return particles and weights adjusted with LRA.

        Runs (weighted) linear regression from xs onto particles to adjust the
        particles.

        Updates the SMC weights according to the new particles.
        """

        adjusted_particels = super()._run_lra(
            theta=particles,
            x=xs,
            observation=observation,
            sample_weight=log_weights.exp() if lra_with_weights else None,
        )

        # Update SMC weights with LRA adjusted weights
        adjusted_log_weights = self._calculate_new_log_weights(
            new_particles=adjusted_particels,
            old_particles=particles,
            old_log_weights=log_weights,
        )

        return adjusted_particels, adjusted_log_weights

    def _run_sass_set_xo(
        self,
        num_particles: int,
        num_pilot_simulations: int,
        x_o: Union[Tensor, ndarray],
        num_iid_samples: int,
        lra: bool = False,
        sass_expansion_degree: int = 1,
    ) -> Callable:
        """Return transform for semi-automatic summary statistics.

        Runs an single round of rejection abc with fixed budget and accepts
        num_particles simulations to run the regression for sass.

        Sets self.x_o once the x_shape can be derived from simulations.
        """
        (
            pilot_particles,
            _,
            _,
            pilot_xs,
        ) = self._set_xo_and_sample_initial_population(
            x_o, num_particles, num_pilot_simulations, num_iid_samples
        )
        assert self.x_o is not None, "x_o not set yet."

        # Adjust with LRA.
        if lra:
            pilot_particles = super()._run_lra(pilot_particles, pilot_xs, self.x_o)
        sass_transform = super()._get_sass_transform(
            pilot_particles,
            pilot_xs,
            expansion_degree=sass_expansion_degree,
            sample_weight=None,
        )
        return sass_transform

    def _get_particle_ranges(
        self, particles: Tensor, weights: Tensor, samples_per_dim: int = 100
    ) -> Tensor:
        """Return range of particles in each parameter dimension."""

        # get weighted samples
        samples = self.sample_from_population_with_weights(
            particles,
            weights,
            num_samples=samples_per_dim * particles.shape[1],
        )

        # Variance spans the range of particles for every dimension.
        particle_ranges = samples.max(0).values - samples.min(0).values
        assert particle_ranges.ndim < 2
        return particle_ranges

__call__(x_o, num_particles, num_initial_pop, num_simulations, epsilon_decay, distance_based_decay=False, ess_min=None, kernel_variance_scale=1.0, use_last_pop_samples=True, return_summary=False, kde=False, kde_kwargs=None, kde_sample_weights=False, lra=False, lra_with_weights=False, sass=False, sass_fraction=0.25, sass_expansion_degree=1, num_iid_samples=1)

Run SMCABC and return accepted parameters or KDE object fitted on them.

Parameters:

Name	Type	Description	Default
x_o	Union[Tensor, ndarray]	Observed data.	required
num_particles	int	Number of particles in each population.	required
num_initial_pop	int	Number of simulations used for initial population.	required
num_simulations	int	Total number of possible simulations.	required
epsilon_decay	float	Factor with which the acceptance threshold \(\epsilon\) decays.	required
distance_based_decay	bool	Whether the \(\epsilon\) decay is constant over populations or calculated from the previous populations distribution of distances.	False
ess_min	Optional[float]	Threshold of effective sampling size for resampling weights. Not used when None (default).	None
kernel_variance_scale	float	Factor for scaling the perturbation kernel variance.	1.0
use_last_pop_samples	bool	Whether to fill up the current population with samples from the previous population when the budget is used up. If False, the current population is discarded and the previous population is returned.	True
lra	bool	Whether to run linear regression adjustment as in Beaumont et al. 2002	False
lra_with_weights	bool	Whether to run lra as weighted linear regression with SMC weights	False
sass	bool	Whether to determine semi-automatic summary statistics (sass) as in Fearnhead & Prangle 2012.	False
sass_fraction	float	Fraction of simulation budget used for the initial sass run.	0.25
sass_expansion_degree	int	Degree of the polynomial feature expansion for the sass regression, default 1 - no expansion.	1
kde	bool	Whether to run KDE on the accepted parameters to return a KDE object from which one can sample.	False
kde_kwargs	Optional[Dict[str, Any]]	kwargs for performing KDE: ‘bandwidth=’; either a float, or a string naming a bandwidth heuristics, e.g., ‘cv’ (cross validation), ‘silvermann’ or ‘scott’, default ‘cv’. ‘transform’: transform applied to the parameters before doing KDE. ‘sample_weights’: weights associated with samples. See ‘get_kde’ for more details	None
kde_sample_weights	bool	Whether perform weighted KDE with SMC weights or on raw particles.	False
return_summary	bool	Whether to return a dictionary with all accepted particles, weights, etc. at the end.	False
num_iid_samples	int	Number of simulations per parameter. Choose num_iid_samples>1, if you have chosen a statistical distance that evaluates sets of simulations against a set of reference observations instead of a single data-point comparison.	1

Returns:

Name	Type	Description
theta	if kde False	accepted parameters of the last population.
kde	if kde True	KDE object fitted on accepted parameters, from which one can .sample() and .log_prob().
summary	if return_summary True	dictionary containing the accepted paramters (if kde True), distances and simulated data x of all populations.

Source code in sbi/inference/abc/smcabc.py

def __call__(
    self,
    x_o: Union[Tensor, ndarray],
    num_particles: int,
    num_initial_pop: int,
    num_simulations: int,
    epsilon_decay: float,
    distance_based_decay: bool = False,
    ess_min: Optional[float] = None,
    kernel_variance_scale: float = 1.0,
    use_last_pop_samples: bool = True,
    return_summary: bool = False,
    kde: bool = False,
    kde_kwargs: Optional[Dict[str, Any]] = None,
    kde_sample_weights: bool = False,
    lra: bool = False,
    lra_with_weights: bool = False,
    sass: bool = False,
    sass_fraction: float = 0.25,
    sass_expansion_degree: int = 1,
    num_iid_samples: int = 1,
) -> Union[Tensor, KDEWrapper, Tuple[Tensor, dict], Tuple[KDEWrapper, dict]]:
    r"""Run SMCABC and return accepted parameters or KDE object fitted on them.

    Args:
        x_o: Observed data.
        num_particles: Number of particles in each population.
        num_initial_pop: Number of simulations used for initial population.
        num_simulations: Total number of possible simulations.
        epsilon_decay: Factor with which the acceptance threshold $\epsilon$ decays.
        distance_based_decay: Whether the $\epsilon$ decay is constant over
            populations or calculated from the previous populations distribution of
            distances.
        ess_min: Threshold of effective sampling size for resampling weights. Not
            used when None (default).
        kernel_variance_scale: Factor for scaling the perturbation kernel variance.
        use_last_pop_samples: Whether to fill up the current population with
            samples from the previous population when the budget is used up. If
            False, the current population is discarded and the previous population
            is returned.
        lra: Whether to run linear regression adjustment as in Beaumont et al. 2002
        lra_with_weights: Whether to run lra as weighted linear regression with SMC
            weights
        sass: Whether to determine semi-automatic summary statistics (sass) as in
            Fearnhead & Prangle 2012.
        sass_fraction: Fraction of simulation budget used for the initial sass run.
        sass_expansion_degree: Degree of the polynomial feature expansion for the
            sass regression, default 1 - no expansion.
        kde: Whether to run KDE on the accepted parameters to return a KDE
            object from which one can sample.
        kde_kwargs: kwargs for performing KDE:
            'bandwidth='; either a float, or a string naming a bandwidth
            heuristics, e.g., 'cv' (cross validation), 'silvermann' or 'scott',
            default 'cv'.
            'transform': transform applied to the parameters before doing KDE.
            'sample_weights': weights associated with samples. See 'get_kde' for
            more details
        kde_sample_weights: Whether perform weighted KDE with SMC weights or on raw
            particles.
        return_summary: Whether to return a dictionary with all accepted particles,
            weights, etc. at the end.
        num_iid_samples: Number of simulations per parameter. Choose
            `num_iid_samples>1`, if you have chosen a statistical distance that
            evaluates sets of simulations against a set of reference observations
            instead of a single data-point comparison.

    Returns:
        theta (if kde False): accepted parameters of the last population.
        kde (if kde True): KDE object fitted on accepted parameters, from which one
            can .sample() and .log_prob().
        summary (if return_summary True): dictionary containing the accepted
            paramters (if kde True), distances and simulated data x of all
            populations.
    """

    pop_idx = 0
    self.num_simulations = num_simulations * num_iid_samples
    if kde_kwargs is None:
        kde_kwargs = {}
    assert isinstance(epsilon_decay, float) and epsilon_decay > 0.0
    assert not (self.distance.requires_iid_data and lra), (
        "Currently there is no support to run inference "
    )
    "on multiple observations together with lra."
    assert not (self.distance.requires_iid_data and sass), (
        "Currently there is no support to run inference "
    )
    "on multiple observations together with sass."

    # Pilot run for SASS.
    if sass:
        num_pilot_simulations = int(sass_fraction * num_simulations)
        self.logger.info(
            "Running SASS with %s pilot samples.", num_pilot_simulations
        )
        sass_transform = self._run_sass_set_xo(
            num_particles,
            num_pilot_simulations,
            x_o,
            num_iid_samples,
            lra,
            sass_expansion_degree,
        )
        # Udpate simulator and xo
        x_o = sass_transform(self.x_o)

        def sass_simulator(theta):
            self.simulation_counter += theta.shape[0]
            return sass_transform(self._batched_simulator(theta))

        self._simulate_with_budget = sass_simulator

    # run initial population
    particles, epsilon, distances, x = self._set_xo_and_sample_initial_population(
        x_o, num_particles, num_initial_pop, num_iid_samples
    )
    log_weights = torch.log(1 / num_particles * torch.ones(num_particles))

    self.logger.info((
        "population=%s, eps=%s, ess=%s, num_sims=%s",
        pop_idx,
        epsilon,
        1.0,
        num_initial_pop,
    ))

    all_particles = [particles]
    all_log_weights = [log_weights]
    all_distances = [distances]
    all_epsilons = [epsilon]
    all_x = [x]

    while self.simulation_counter < self.num_simulations:
        pop_idx += 1
        # Decay based on quantile of distances from previous pop.
        if distance_based_decay:
            epsilon = self._get_next_epsilon(
                all_distances[pop_idx - 1], epsilon_decay
            )
        # Constant decay.
        else:
            epsilon *= epsilon_decay

        # Get kernel variance from previous pop.
        self.kernel_variance = self._get_kernel_variance(
            all_particles[pop_idx - 1],
            torch.exp(all_log_weights[pop_idx - 1]),
            samples_per_dim=500,
            kernel_variance_scale=kernel_variance_scale,
        )
        particles, log_weights, distances, x = self._sample_next_population(
            particles=all_particles[pop_idx - 1],
            log_weights=all_log_weights[pop_idx - 1],
            distances=all_distances[pop_idx - 1],
            epsilon=epsilon,
            x=all_x[pop_idx - 1],
            num_iid_samples=num_iid_samples,
            use_last_pop_samples=use_last_pop_samples,
        )

        # Resample population if effective sampling size is too small.
        if ess_min is not None:
            particles, log_weights = self._resample_if_ess_too_small(
                particles, log_weights, ess_min, pop_idx
            )

        self.logger.info((
            "population=%s done: eps={epsilon:.6f}, num_sims=%s.",
            pop_idx,
            epsilon,
            self.simulation_counter,
        ))

        # collect results
        all_particles.append(particles)
        all_log_weights.append(log_weights)
        all_distances.append(distances)
        all_epsilons.append(epsilon)
        all_x.append(x)

    # Maybe run LRA and adjust weights.
    if lra:
        self.logger.info("Running Linear regression adjustment.")
        adjusted_particles, _ = self._run_lra_update_weights(
            particles=all_particles[-1],
            xs=all_x[-1],
            observation=process_x(x_o),
            log_weights=all_log_weights[-1],
            lra_with_weights=lra_with_weights,
        )
        final_particles = adjusted_particles
    else:
        final_particles = all_particles[-1]

    if kde:
        self.logger.info(
            """KDE on %s samples with bandwidth option %s. Beware that KDE can give
            unreliable results when used with too few samples and in high
            dimensions.""",
            final_particles.shape[0],
            kde_kwargs.get("bandwidth", "cv"),
        )
        # Maybe get particles weights from last population for weighted KDE.
        if kde_sample_weights:
            kde_kwargs["sample_weights"] = all_log_weights[-1].exp()

        kde_dist = get_kde(final_particles, **kde_kwargs)

        if return_summary:
            return (
                kde_dist,
                dict(
                    particles=all_particles,
                    weights=all_log_weights,
                    epsilons=all_epsilons,
                    distances=all_distances,
                    xs=all_x,
                ),
            )
        else:
            return kde_dist

    if return_summary:
        return (
            final_particles,
            dict(
                particles=all_particles,
                weights=all_log_weights,
                epsilons=all_epsilons,
                distances=all_distances,
                xs=all_x,
            ),
        )
    else:
        return final_particles

__init__(simulator, prior, distance='l2', requires_iid_data=None, distance_kwargs=None, num_workers=1, simulation_batch_size=1, distance_batch_size=-1, show_progress_bars=True, kernel='gaussian', algorithm_variant='C')

Sequential Monte Carlo Approximate Bayesian Computation.

We distinguish between three different SMC methods here

• A: Toni et al. 2010 (Phd Thesis)
• B: Sisson et al. 2007 (with correction from 2009)
• C: Beaumont et al. 2009

In Toni et al. 2010 we find an overview of the differences on page 34: - B: same as A except for resampling of weights if the effective sampling size is too small. - C: same as A except for calculation of the covariance of the perturbation kernel: the kernel covariance is a scaled version of the covariance of the previous population.

Parameters:

Name	Type	Description	Default
simulator	Callable	A function that takes parameters \(\theta\) and maps them to simulations, or observations, x, \(\mathrm{sim}(\theta)\to x\). Any regular Python callable (i.e. function or class with __call__ method) can be used.	required
prior	Distribution	A probability distribution that expresses prior knowledge about the parameters, e.g. which ranges are meaningful for them. Any object with .log_prob()and .sample() (for example, a PyTorch distribution) can be used.	required
distance	Union[str, Callable]	Distance function to compare observed and simulated data. Can be a custom callable function or one of l1, l2, mse, mmd, wasserstein.	'l2'
requires_iid_data	Optional[bool]	Whether to allow conditioning on iid sampled data or not. Typically, this information is inferred by the choice of the distance, but in case a custom distance is used, this information is pivotal.	None
distance_kwargs	Optional[Dict]	Configurations parameters for the distances. In particular useful for the MMD and Wasserstein distance.	None
num_workers	int	Number of parallel workers to use for simulations.	1
simulation_batch_size	int	Number of parameter sets that the simulator maps to data x at once. If None, we simulate all parameter sets at the same time. If >= 1, the simulator has to process data of shape (simulation_batch_size, parameter_dimension).	1
distance_batch_size	int	Number of simulations that the distance function evaluates against the reference observations at once. If -1, we evaluate all simulations at the same time.	-1
show_progress_bars	bool	Whether to show a progressbar during simulation and sampling.	True
kernel	Optional[str]	Perturbation kernel.	'gaussian'
algorithm_variant	str	Indicating the choice of algorithm variant, A, B, or C.	'C'

Source code in sbi/inference/abc/smcabc.py

def __init__(
    self,
    simulator: Callable,
    prior: Distribution,
    distance: Union[str, Callable] = "l2",
    requires_iid_data: Optional[bool] = None,
    distance_kwargs: Optional[Dict] = None,
    num_workers: int = 1,
    simulation_batch_size: int = 1,
    distance_batch_size: int = -1,
    show_progress_bars: bool = True,
    kernel: Optional[str] = "gaussian",
    algorithm_variant: str = "C",
):
    r"""Sequential Monte Carlo Approximate Bayesian Computation.

    We distinguish between three different SMC methods here:
        - A: Toni et al. 2010 (Phd Thesis)
        - B: Sisson et al. 2007 (with correction from 2009)
        - C: Beaumont et al. 2009

    In Toni et al. 2010 we find an overview of the differences on page 34:
        - B: same as A except for resampling of weights if the effective sampling
            size is too small.
        - C: same as A except for calculation of the covariance of the perturbation
            kernel: the kernel covariance is a scaled version of the covariance of
            the previous population.

    Args:
        simulator: A function that takes parameters $\theta$ and maps them to
            simulations, or observations, `x`, $\mathrm{sim}(\theta)\to x$. Any
            regular Python callable (i.e. function or class with `__call__` method)
            can be used.
        prior: A probability distribution that expresses prior knowledge about the
            parameters, e.g. which ranges are meaningful for them. Any
            object with `.log_prob()`and `.sample()` (for example, a PyTorch
            distribution) can be used.
        distance: Distance function to compare observed and simulated data. Can be
            a custom callable function or one of `l1`, `l2`, `mse`,
            `mmd`, `wasserstein`.
        requires_iid_data: Whether to allow conditioning on iid sampled data or not.
            Typically, this information is inferred by the choice of the distance,
            but in case a custom distance is used, this information is pivotal.
        distance_kwargs: Configurations parameters for the distances. In particular
            useful for the MMD and Wasserstein distance.
        num_workers: Number of parallel workers to use for simulations.
        simulation_batch_size: Number of parameter sets that the simulator
            maps to data x at once. If None, we simulate all parameter sets at the
            same time. If >= 1, the simulator has to process data of shape
            (simulation_batch_size, parameter_dimension).
        distance_batch_size: Number of simulations that the distance function
            evaluates against the reference observations at once. If -1, we evaluate
            all simulations at the same time.
        show_progress_bars: Whether to show a progressbar during simulation and
            sampling.
        kernel: Perturbation kernel.
        algorithm_variant: Indicating the choice of algorithm variant, A, B, or C.
    """

    super().__init__(
        simulator=simulator,
        prior=prior,
        distance=distance,
        requires_iid_data=requires_iid_data,
        distance_kwargs=distance_kwargs,
        num_workers=num_workers,
        simulation_batch_size=simulation_batch_size,
        distance_batch_size=distance_batch_size,
        show_progress_bars=show_progress_bars,
    )

    kernels = ("gaussian", "uniform")
    assert kernel in kernels, (
        f"Kernel '{kernel}' not supported. Choose one from {kernels}."
    )
    self.kernel = kernel

    algorithm_variants = ("A", "B", "C")
    assert algorithm_variant in algorithm_variants, (
        f"SMCABC variant '{algorithm_variant}' not supported, choose one from"
        " {algorithm_variants}."
    )
    self.algorithm_variant = algorithm_variant
    self.distance_to_x0 = None
    self.simulation_counter = 0
    self.num_simulations = 0
    self.kernel_variance = None

    # Define simulator that keeps track of budget.
    def simulate_with_budget(theta):
        self.simulation_counter += theta.shape[0]
        return self._batched_simulator(theta)

    self._simulate_with_budget = simulate_with_budget

sample_from_population_with_weights(particles, weights, num_samples=1) staticmethod

Return samples from particles sampled with weights.

Source code in sbi/inference/abc/smcabc.py

@staticmethod
def sample_from_population_with_weights(
    particles: Tensor, weights: Tensor, num_samples: int = 1
) -> Tensor:
    """Return samples from particles sampled with weights."""

    # define multinomial with weights as probs
    multi = Multinomial(probs=weights)
    # sample num samples, with replacement
    samples = multi.sample(sample_shape=torch.Size((num_samples,)))
    # get indices of success trials
    indices = torch.where(samples)[1]
    # return those indices from trace
    return particles[indices]

Helpers

simulate_for_sbi(simulator, proposal, num_simulations, num_workers=1, simulation_batch_size=1, seed=None, show_progress_bar=True)

Returns pairs :math:(\theta, x) by sampling proposal and running simulations.

This function performs two steps:

• Sample parameters \(\theta\) from the proposal.
• Simulate these parameters to obtain \(x\).

Parameters:

Name	Type	Description	Default
simulator	Callable	A function that takes parameters \(\theta\) and maps them to simulations, or observations, x, \(\text{sim}(\theta)\to x\). Any regular Python callable (i.e. function or class with __call__ method) can be used. Note that the simulator should be able to handle numpy arrays for efficient parallelization. You can use process_simulator to ensure this.	required
proposal	Any	Probability distribution that the parameters \(\theta\) are sampled from.	required
num_simulations	int	Number of simulations that are run.	required
num_workers	int	Number of parallel workers to use for simulations.	1
simulation_batch_size	Union[int, None]	Number of parameter sets of shape (simulation_batch_size, parameter_dimension) that the simulator receives per call. If None, we set simulation_batch_size=num_simulations and simulate all parameter sets with one call. Otherwise, we construct batches of parameter sets and distribute them among num_workers.	1
seed	Optional[int]	Seed for reproducibility.	None
show_progress_bar	bool	Whether to show a progress bar for simulating. This will not affect whether there will be a progressbar while drawing samples from the proposal.	True

Returns: Sampled parameters \(\theta\) and simulation-outputs \(x\).

Source code in sbi/utils/simulation_utils.py

def simulate_for_sbi(
    simulator: Callable,
    proposal: Any,
    num_simulations: int,
    num_workers: int = 1,
    simulation_batch_size: Union[int, None] = 1,
    seed: Optional[int] = None,
    show_progress_bar: bool = True,
) -> Tuple[Tensor, Tensor]:
    r"""Returns pairs :math:`(\theta, x)` by sampling proposal and running simulations.

    This function performs two steps:

    - Sample parameters $\theta$ from the `proposal`.
    - Simulate these parameters to obtain $x$.

    Args:
        simulator: A function that takes parameters $\theta$ and maps them to
            simulations, or observations, `x`, $\text{sim}(\theta)\to x$. Any
            regular Python callable (i.e. function or class with `__call__` method)
            can be used. Note that the simulator should be able to handle numpy
            arrays for efficient parallelization. You can use
            `process_simulator` to ensure this.
        proposal: Probability distribution that the parameters $\theta$ are sampled
            from.
        num_simulations: Number of simulations that are run.
        num_workers: Number of parallel workers to use for simulations.
        simulation_batch_size: Number of parameter sets of shape
            (simulation_batch_size, parameter_dimension) that the simulator
            receives per call. If None, we set
            simulation_batch_size=num_simulations and simulate all parameter
            sets with one call. Otherwise, we construct batches of parameter
            sets and distribute them among num_workers.
        seed: Seed for reproducibility.
        show_progress_bar: Whether to show a progress bar for simulating. This will not
            affect whether there will be a progressbar while drawing samples from the
            proposal.

    Returns: Sampled parameters $\theta$ and simulation-outputs $x$.
    """

    if num_simulations == 0:
        theta = torch.tensor([], dtype=float32)
        x = torch.tensor([], dtype=float32)

    else:
        # Cast theta to numpy for better joblib performance (seee #1175)
        seed_all_backends(seed)
        theta = proposal.sample((num_simulations,))

        # Parse the simulation_batch_size logic
        if simulation_batch_size is None:
            simulation_batch_size = num_simulations
        else:
            simulation_batch_size = min(simulation_batch_size, num_simulations)

        if num_workers != 1:
            # For multiprocessing, we want to switch to numpy arrays.
            # The batch size will be an approximation, since np.array_split does
            # not take as argument the size of the batch but their total.
            num_batches = num_simulations // simulation_batch_size
            batches = np.array_split(theta.cpu().numpy(), num_batches, axis=0)
            batch_seeds = np.random.randint(low=0, high=1_000_000, size=(len(batches),))

            # define seeded simulator.
            def simulator_seeded(theta: ndarray, seed: int) -> Tensor:
                seed_all_backends(seed)
                return simulator(theta)

            try:  # catch TypeError to give more informative error message
                simulation_outputs: list[Tensor] = [  # pyright: ignore
                    xx
                    for xx in tqdm(
                        Parallel(return_as="generator", n_jobs=num_workers)(
                            delayed(simulator_seeded)(batch, seed)
                            for batch, seed in zip(batches, batch_seeds, strict=False)
                        ),
                        total=num_simulations,
                        disable=not show_progress_bar,
                    )
                ]
            except TypeError as err:
                raise TypeError(
                    "There is a TypeError error in your simulator function. Note: For"
                    " multiprocessing, we switch to numpy arrays. Besides confirming"
                    " your simulator works correctly, make sure to preprocess your"
                    " simulator with `process_simulator` to handle numpy arrays."
                ) from err

        else:
            simulation_outputs: list[Tensor] = []
            batches = torch.split(theta, simulation_batch_size)
            for batch in tqdm(batches, disable=not show_progress_bar):
                simulation_outputs.append(simulator(batch))

        # Correctly format the output
        x = torch.cat(simulation_outputs, dim=0)
        theta = torch.as_tensor(theta, dtype=float32)

    return theta, x

process_prior(prior, custom_prior_wrapper_kwargs=None)

Return PyTorch distribution-like prior from user-provided prior.

NOTE: If the prior argument is a sequence of PyTorch distributions, they will be interpreted as independent prior dimensions wrapped in a :class:MultipleIndependent PyTorch Distribution. In case the elements are not PyTorch distributions, make sure to use :func:process_prior on each element in the list beforehand.

NOTE: Returns a tuple (processed_prior, num_params, whether_prior_returns_numpy). The last two entries in the tuple can be passed on to process_simulator to prepare the simulator. For example, it ensures parameters are cast to numpy or adds a batch dimension to the simulator output, if needed.

Parameters:

Name	Type	Description	Default
prior (		class:~torch.distributions.distribution.Distribution or Sequence[:class:~torch.distributions.distribution.Distribution]): Prior object with .sample() and .log_prob(), or a sequence of such objects.	required
custom_prior_wrapper_kwargs	dict	Additional arguments passed to the wrapper class that processes the prior into a PyTorch Distribution, such as bounds (lower_bound, upper_bound) or argument constraints (arg_constraints).	None

Raises:

Type	Description
AttributeError	If prior objects lack .sample() or .log_prob().

Returns:

Type	Description
Tuple[Distribution, int, bool]	Tuple[torch.distributions.Distribution, int, bool]: - prior: A PyTorch-compatible prior. - theta_numel: Dimensionality of a single sample from the prior. - prior_returns_numpy: Whether the prior originally returned NumPy arrays.

Example:

::

import torch
from torch.distributions import Uniform
from sbi.utils.user_input_checks import process_prior

prior = Uniform(torch.zeros(1), torch.ones(1))
prior, theta_numel, prior_returns_numpy = process_prior(prior)

Source code in sbi/utils/user_input_checks.py

def process_prior(
    prior: Union[Sequence[Distribution], Distribution, rv_frozen, multi_rv_frozen],
    custom_prior_wrapper_kwargs: Optional[Dict] = None,
) -> Tuple[Distribution, int, bool]:
    """
    Return PyTorch distribution-like prior from user-provided prior.

    NOTE: If the prior argument is a sequence of PyTorch distributions, \
    they will be interpreted as independent prior dimensions wrapped in a \
    :class:`MultipleIndependent` PyTorch Distribution. In case the elements \
    are not PyTorch distributions, make sure to use :func:`process_prior` \
    on each element in the list beforehand.

    NOTE: Returns a tuple `(processed_prior, num_params, \
        whether_prior_returns_numpy)`. The last two entries in the tuple\
        can be passed on to `process_simulator` to prepare the simulator. For \
        example, it ensures parameters are cast to numpy or adds a batch \
        dimension to the simulator output, if needed.

    Args:
        prior (:class:`~torch.distributions.distribution.Distribution` \
            or Sequence[:class:`~torch.distributions.distribution.Distribution`]):
            Prior object with `.sample()` and `.log_prob()`, or a sequence \
                of such objects.
        custom_prior_wrapper_kwargs (dict, optional):
            Additional arguments passed to the wrapper class that processes the prior
            into a PyTorch Distribution, such as bounds (`lower_bound`, `upper_bound`)
            or argument constraints (`arg_constraints`).

    Raises:
        AttributeError: If prior objects lack `.sample()` or `.log_prob()`.

    Returns:
        Tuple[torch.distributions.Distribution, int, bool]:
            - `prior`: A PyTorch-compatible prior.
            - `theta_numel`: Dimensionality of a single sample from the prior.
            - `prior_returns_numpy`: Whether the prior originally returned NumPy arrays.

    Example:
    --------

    ::

        import torch
        from torch.distributions import Uniform
        from sbi.utils.user_input_checks import process_prior

        prior = Uniform(torch.zeros(1), torch.ones(1))
        prior, theta_numel, prior_returns_numpy = process_prior(prior)
    """

    # If prior is a sequence, assume independent components and check as PyTorch prior.
    if isinstance(prior, Sequence):
        warnings.warn(
            f"Prior was provided as a sequence of {len(prior)} priors. They will be "
            "interpreted as independent of each other and matched in order to the "
            "components of the parameter.",
            stacklevel=2,
        )
        # process individual priors
        prior = [process_prior(p, custom_prior_wrapper_kwargs)[0] for p in prior]
        return process_pytorch_prior(MultipleIndependent(prior))

    if isinstance(prior, Distribution):
        return process_pytorch_prior(prior)

    # If prior is given as `scipy.stats` object, wrap as PyTorch.
    elif isinstance(prior, (rv_frozen, multi_rv_frozen)):
        raise NotImplementedError(
            "Passing a prior as scipy.stats object is deprecated. "
            "Please pass it as a PyTorch Distribution."
        )

    # Otherwise it is a custom prior - check for `.sample()` and `.log_prob()`.
    else:
        return process_custom_prior(prior, custom_prior_wrapper_kwargs)

process_simulator(user_simulator, prior, is_numpy_simulator)

Returns a simulator that meets the requirements for usage in sbi.

Parameters:

Name	Type	Description	Default
user_simulator	Callable	simulator provided by the user, possibly written in numpy.	required
prior	Distribution	prior as pytorch distribution or processed with :func:process_prior.	required
is_numpy_simulator	bool	whether the simulator needs theta in numpy types, returned from process_prior.	required

Returns:

Name	Type	Description
Callable	Callable	simulator: processed simulator that returns :class:torch.Tensor and can handle batches of parameters.

Example:

::

import torch
from sbi.utils.user_input_checks import process_simulator
from torch.distributions import Uniform
from sbi.utils.user_input_checks import process_prior

prior = Uniform(torch.zeros(1), torch.ones(1))
prior, theta_numel, prior_returns_numpy = process_prior(prior)
simulator = lambda theta: theta + 1
simulator = process_simulator(simulator, prior, prior_returns_numpy)

Source code in sbi/utils/user_input_checks.py

def process_simulator(
    user_simulator: Callable,
    prior: Distribution,
    is_numpy_simulator: bool,
) -> Callable:
    """Returns a simulator that meets the requirements for usage in sbi.

    Args:
        user_simulator (Callable):
            simulator provided by the user, possibly written in numpy.
        prior (torch.distributions.Distribution):
            prior as pytorch distribution or processed with :func:`process_prior`.
        is_numpy_simulator (bool):
            whether the simulator needs theta in numpy types, returned
            from `process_prior`.

    Returns:
        Callable:
            simulator: processed simulator that returns :class:`torch.Tensor` \
                and can handle batches of parameters.

    Example:
    --------

    ::

        import torch
        from sbi.utils.user_input_checks import process_simulator
        from torch.distributions import Uniform
        from sbi.utils.user_input_checks import process_prior

        prior = Uniform(torch.zeros(1), torch.ones(1))
        prior, theta_numel, prior_returns_numpy = process_prior(prior)
        simulator = lambda theta: theta + 1
        simulator = process_simulator(simulator, prior, prior_returns_numpy)
    """

    assert isinstance(user_simulator, Callable), "Simulator must be a function."

    joblib_simulator = wrap_as_joblib_efficient_simulator(
        user_simulator, prior, is_numpy_simulator
    )

    batch_simulator = ensure_batched_simulator(joblib_simulator, prior)

    return batch_simulator