implement different input modes for plotting functions

Author: vpratz

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -1,10 +1,9 @@
 from collections.abc import Callable, Mapping, Sequence
 
 import numpy as np
-import keras
 import matplotlib.pyplot as plt
 
-from ...utils.plot_utils import prepare_plot_data, add_titles_and_labels, prettify_subplots
+from ...utils.plot_utils import prepare_plot_data, add_titles_and_labels, prettify_subplots, compute_test_quantities
 from ...utils.ecdf import simultaneous_ecdf_bands
 from ...utils.ecdf.ranks import fractional_ranks, distance_ranks
 
@@ -136,38 +135,17 @@ def calibration_ecdf(
 
     # Optionally, compute and prepend test quantities from draws
     if test_quantities is not None:
-        test_quantities_estimates = {}
-        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)
-
-            # 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 isinstance(variable_names, list):
-            variable_names = test_quantities_names + variable_names
-
-        variable_keys = test_quantities_names + variable_keys
-        estimates = test_quantities_estimates | estimates
-        targets = test_quantities_targets | targets
+        updated_data = compute_test_quantities(
+            targets=targets,
+            estimates=estimates,
+            variable_keys=variable_keys,
+            variable_names=variable_names,
+            test_quantities=test_quantities,
+        )
+        variable_names = updated_data["variable_names"]
+        variable_keys = updated_data["variable_keys"]
+        estimates = updated_data["estimates"]
+        targets = updated_data["targets"]
 
     plot_data = prepare_plot_data(
         estimates=estimates,

bayesflow/diagnostics/plots/pairs_quantity.py

@@ -1,4 +1,4 @@
-from collections.abc import Sequence, Mapping
+from collections.abc import Callable, Sequence, Mapping
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -7,18 +7,22 @@
 import pandas as pd
 import seaborn as sns
 
-from bayesflow.utils.dict_utils import make_variable_array
+
+from .plot_quantity import _prepare_values
 
 
 def pairs_quantity(
-    values: Mapping[str, np.ndarray] | np.ndarray,
+    values: Mapping[str, np.ndarray] | np.ndarray | Callable,
+    *,
     targets: Mapping[str, np.ndarray] | np.ndarray,
     variable_keys: Sequence[str] = None,
     variable_names: Sequence[str] = None,
+    estimates: Mapping[str, np.ndarray] | np.ndarray | None = None,
+    test_quantities: dict[str, Callable] = None,
     height: float = 2.5,
     cmap: str | matplotlib.colors.Colormap = "viridis",
     alpha: float = 0.9,
-    label: str = "",
+    label: str = None,
     label_fontsize: int = 14,
     tick_fontsize: int = 12,
     colorbar_label_fontsize: int = 14,
@@ -28,6 +32,7 @@ def pairs_quantity(
     colorbar_offset: float = 0.06,
     vmin: float = None,
     vmax: float = None,
+    default_name: str = "v",
     **kwargs,
 ) -> sns.PairGrid:
     """
@@ -38,25 +43,59 @@ def pairs_quantity(
     each parameter is plotted on the diagonal. Each column displays the
     values of corresponding to the parameter in the column.
 
+    The function supports the following different combinations to pass
+    or compute the values:
+
+    1. pass `values` as an array of shape (num_datasets,) or (num_datasets, num_variables)
+    2. pass `values` as a dictionary with the keys 'values', 'metric_name' and 'variable_names'
+       as provided by the metrics functions. Note that the functions have to be called
+       without aggregation to obtain value per dataset.
+    3. pass a function to `values`, as well as `estimates`. The function should have the
+       signature fn(estimates, targets, [aggregation]) and return an object like the
+       `values` described in the previous options.
+
     Parameters
     ----------
-    values      : dict[str, np.ndarray],
-        The value of the quantity to plot.
-    targets     : dict[str, np.ndarray],
+    values      : dict[str, np.ndarray] | np.ndarray | Callable,
+        The value of the quantity to plot. One of the following:
+
+        1. an array of shape (num_datasets,) or (num_datasets, num_variables)
+        2. a dictionary with the keys 'values', 'metric_name' and 'variable_names'
+           as provided by the metrics functions. Note that the functions have to be called
+           without aggregation to obtain value per dataset.
+        3. a callable, requires passing `estimates` as well. The function should have the
+           signature fn(estimates, targets, [aggregation]) and return an object like the
+           ones described in the previous options.
+    targets     : dict[str, np.ndarray] | np.ndarray,
         The parameter values plotted on the axis.
     variable_keys       : list or None, optional, default: None
        Select keys from the dictionary provided in samples.
        By default, select all keys.
     variable_names    : list or None, optional, default: None
         The parameter names for nice plot titles. Inferred if None
+    estimates      : np.ndarray of shape (n_data_sets, n_post_draws, n_params), optional, default:  None
+        The posterior draws obtained from n_data_sets. Can only be supplied if
+        `values` is of type Callable.
+    test_quantities   : dict or None, optional, default: None
+        A dict that maps plot titles to functions that compute
+        test quantities based on estimate/target draws.
+
+        The dict keys are automatically added to ``variable_keys``
+        and ``variable_names``.
+        Test quantity functions are expected to accept a dict of draws with
+        shape ``(batch_size, ...)`` as the first (typically only)
+        positional argument and return an NumPy array of shape
+        ``(batch_size,)``.
+        The functions do not have to deal with an additional
+        sample dimension, as appropriate reshaping is done internally.
     height      : float, optional, default: 2.5
         The height of the pair plot
     cmap       : str or Colormap, default: "viridis"
         The colormap for the plot.
     alpha       : float in [0, 1], optional, default: 0.9
         The opacity of the plot
-    label       : str, optional, default: ""
-        Label for the dataset to plot
+    label       : str, optional, default: None
+        Label for the dataset to plot.
     label_fontsize    : int, optional, default: 14
         The font size of the x and y-label texts (parameter names)
     tick_fontsize     : int, optional, default: 12
@@ -77,21 +116,44 @@ def pairs_quantity(
     vmax : float, optional, default: None
         Maximum value for the colormap. If None, the maximum value is
         determined from `values`.
+    default_name      : str, optional (default = "v")
+        The default name to use for estimates if None provided
     **kwargs    : dict, optional
         Additional keyword arguments passed to the sns.PairGrid constructor
+
+    Returns
+    -------
+    plt.Figure
+        The figure instance
+
+    Raises
+    ------
+    ValueError
+        If a callable is supplied as `values`, but `estimates` is None.
     """
-    values = make_variable_array(
-        values,
+
+    if isinstance(values, Callable) and estimates is None:
+        raise ValueError("Supplied a callable as `values`, but not `estimates`.")
+
+    d = _prepare_values(
+        values=values,
+        targets=targets,
+        estimates=estimates,
         variable_keys=variable_keys,
         variable_names=variable_names,
+        test_quantities=test_quantities,
+        label=label,
+        default_name=default_name,
     )
-    variable_names = values.variable_names
-    variable_keys = values.variable_keys
-    targets = make_variable_array(
-        targets,
-        variable_keys=variable_keys,
-        variable_names=variable_names,
+    (values, targets, variable_keys, variable_names, test_quantities, label) = (
+        d["values"],
+        d["targets"],
+        d["variable_keys"],
+        d["variable_names"],
+        d["test_quantities"],
+        d["label"],
     )
+
     # Convert samples to pd.DataFrame
     data_to_plot = pd.DataFrame(targets, columns=variable_names)
 
@@ -110,11 +172,12 @@ def pairs_quantity(
     dim = g.axes.shape[0]
     for i in range(dim):
         for j in range(dim):
+            # if one value for each variable is supplied, use it for the corresponding column
+            row_values = values[:, j] if values.ndim == 2 else values
+
             if i == j:
                 ax = g.axes[i, j].twinx()
-                ax.scatter(
-                    targets[:, i], values[:, i], c=values[:, i], cmap=cmap, s=4, vmin=vmin, vmax=vmax, alpha=alpha
-                )
+                ax.scatter(targets[:, i], values[:, i], c=row_values, cmap=cmap, s=4, vmin=vmin, vmax=vmax, alpha=alpha)
                 ax.spines["left"].set_visible(False)
                 ax.spines["top"].set_visible(False)
                 ax.tick_params(axis="both", which="major", labelsize=tick_fontsize)
@@ -132,7 +195,7 @@ def pairs_quantity(
                 g.axes[i, j].scatter(
                     targets[:, j],
                     targets[:, i],
-                    c=values[:, j],
+                    c=row_values,
                     cmap=cmap,
                     s=4,
                     vmin=vmin,