Add examples to documentation (#1548)

* Example BoxUniform

* add imports BoxUniform example

* update documentation and add example for process_prior, process_simulator, BoxUniform, and MultipleIndependent

Author: jorobledo

docs/conf.py

@@ -40,6 +40,7 @@
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
+    "torch": ("https://pytorch.org/docs/stable/", None),
 }
 
 source_suffix = {'.rst': 'restructuredtext', '.myst': 'myst-nb', '.ipynb': 'myst-nb'}

sbi/utils/torchutils.py

@@ -285,23 +285,48 @@ def __init__(
     ):
         """Multidimensional uniform distribution defined on a box.
 
-        A `Uniform` distribution initialized with e.g. a parameter vector low or high of
-         length 3 will result in a /batch/ dimension of length 3. A log_prob evaluation
-         will then output three numbers, one for each of the independent Uniforms in
-         the batch. Instead, a `BoxUniform` initialized in the same way has three
-         /event/ dimensions, and returns a scalar log_prob corresponding to whether
-         the evaluated point is in the box defined by low and high or outside.
-
-        Refer to torch.distributions.Uniform and torch.distributions.Independent for
-         further documentation.
+        A :class:`~torch.distributions.uniform.Uniform` distribution initialized \
+        with e.g. a parameter vector low or high of length 3 will result \
+        in a *batch* dimension of length 3. A log_prob evaluation will then \
+        output three numbers, one for each of the independent Uniforms in the \
+        batch. Instead, a :class:`BoxUniform` initialized in the same way has three \
+        *event* dimensions, and returns a scalar log_prob corresponding to whether \
+        the evaluated point is in the box defined by low and high or outside.
+
+        Refer to :class:`~torch.distributions.uniform.Uniform`\
+        and :class:`~torch.distributions.independent.Independent` for \
+        further documentation.
 
         Args:
             low: lower range (inclusive).
             high: upper range (exclusive).
-            reinterpreted_batch_ndims (int): the number of batch dims to
-                reinterpret as event dims.
-            device: device of the prior, inferred from low arg, defaults to "cpu",
-                should match the training device when used in SBI.
+            reinterpreted_batch_ndims (int): the number of batch dims to \
+            reinterpret as event dims.
+            device (Optional): device of the prior, inferred from low arg, \
+            defaults to "cpu", should match the training device when used in SBI.
+
+        Example:
+        --------
+
+        ::
+
+            import torch
+            from sbi.utils.torchutils import BoxUniform
+
+            # Define lower bounds
+            low = torch.tensor([0.0, 0.0, 0.0])
+
+            # Define upper bounds
+            high = torch.tensor([1.0, 1.0, 1.0])
+
+            box_uniform = BoxUniform(low, high)
+
+            # Sample from the box_uniform
+            N_samples = 100
+            sample = box_uniform.sample((N_samples,))
+
+            # Evaluate the log probability of the sample
+            log_prob = box_uniform.log_prob(sample)
         """
 
         # Type checks.
@@ -342,6 +367,15 @@ def to(self, device: Union[str, torch.device]) -> None:
 
         Args:
             device: Target device (e.g., "cpu", "cuda", "mps").
+
+        Example:
+        --------
+
+        ::
+
+            device = "cuda"
+            prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2))
+            prior.to(device) #inplace
         """
         # Update the device attribute
         self.device = device

sbi/utils/user_input_checks.py

@@ -39,34 +39,51 @@ 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.
+    """
+    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 `MultipleIndependent`
-    pytorch Distribution. In case the elements are not PyTorch distributions, make sure
-    to use process_prior on each element in the list beforehand.
+    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 as well. For example, it will take care of casting parameters to numpy
-    or adding a batch dimension to the simulator output, if needed.
+    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: Prior object with `.sample()` and `.log_prob()` as provided by the user,
-            or a sequence of such objects.
-        custom_prior_wrapper_kwargs: kwargs to be passed to the class that wraps a
-            custom prior into a pytorch Distribution, e.g., for passing bounds for a
-            prior with bounded support (lower_bound, upper_bound), or argument
-            constraints.
-            (arg_constraints), see pytorch.distributions.Distribution for more info.
+        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 lacks `.sample()` or `.log_prob()`.
+        AttributeError: If prior objects lack `.sample()` or `.log_prob()`.
 
     Returns:
-        prior: Prior that emits samples and evaluates log prob as PyTorch Tensors.
-        theta_numel: Number of parameters - elements in a single sample from the prior.
-        prior_returns_numpy: Whether the return type of the prior was a Numpy array.
+        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.
@@ -456,14 +473,33 @@ def process_simulator(
     """Returns a simulator that meets the requirements for usage in sbi.
 
     Args:
-        user_simulator: simulator provided by the user, possibly written in numpy.
-        prior: prior as pytorch distribution or processed with `process_prior`.
-        is_numpy_simulator: whether the simulator needs theta in numpy types, returned
+        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:
-        simulator: processed simulator that returns `torch.Tensor` can handle batches
-            of parameters.
+        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."

sbi/utils/user_input_checks_utils.py

@@ -209,31 +209,43 @@ def to(self, device: Union[str, torch.device]) -> None:
 
 
 class MultipleIndependent(Distribution):
-    """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution.
-
-    Every element of the sequence is treated as independent from the other elements.
-    Single elements can be multivariate with dependent dimensions.
-
-    Example:
-
-    ::
-
-        import torch
-        from torch.distributions import Gamma, Beta, MultivariateNormal
-        prior = MultipleIndependent([
-            Gamma(torch.zeros(1), torch.ones(1)),
-            Beta(torch.zeros(1), torch.ones(1)),
-            MultivariateNormal(torch.ones(2), torch.tensor([[1, .1], [.1, 1.]]))
-        ])
-    """
+    """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution."""
 
     def __init__(
         self,
         dists: Sequence[Distribution],
-        validate_args=None,
+        validate_args: Optional[bool] = None,
         arg_constraints: Optional[Dict[str, constraints.Constraint]] = None,
         device: Optional[str] = None,
     ):
+        """Joint distribution of multiple independent :class:`torch.distributions`.
+
+        Every element of the sequence is treated as independent from the \
+        other elements. Single elements can be multivariate with dependent dimensions.
+
+        Args:
+            dists: Sequence of PyTorch distributions.
+            validate_args (Optional): If True, the distribution checks its parameters.
+            arg_constraints (Optional): Dictionary of constraints for the parameters \
+                of the distribution.
+            device (Optional): Device to move the distribution to. If None, \
+                the distribution is moved to the CPU.
+
+        Example:
+        --------
+
+        ::
+
+            import torch
+            from torch.distributions import Gamma, Beta, MultivariateNormal
+            from sbi.utils.user_input_checks_utils import MultipleIndependent
+
+            prior = MultipleIndependent([
+                Gamma(torch.zeros(1), torch.ones(1)),
+                Beta(torch.zeros(1), torch.ones(1)),
+                MultivariateNormal(torch.ones(2), torch.tensor([[1, .1], [.1, 1.]]))
+            ])
+        """
         self._check_distributions(dists)
         if validate_args is not None:
             [d.set_default_validate_args(validate_args) for d in dists]
@@ -280,9 +292,9 @@ def _check_distribution(self, dist: Distribution):
         )
         assert isinstance(
             dist, Distribution
-        ), """priors passed to MultipleIndependent must be PyTorch distributions. Make
-            sure to process custom priors individually using process_prior before
-            passing them in a list to process_prior."""
+        ), """priors passed to MultipleIndependent must be PyTorch distributions. Make \
+            sure to process custom priors individually using :func:`process_prior` \
+            before passing them in a list to :func:`process_prior`."""
         # Make sure batch shape is smaller or equal to 1.
         assert dist.batch_shape in (
             torch.Size([1]),