rename 's' to 'markersize', add tests

Author: vpratz

bayesflow/diagnostics/plots/pairs_quantity.py

@@ -22,7 +22,7 @@ def pairs_quantity(
     height: float = 2.5,
     cmap: str | matplotlib.colors.Colormap = "viridis",
     alpha: float = 0.9,
-    s: float = 8.0,
+    markersize: float = 8.0,
     marker: str = "o",
     label: str = None,
     label_fontsize: int = 14,
@@ -81,6 +81,7 @@ def pairs_quantity(
     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.
+        Can only be supplied if `values` is a function.
 
         The dict keys are automatically added to ``variable_keys``
         and ``variable_names``.
@@ -96,7 +97,7 @@ def pairs_quantity(
         The colormap for the plot.
     alpha       : float in [0, 1], optional, default: 0.9
         The opacity of the plot
-    s                 : float, optional, default: 8.0
+    markersize        : float, optional, default: 8.0
         The marker size in points**2 for the scatter plot.
     marker            : str, optional, default: 'o'
         The marker for the scatter plot.
@@ -139,7 +140,13 @@ def pairs_quantity(
     """
 
     if isinstance(values, Callable) and estimates is None:
-        raise ValueError("Supplied a callable as `values`, but not `estimates`.")
+        raise ValueError("Supplied a callable as `values`, but no `estimates`.")
+    if not isinstance(values, Callable) and test_quantities is not None:
+        raise ValueError(
+            "Supplied `test_quantities`, but `values` is not a function. "
+            "As the values have to be calculated for the test quantities, "
+            "passing a function is required."
+        )
 
     d = _prepare_values(
         values=values,
@@ -188,7 +195,7 @@ def pairs_quantity(
                     values[:, i],
                     c=row_values,
                     cmap=cmap,
-                    s=s,
+                    s=markersize,
                     marker=marker,
                     vmin=vmin,
                     vmax=vmax,
@@ -213,7 +220,7 @@ def pairs_quantity(
                     targets[:, i],
                     c=row_values,
                     cmap=cmap,
-                    s=s,
+                    s=markersize,
                     vmin=vmin,
                     vmax=vmax,
                     alpha=alpha,
@@ -252,4 +259,4 @@ def inches_to_figure(fig, values):
         g.axes[i, 0].set_ylabel(variable_names[i], fontsize=label_fontsize)
         g.axes[dim - 1, i].set_xlabel(variable_names[i], fontsize=label_fontsize)
 
-    return g.figure
+    return g

bayesflow/diagnostics/plots/plot_quantity.py

@@ -26,7 +26,7 @@ def plot_quantity(
     title_fontsize: int = 18,
     tick_fontsize: int = 12,
     color: str = "#132a70",
-    s: float = 25.0,
+    markersize: float = 25.0,
     marker: str = "o",
     alpha: float = 0.5,
     xlabel: str = "Ground truth",
@@ -74,6 +74,7 @@ def plot_quantity(
     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.
+        Can only be supplied if `values` is a function.
 
         The dict keys are automatically added to ``variable_keys``
         and ``variable_names``.
@@ -93,7 +94,7 @@ def plot_quantity(
         The font size of the axis ticklabels
     color             : str, optional, default: '#8f2727'
         The color for the true vs. estimated scatter points and error bars
-    s                 : float, optional, default: 25.0
+    markersize        : float, optional, default: 25.0
         The marker size in points**2 for the scatter plot.
     marker            : str, optional, default: 'o'
         The marker for the scatter plot.
@@ -117,7 +118,13 @@ def plot_quantity(
     """
 
     if isinstance(values, Callable) and estimates is None:
-        raise ValueError("Supplied a callable as `values`, but not `estimates`.")
+        raise ValueError("Supplied a callable as `values`, but no `estimates`.")
+    if not isinstance(values, Callable) and test_quantities is not None:
+        raise ValueError(
+            "Supplied `test_quantities`, but `values` is not a function. "
+            "As the values have to be calculated for the test quantities, "
+            "passing a function is required."
+        )
 
     d = _prepare_values(
         values=values,
@@ -152,7 +159,7 @@ def plot_quantity(
         if i >= num_variables:
             break
 
-        ax.scatter(targets[:, i], values[:, i], color=color, alpha=alpha, s=s, marker=marker)
+        ax.scatter(targets[:, i], values[:, i], color=color, alpha=alpha, s=markersize, marker=marker)
 
     prettify_subplots(axes, num_subplots=num_variables, tick_fontsize=tick_fontsize)
 
@@ -242,16 +249,10 @@ def _prepare_values(
             default_name=default_name,
         )
     except ValueError:
-        if test_quantities is not None and not is_values_callable:
-            raise ValueError(
-                "`test_quantities` requires specifying `values` as callable and passing `estimates "
-                "to enable the computation of the values for each test quantity."
-            )
         raise ValueError(
             "Length of 'variable_names' and number of variables do not match. "
             "Did you forget to specify `variable_keys`?"
         )
-
     variable_names = targets.variable_names
     variable_keys = targets.variable_keys
 
@@ -266,11 +267,6 @@ def _prepare_values(
             default_name=default_name,
         )
     except ValueError:
-        if test_quantities is not None and not is_values_callable:
-            raise ValueError(
-                "`test_quantities` requires specifying `values` as callable and passing `estimates "
-                "to enable the computation of the values for each test quantity."
-            )
         raise ValueError(
             "Length of 'variable_names' and number of variables do not match. "
             "Did you forget to specify `variable_keys`?"

tests/test_diagnostics/test_diagnostics_plots.py

@@ -138,6 +138,102 @@ def test_pairs_posterior(random_estimates, random_targets, random_priors):
         )
 
 
+def test_pairs_quantity(random_estimates, random_targets, random_priors):
+    # test test_quantities and label assignment
+    key = next(iter(random_estimates.keys()))
+    test_quantities = {
+        "a": lambda data: np.sum(data[key], axis=-1),
+        "b": lambda data: np.prod(data[key], axis=-1),
+    }
+    out = bf.diagnostics.plots.pairs_quantity(
+        values=bf.diagnostics.posterior_contraction,
+        estimates=random_estimates,
+        targets=random_targets,
+        test_quantities=test_quantities,
+    )
+
+    num_vars = num_variables(random_estimates) + len(test_quantities)
+    assert out.axes.shape == (num_vars, num_vars)
+    assert out.axes[0, 0].get_ylabel() == "a"
+    assert out.axes[2, 0].get_ylabel() == "beta_0"
+    assert out.axes[4, 4].get_xlabel() == "sigma"
+
+    values = bf.diagnostics.posterior_contraction(estimates=random_estimates, targets=random_targets, aggregation=None)
+
+    bf.diagnostics.plots.pairs_quantity(
+        values,
+        targets=random_targets,
+    )
+
+    raw_values = np.random.normal(size=values["values"].shape)
+    out = bf.diagnostics.plots.pairs_quantity(raw_values, targets=random_targets, variable_keys=["beta", "sigma"])
+    assert out.axes.shape == (3, 3)
+
+    with pytest.raises(ValueError):
+        bf.diagnostics.plots.pairs_quantity(raw_values, targets=random_targets)
+
+    with pytest.raises(ValueError):
+        bf.diagnostics.plots.pairs_quantity(
+            values=values,
+            estimates=random_estimates,
+            targets=random_targets,
+            test_quantities=test_quantities,
+        )
+
+    with pytest.raises(ValueError):
+        bf.diagnostics.plots.pairs_quantity(
+            values=bf.diagnostics.posterior_contraction,
+            targets=random_targets,
+        )
+
+
+def test_plot_quantity(random_estimates, random_targets, random_priors):
+    # test test_quantities and label assignment
+    key = next(iter(random_estimates.keys()))
+    test_quantities = {
+        "a": lambda data: np.sum(data[key], axis=-1),
+        "b": lambda data: np.prod(data[key], axis=-1),
+    }
+    out = bf.diagnostics.plots.plot_quantity(
+        values=bf.diagnostics.posterior_contraction,
+        estimates=random_estimates,
+        targets=random_targets,
+        test_quantities=test_quantities,
+    )
+
+    num_vars = num_variables(random_estimates) + len(test_quantities)
+    assert len(out.axes) == num_vars
+    assert out.axes[0].title._text == "a"
+
+    values = bf.diagnostics.posterior_contraction(estimates=random_estimates, targets=random_targets, aggregation=None)
+
+    bf.diagnostics.plots.plot_quantity(
+        values,
+        targets=random_targets,
+    )
+
+    raw_values = np.random.normal(size=values["values"].shape)
+    out = bf.diagnostics.plots.plot_quantity(raw_values, targets=random_targets, variable_keys=["beta", "sigma"])
+    assert len(out.axes) == 3
+
+    with pytest.raises(ValueError):
+        bf.diagnostics.plots.plot_quantity(raw_values, targets=random_targets)
+
+    with pytest.raises(ValueError):
+        bf.diagnostics.plots.plot_quantity(
+            values=values,
+            estimates=random_estimates,
+            targets=random_targets,
+            test_quantities=test_quantities,
+        )
+
+    with pytest.raises(ValueError):
+        bf.diagnostics.plots.plot_quantity(
+            values=bf.diagnostics.posterior_contraction,
+            targets=random_targets,
+        )
+
+
 def test_mc_calibration(pred_models, true_models, model_names):
     out = bf.diagnostics.plots.mc_calibration(pred_models, true_models, model_names=model_names)
     assert len(out.axes) == pred_models.shape[-1]