adapt docs, minor stylistic changes

Author: vpratz

bayesflow/diagnostics/plots/recovery.py

@@ -1,4 +1,4 @@
-from collections.abc import Sequence, Mapping
+from collections.abc import Sequence, Mapping, Callable
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -12,10 +12,10 @@ def recovery(
     targets: Mapping[str, np.ndarray] | np.ndarray,
     variable_keys: Sequence[str] = None,
     variable_names: Sequence[str] = None,
-    point_agg=np.median,
-    uncertainty_agg=credible_interval,
-    point_agg_kwargs=None,
-    uncertainty_agg_kwargs=None,
+    point_agg: Callable = np.median,
+    uncertainty_agg: Callable = credible_interval,
+    point_agg_kwargs: dict = None,
+    uncertainty_agg_kwargs: dict = None,
     add_corr: bool = True,
     figsize: Sequence[int] = None,
     label_fontsize: int = 16,
@@ -58,13 +58,14 @@ def recovery(
        By default, select all keys.
     variable_names    : list or None, optional, default: None
         The individual parameter names for nice plot titles. Inferred if None
-    point_agg         : function to compute point estimates. Default: median
-    uncertainty_agg   : function to compute uncertainty interval bounds.
-        Default: credible_interval with coverage probability 95%.
+    point_agg         : callable, optional, default: median
+        Function to compute point estimates.
+    uncertainty_agg   : callable, optional, default: credible_interval with coverage probability 95%
+        Function to compute uncertainty interval bounds.
     point_agg_kwargs : Optional dictionary of further arguments passed to point_agg.
     uncertainty_agg_kwargs : Optional dictionary of further arguments passed to uncertainty_agg.
         For example, to change the coverage probability of credible_interval to 50%,
-        use uncertainty_agg_kwargs = dict(prob = 0.5)
+        use uncertainty_agg_kwargs = dict(prob=0.5)
     add_corr          : boolean, default: True
         Should correlations between estimates and ground truth values be shown?
     figsize           : tuple or None, optional, default : None
@@ -112,11 +113,8 @@ def recovery(
     estimates = plot_data.pop("estimates")
     targets = plot_data.pop("targets")
 
-    if point_agg_kwargs is None:
-        point_agg_kwargs = {}
-
-    if uncertainty_agg_kwargs is None:
-        uncertainty_agg_kwargs = {}
+    point_agg_kwargs = point_agg_kwargs or {}
+    uncertainty_agg_kwargs = uncertainty_agg_kwargs or {}
 
     # Compute point estimates and uncertainties
     point_estimate = point_agg(estimates, axis=1, **point_agg_kwargs)