Merge remote-tracking branch 'upstream/main' into feature/lc2st-mlp-gpu-support

Author: Dev-Sudarshan

README.md

@@ -10,7 +10,7 @@
 ## `sbi`: Simulation-Based Inference
 
 [Getting Started](https://sbi.readthedocs.io/en/latest/tutorials/00_getting_started.html) |
-[Documentation](https://sbi.readthedocs.io/en/latest/) | [Discord Server](https://discord.gg/eEeVPSvWKy)
+[Documentation](https://sbi.readthedocs.io/en/latest/) | [Discord Server](https://discord.gg/VPkV7XPj7k)
 
 `sbi` is a Python package for simulation-based inference, designed to meet the needs of
 both researchers and practitioners. Whether you need fine-grained control or an

docs/how_to_guide.rst

@@ -57,6 +57,7 @@ Sampling
    how_to_guide/09_sampler_interface.ipynb
    how_to_guide/10_refine_posterior_with_importance_sampling.ipynb
    how_to_guide/11_iid_sampling_with_nle_or_nre.ipynb
+   how_to_guide/23_using_pyro_with_sbi.ipynb
 
 
 Diagnostics

docs/how_to_guide/23_using_pyro_with_sbi.ipynb

@@ -106,21 +106,47 @@ def sample(
         sample_shape: Shape = torch.Size(),
         x: Optional[Tensor] = None,
         show_progress_bars: bool = True,
-        mcmc_method: Optional[str] = None,
-        mcmc_parameters: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
     ) -> Tensor:
-        """See child classes for docstring."""
+        r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
+
+        Args:
+            sample_shape: Shape of samples to draw.
+            x: Conditioning observation $x_o$. If not provided, uses the default
+                `x` set via `.set_default_x()`.
+            show_progress_bars: Whether to show a progress bar during sampling.
+            **kwargs: Additional keyword arguments passed to the specific
+                posterior's sampling method. See the docstring of the specific
+                posterior class for available options.
+
+        Returns:
+            Samples from the posterior with shape `(*sample_shape, *theta_shape)`.
+        """
         pass
 
     @abstractmethod
     def sample_batched(
         self,
         sample_shape: Shape,
         x: Tensor,
-        max_sampling_batch_size: int = 10_000,
         show_progress_bars: bool = True,
+        **kwargs: Any,
     ) -> Tensor:
-        """See child classes for docstring."""
+        r"""Draw samples from the posteriors for a batch of different xs.
+
+        Given a batch of observations `[x_1, ..., x_B]`, this method samples from
+        posteriors $p(\theta|x_1), \ldots, p(\theta|x_B)$ in a vectorized manner.
+
+        Args:
+            sample_shape: Shape of samples to draw for each observation.
+            x: Batch of observations with shape `(batch_dim, *event_shape_x)`.
+            show_progress_bars: Whether to show a progress bar during sampling.
+            **kwargs: Additional keyword arguments passed to the specific
+                posterior's sampling method.
+
+        Returns:
+            Samples with shape `(*sample_shape, batch_dim, *theta_shape)`.
+        """
         pass
 
     @property

sbi/inference/posteriors/base_posterior.py

@@ -137,19 +137,20 @@ def sample(
         sample_shape: Shape = torch.Size(),
         x: Optional[Tensor] = None,
         max_sampling_batch_size: int = 10_000,
-        sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
         reject_outside_prior: bool = True,
         max_sampling_time: Optional[float] = None,
+        return_partial_on_timeout: bool = False,
     ) -> Tensor:
-        r"""Return samples from posterior distribution $p(\theta|x)$.
+        r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            sample_with: This argument only exists to keep backward-compatibility with
-                `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
+            max_sampling_batch_size: Maximum batch size for rejection sampling.
             show_progress_bars: Whether to show sampling progress monitor.
             reject_outside_prior: If True (default), rejection sampling is used to
                 ensure samples lie within the prior support. If False, samples are drawn
@@ -159,6 +160,10 @@ def sample(
                 If exceeded, sampling is aborted and a RuntimeError is raised. Only
                 applies when `reject_outside_prior=True` (no effect otherwise since
                 direct sampling is fast).
+            return_partial_on_timeout: If True and `max_sampling_time` is exceeded,
+                return the samples collected so far instead of raising a RuntimeError.
+                A warning will be issued. Only applies when `reject_outside_prior=True`
+                (default).
         """
         num_samples = torch.Size(sample_shape).numel()
         x = self._x_else_default_x(x)
@@ -180,13 +185,6 @@ def sample(
             else max_sampling_batch_size
         )
 
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                f"`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
-
         if reject_outside_prior:
             # Normal rejection behavior.
             samples = rejection.accept_reject_sample(
@@ -198,6 +196,7 @@ def sample(
                 proposal_sampling_kwargs={"condition": x},
                 alternative_method="build_posterior(..., sample_with='mcmc')",
                 max_sampling_time=max_sampling_time,
+                return_partial_on_timeout=return_partial_on_timeout,
             )[0]
         else:
             # Bypass rejection sampling entirely.
@@ -217,10 +216,12 @@ def sample_batched(
         show_progress_bars: bool = True,
         reject_outside_prior: bool = True,
         max_sampling_time: Optional[float] = None,
+        return_partial_on_timeout: bool = False,
     ) -> Tensor:
-        r"""Given a batch of observations [x_1, ..., x_B] this function samples from
-        posteriors $p(\theta|x_1)$, ... ,$p(\theta|x_B)$, in a batched (i.e. vectorized)
-        manner.
+        r"""Draw samples from the posteriors for a batch of different xs.
+
+        Given a batch of observations `[x_1, ..., x_B]`, this method samples from
+        posteriors $p(\theta|x_1), \ldots, p(\theta|x_B)$ in a vectorized manner.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from the posterior
@@ -236,6 +237,9 @@ def sample_batched(
             max_sampling_time: Optional maximum allowed sampling time in seconds.
                 If exceeded, sampling is aborted and a RuntimeError is raised. Only
                 applies when `reject_outside_prior=True`.
+            return_partial_on_timeout: If True and `max_sampling_time` is exceeded,
+                return the samples collected so far instead of raising a RuntimeError.
+                A warning will be issued. Only applies when `reject_outside_prior=True`.
 
         Returns:
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
@@ -282,6 +286,7 @@ def sample_batched(
                 proposal_sampling_kwargs={"condition": x},
                 alternative_method="build_posterior(..., sample_with='mcmc')",
                 max_sampling_time=max_sampling_time,
+                return_partial_on_timeout=return_partial_on_timeout,
             )[0]
         else:
             # Bypass rejection sampling entirely.

sbi/inference/posteriors/direct_posterior.py

@@ -190,14 +190,14 @@ def sample(
         method: Optional[str] = None,
         oversampling_factor: int = 32,
         max_sampling_batch_size: int = 10_000,
-        sample_with: Optional[str] = None,
         show_progress_bars: bool = False,
     ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
-        """Return samples from the approximate posterior distribution.
+        """Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
         Args:
             sample_shape: Shape of samples that are drawn from posterior.
-            x: Observed data.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
             method: Either of [`sir`|`importance`]. This sets the behavior of the
                 `.sample()` method. With `sir`, approximate posterior samples are
                 generated with sampling importance resampling (SIR). With
@@ -212,13 +212,6 @@ def sample(
 
         method = self.method if method is None else method
 
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                f"`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
-
         self.potential_fn.set_x(self._x_else_default_x(x))
 
         if method == "sir":
@@ -255,8 +248,6 @@ def _importance_sample(
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior.
-            sample_with: This argument only exists to keep backward-compatibility with
-                `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
             show_progress_bars: Whether to show sampling progress monitor.
 
         Returns:
@@ -286,13 +277,10 @@ def _sir_sample(
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            x: Observed data.
-            sample_with: This argument only exists to keep backward-compatibility with
-                `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
-            oversampling_factor: Number of proposed samples form which only one is
+            oversampling_factor: Number of proposed samples from which only one is
                 selected based on its importance weight.
-            max_sampling_batch_size: The batchsize of samples being drawn from
-                the proposal at every iteration. Used only in `sir_sample()`.
+            max_sampling_batch_size: The batch size of samples being drawn from
+                the proposal at every iteration.
             show_progress_bars: Whether to show sampling progress monitor.
 
         Returns:

sbi/inference/posteriors/importance_posterior.py

@@ -257,29 +257,36 @@ def sample(
         num_chains: Optional[int] = None,
         init_strategy: Optional[str] = None,
         init_strategy_parameters: Optional[Dict[str, Any]] = None,
-        init_strategy_num_candidates: Optional[int] = None,
-        mcmc_parameters: Optional[Dict] = None,
-        mcmc_method: Optional[str] = None,
-        sample_with: Optional[str] = None,
         num_workers: Optional[int] = None,
         mp_context: Optional[str] = None,
         show_progress_bars: bool = True,
     ) -> Tensor:
-        r"""Return samples from posterior distribution $p(\theta|x)$ with MCMC.
-
-        Check the `__init__()` method for a description of all arguments as well as
-        their default values.
+        r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            mcmc_parameters: Dictionary that is passed only to support the API of
-                `sbi` v0.17.2 or older.
-            mcmc_method: This argument only exists to keep backward-compatibility with
-                `sbi` v0.17.2 or older. Please use `method` instead.
-            sample_with: This argument only exists to keep backward-compatibility with
-                `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
+            method: MCMC method to use. One of `slice_np`, `slice_np_vectorized`,
+                `hmc_pyro`, `nuts_pyro`, `slice_pymc`, `hmc_pymc`, `nuts_pymc`.
+                If not provided, uses the method specified at initialization.
+            thin: Thinning factor for the chain. If not provided, uses the value
+                specified at initialization.
+            warmup_steps: Number of warmup steps to discard. If not provided, uses
+                the value specified at initialization.
+            num_chains: Number of MCMC chains to run. If not provided, uses the
+                value specified at initialization.
+            init_strategy: Initialization strategy for chains (`proposal`, `sir`,
+                or `resample`). If not provided, uses the value specified at
+                initialization.
+            init_strategy_parameters: Parameters for the initialization strategy.
+                If not provided, uses the value specified at initialization.
+            num_workers: Number of CPU cores for parallelization. If not provided,
+                uses the value specified at initialization.
+            mp_context: Multiprocessing context (`fork` or `spawn`). If not provided,
+                uses the value specified at initialization.
             show_progress_bars: Whether to show sampling progress monitor.
 
         Returns:
@@ -301,46 +308,6 @@ def sample(
             if init_strategy_parameters is None
             else init_strategy_parameters
         )
-        if init_strategy_num_candidates is not None:
-            warn(
-                f"Passing `init_strategy_num_candidates` is deprecated as of sbi \
-                v0.19.0. Instead, use e.g., \
-                `init_strategy_parameters={'num_candidate_samples': 1000}`",
-                stacklevel=2,
-            )
-            self.init_strategy_parameters["num_candidate_samples"] = (
-                init_strategy_num_candidates
-            )
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                "`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
-        if mcmc_method is not None:
-            warn(
-                "You passed `mcmc_method` to `.sample()`. As of sbi v0.18.0, this "
-                "is deprecated and will be removed in a future release. Use `method` "
-                "instead of `mcmc_method`.",
-                stacklevel=2,
-            )
-            method = mcmc_method
-        if mcmc_parameters:
-            warn(
-                "You passed `mcmc_parameters` to `.sample()`. As of sbi v0.18.0, this "
-                "is deprecated and will be removed in a future release. Instead, pass "
-                "the variable to `.sample()` directly, e.g. "
-                "`posterior.sample((1,), num_chains=5)`.",
-                stacklevel=2,
-            )
-        # The following lines are only for backwards compatibility with sbi v0.17.2 or
-        # older.
-        m_p = mcmc_parameters or {}  # define to shorten the variable name
-        method = _maybe_use_dict_entry(method, "mcmc_method", m_p)
-        thin = _maybe_use_dict_entry(thin, "thin", m_p)
-        warmup_steps = _maybe_use_dict_entry(warmup_steps, "warmup_steps", m_p)
-        num_chains = _maybe_use_dict_entry(num_chains, "num_chains", m_p)
-        init_strategy = _maybe_use_dict_entry(init_strategy, "init_strategy", m_p)
         self.potential_ = self._prepare_potential(method)  # type: ignore
 
         initial_params = self._get_initial_params(
@@ -415,9 +382,10 @@ def sample_batched(
         mp_context: Optional[str] = None,
         show_progress_bars: bool = True,
     ) -> Tensor:
-        r"""Given a batch of observations [x_1, ..., x_B] this function samples from
-        posteriors $p(\theta|x_1)$, ... ,$p(\theta|x_B)$, in a batched (i.e. vectorized)
-        manner.
+        r"""Draw samples from the posteriors for a batch of different xs.
+
+        Given a batch of observations `[x_1, ..., x_B]`, this method samples from
+        posteriors $p(\theta|x_1), \ldots, p(\theta|x_B)$ in a vectorized manner.
 
         Check the `__init__()` method for a description of all arguments as well as
         their default values.
@@ -1134,24 +1102,6 @@ def _process_thin_default(thin: int) -> int:
     return thin
 
 
-def _maybe_use_dict_entry(default: Any, key: str, dict_to_check: Dict) -> Any:
-    """Returns `default` if `key` is not in the dict and otherwise the dict entry.
-
-    This method exists only to keep backwards compatibility with `sbi` v0.17.2 or
-    older. It allows passing `mcmc_parameters` to `.sample()`.
-
-    Args:
-        default: The default value if `key` is not in `dict_to_check`.
-        key: The key for which to check in `dict_to_check`.
-        dict_to_check: The dictionary to be checked.
-
-    Returns:
-        The potentially replaced value.
-    """
-    attribute = dict_to_check.get(key, default)
-    return attribute
-
-
 def _num_required_args(func):
     """
     Utility for counting the number of positional args in a function.