move compute_test_quantities to dict_utils, add docs

Author: vpratz

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -3,7 +3,8 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
-from ...utils.plot_utils import prepare_plot_data, add_titles_and_labels, prettify_subplots, compute_test_quantities
+from ...utils.dict_utils import compute_test_quantities
+from ...utils.plot_utils import prepare_plot_data, add_titles_and_labels, prettify_subplots
 from ...utils.ecdf import simultaneous_ecdf_bands
 from ...utils.ecdf.ranks import fractional_ranks, distance_ranks
 

bayesflow/diagnostics/plots/plot_quantity.py

@@ -3,13 +3,12 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from bayesflow.utils.dict_utils import make_variable_array, dicts_to_arrays, filter_kwargs
+from bayesflow.utils.dict_utils import make_variable_array, dicts_to_arrays, filter_kwargs, compute_test_quantities
 from bayesflow.utils.plot_utils import (
     add_titles_and_labels,
     make_figure,
     set_layout,
     prettify_subplots,
-    compute_test_quantities,
 )
 from bayesflow.utils.validators import check_estimates_prior_shapes
 
@@ -166,15 +165,21 @@ def plot_quantity(
 
 def _prepare_values(
     *,
-    values,
-    targets,
-    estimates,
-    variable_keys,
-    variable_names,
-    test_quantities,
-    label,
-    default_name,
+    values: Mapping[str, np.ndarray] | np.ndarray | Callable,
+    targets: Mapping[str, np.ndarray] | np.ndarray,
+    estimates: Mapping[str, np.ndarray] | np.ndarray | None,
+    variable_keys: Sequence[str],
+    variable_names: Sequence[str],
+    test_quantities: dict[str, Callable],
+    label: str | None,
+    default_name: str,
 ):
+    """
+    Provate helper function to compute/extract the values required for plotting
+    a quantity.
+
+    Refer to pairs_quantity and plot_quantity for details.
+    """
     is_values_callable = isinstance(values, Callable)
     # Optionally, compute and prepend test quantities from draws
     if test_quantities is not None:
@@ -190,7 +195,6 @@ def _prepare_values(
         estimates = updated_data["estimates"]
         targets = updated_data["targets"]
 
-    # input option 3
     if estimates is not None:
         if is_values_callable:
             values = values(estimates=estimates, targets=targets, **filter_kwargs({"aggregation": None}, values))
@@ -210,7 +214,6 @@ def _prepare_values(
         if test_quantities is None:
             variable_names = variable_names or estimates.variable_names
 
-    # input option 2
     if all([key in values for key in ["values", "metric_name", "variable_names"]]):
         # output of a metric function
         label = values["metric_name"] if label is None else label

bayesflow/utils/dict_utils.py

@@ -344,3 +344,55 @@ def squeeze_inner_estimates_dict(estimates):
         return estimates["value"]
     else:
         return estimates
+
+
+def compute_test_quantities(
+    targets: Mapping[str, np.ndarray] | np.ndarray,
+    estimates: Mapping[str, np.ndarray] | np.ndarray | None = None,
+    variable_keys: Sequence[str] = None,
+    variable_names: Sequence[str] = None,
+    test_quantities: dict[str, Callable] = None,
+):
+    """Compute additional test quantities for given targets and estimates."""
+    import keras
+
+    test_quantities_estimates = {} if estimates is not None else None
+    test_quantities_targets = {}
+
+    for key, test_quantity_fn in test_quantities.items():
+        # Apply test_quantity_func to ground-truths
+        tq_targets = test_quantity_fn(data=targets)
+        test_quantities_targets[key] = np.expand_dims(tq_targets, axis=1)
+
+        if estimates is not None:
+            # Flatten estimates for batch processing in test_quantity_fn, apply function, and restore shape
+            num_conditions, num_samples = next(iter(estimates.values())).shape[:2]
+            flattened_estimates = keras.tree.map_structure(
+                lambda t: np.reshape(t, (num_conditions * num_samples, *t.shape[2:]))
+                if isinstance(t, np.ndarray)
+                else t,
+                estimates,
+            )
+            flat_tq_estimates = test_quantity_fn(data=flattened_estimates)
+            test_quantities_estimates[key] = np.reshape(flat_tq_estimates, (num_conditions, num_samples, 1))
+
+    # Add custom test quantities to variable keys and names for plotting
+    # keys and names are set to the test_quantities dict keys
+    test_quantities_names = list(test_quantities.keys())
+
+    if variable_keys is None:
+        variable_keys = list(estimates.keys() if estimates is not None else targets.keys())
+    if isinstance(variable_names, list):
+        variable_names = test_quantities_names + variable_names
+
+    variable_keys = test_quantities_names + variable_keys
+    if estimates is not None:
+        estimates = test_quantities_estimates | estimates
+    targets = test_quantities_targets | targets
+
+    return {
+        "variable_keys": variable_keys,
+        "estimates": estimates,
+        "targets": targets,
+        "variable_names": variable_names,
+    }

bayesflow/utils/plot_utils.py

@@ -1,5 +1,4 @@
 from typing import Sequence, Any, Mapping
-from collections.abc import Callable
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -427,54 +426,3 @@ def create_legends(
             frameon=False,
             fontsize=legend_fontsize,
         )
-
-
-def compute_test_quantities(
-    targets: Mapping[str, np.ndarray] | np.ndarray,
-    estimates: Mapping[str, np.ndarray] | np.ndarray | None = None,
-    variable_keys: Sequence[str] = None,
-    variable_names: Sequence[str] = None,
-    test_quantities: dict[str, Callable] = None,
-):
-    import keras
-
-    test_quantities_estimates = {} if estimates is not None else None
-    test_quantities_targets = {}
-
-    for key, test_quantity_fn in test_quantities.items():
-        # Apply test_quantity_func to ground-truths
-        tq_targets = test_quantity_fn(data=targets)
-        test_quantities_targets[key] = np.expand_dims(tq_targets, axis=1)
-
-        if estimates is not None:
-            # Flatten estimates for batch processing in test_quantity_fn, apply function, and restore shape
-            num_conditions, num_samples = next(iter(estimates.values())).shape[:2]
-            flattened_estimates = keras.tree.map_structure(
-                lambda t: np.reshape(t, (num_conditions * num_samples, *t.shape[2:]))
-                if isinstance(t, np.ndarray)
-                else t,
-                estimates,
-            )
-            flat_tq_estimates = test_quantity_fn(data=flattened_estimates)
-            test_quantities_estimates[key] = np.reshape(flat_tq_estimates, (num_conditions, num_samples, 1))
-
-    # Add custom test quantities to variable keys and names for plotting
-    # keys and names are set to the test_quantities dict keys
-    test_quantities_names = list(test_quantities.keys())
-
-    if variable_keys is None:
-        variable_keys = list(estimates.keys() if estimates is not None else targets.keys())
-    if isinstance(variable_names, list):
-        variable_names = test_quantities_names + variable_names
-
-    variable_keys = test_quantities_names + variable_keys
-    if estimates is not None:
-        estimates = test_quantities_estimates | estimates
-    targets = test_quantities_targets | targets
-
-    return {
-        "variable_keys": variable_keys,
-        "estimates": estimates,
-        "targets": targets,
-        "variable_names": variable_names,
-    }