Add kwargs for agg functions

Author: stefanradev93

bayesflow/diagnostics/plots/recovery.py

@@ -14,7 +14,6 @@ def recovery(
     variable_names: Sequence[str] = None,
     point_agg=np.median,
     uncertainty_agg=credible_interval,
-    prob=0.95,
     add_corr: bool = True,
     figsize: Sequence[int] = None,
     label_fontsize: int = 16,
@@ -108,10 +107,10 @@ def recovery(
     targets = plot_data.pop("targets")
 
     # Compute point estimates and uncertainties
-    point_estimate = point_agg(estimates, axis=1)
+    point_estimate = point_agg(estimates, axis=1, **kwargs.get("point_agg_kwargs", {}))
 
     if uncertainty_agg is not None:
-        u = uncertainty_agg(estimates, prob=prob, axis=1)
+        u = uncertainty_agg(estimates, axis=1, **kwargs.get("uncertainty_agg_kwargs", {}))
         # compute lower and upper error
         u[0, :, :] = point_estimate - u[0, :, :]
         u[1, :, :] = u[1, :, :] - point_estimate

bayesflow/utils/numpy_utils.py

@@ -45,7 +45,7 @@ def softplus(x: np.ndarray, beta: float = 1.0, threshold: float = 20.0) -> np.nd
     return np.where(beta * x > threshold, x, np.log1p(exp_beta_x) / beta)
 
 
-def credible_interval(x: np.ndarray, prob: float = 0.95, axis: Sequence[int] = None, **kwargs) -> np.ndarray:
+def credible_interval(x: np.ndarray, prob: float = 0.95, axis: Sequence[int] | int = None, **kwargs) -> np.ndarray:
     """
     Compute credible interval from samples using quantiles.
 
@@ -69,7 +69,7 @@ def credible_interval(x: np.ndarray, prob: float = 0.95, axis: Sequence[int] = N
     --------
     >>> import numpy as np
     >>> # Simulate posterior samples
-    >>> samples = np.random.normal(10, 1000, 3)
+    >>> samples = np.random.normal(size=(10, 1000, 3))
 
     >>> # Different coverage probabilities
     >>> credible_interval(samples, prob=0.5, axis=1)  # 50% CI