Improve diagnostics

Author: stefanradev93

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -15,13 +15,13 @@ def calibration_ecdf(
     variable_keys: Sequence[str] = None,
     variable_names: Sequence[str] = None,
     test_quantities: dict[str, Callable] = None,
-    difference: bool = False,
+    difference: bool = True,
     stacked: bool = False,
     rank_type: str | np.ndarray = "fractional",
     figsize: Sequence[float] = None,
     label_fontsize: int = 16,
     legend_fontsize: int = 14,
-    legend_location: str = "upper right",
+    legend_location: str = "lower right",
     title_fontsize: int = 18,
     tick_fontsize: int = 12,
     rank_ecdf_color: str = "#132a70",
@@ -59,7 +59,7 @@ def calibration_ecdf(
         The posterior draws obtained from n_data_sets
     targets     : np.ndarray of shape (n_data_sets, n_params)
         The prior draws obtained for generating n_data_sets
-    difference        : bool, optional, default: False
+    difference        : bool, optional, default: True
         If `True`, plots the ECDF difference.
         Enables a more dynamic visualization range.
     stacked           : bool, optional, default: False
@@ -98,7 +98,9 @@ def calibration_ecdf(
     label_fontsize    : int, optional, default: 16
         The font size of the y-label and y-label texts
     legend_fontsize   : int, optional, default: 14
-        The font size of the legend text
+        The font size of the legend text.
+    legend_location : str, optional, default: 'lower right
+        The location of the legend.
     title_fontsize    : int, optional, default: 18
         The font size of the title text.
         Only relevant if `stacked=False`
@@ -211,11 +213,13 @@ def calibration_ecdf(
     else:
         titles = ["Stacked ECDFs"]
 
-    for ax, title in zip(plot_data["axes"].flat, titles):
+    for i, (ax, title) in enumerate(zip(plot_data["axes"].flat, titles)):
         ax.fill_between(z, L, U, color=fill_color, alpha=0.2, label=rf"{int((1 - alpha) * 100)}$\%$ Confidence Bands")
-        ax.legend(fontsize=legend_fontsize, loc=legend_location)
         ax.set_title(title, fontsize=title_fontsize)
 
+        if i == 0:
+            ax.legend(fontsize=legend_fontsize, loc=legend_location)
+
     prettify_subplots(plot_data["axes"], num_subplots=plot_data["num_variables"], tick_fontsize=tick_fontsize)
 
     add_titles_and_labels(

bayesflow/diagnostics/plots/calibration_ecdf_from_quantiles.py

@@ -14,12 +14,12 @@ def calibration_ecdf_from_quantiles(
     quantiles_key: str = "quantiles",
     variable_keys: Sequence[str] = None,
     variable_names: Sequence[str] = None,
-    difference: bool = False,
+    difference: bool = True,
     stacked: bool = False,
     figsize: Sequence[float] = None,
     label_fontsize: int = 16,
     legend_fontsize: int = 14,
-    legend_location: str = "upper right",
+    legend_location: str = "lower right",
     title_fontsize: int = 18,
     tick_fontsize: int = 12,
     rank_ecdf_color: str = "#132a70",
@@ -69,7 +69,7 @@ def calibration_ecdf_from_quantiles(
     variable_names    : list or None, optional, default: None
         The parameter names for nice plot titles.
         Inferred if None. Only relevant if `stacked=False`.
-    difference        : bool, optional, default: False
+    difference        : bool, optional, default: True
         If `True`, plots the ECDF difference.
         Enables a more dynamic visualization range.
     stacked           : bool, optional, default: False
@@ -82,7 +82,9 @@ def calibration_ecdf_from_quantiles(
     label_fontsize    : int, optional, default: 16
         The font size of the y-label and y-label texts
     legend_fontsize   : int, optional, default: 14
-        The font size of the legend text
+        The font size of the legend text.
+    legend_location : str, optional, default: 'lower right
+        The location of the legend.
     title_fontsize    : int, optional, default: 18
         The font size of the title text.
         Only relevant if `stacked=False`

bayesflow/diagnostics/plots/coverage.py

@@ -17,7 +17,7 @@ def coverage(
     legend_fontsize: int = 14,
     title_fontsize: int = 18,
     tick_fontsize: int = 12,
-    legend_location: str = "upper right",
+    legend_location: str = "lower right",
     color: str = "#132a70",
     num_col: int = None,
     num_row: int = None,
@@ -41,7 +41,7 @@ def coverage(
         The posterior draws obtained from num_datasets
     targets : np.ndarray of shape (num_datasets, num_params)
         The true parameter values used for generating num_datasets
-    difference : bool, optional, default: False
+    difference : bool, optional, default: True
         If True, plots the difference between empirical coverage and ideal coverage
         (coverage - width), making deviations from ideal calibration more visible.
         If False, plots the standard coverage plot.
@@ -60,6 +60,8 @@ def coverage(
         The font size of the title text
     tick_fontsize : int, optional, default: 12
         The font size of the axis ticklabels
+    legend_location : str, optional, default: 'upper right
+        The location of the legend.
     color : str, optional, default: '#132a70'
         The color for the coverage line
     num_row : int, optional, default: None
@@ -132,7 +134,7 @@ def coverage(
             )
 
             # Plot ideal coverage difference line (y = 0)
-            ax.axhline(y=0, color="skyblue", linewidth=2.0, label="Ideal Coverage")
+            ax.axhline(y=0, color="black", linestyle="dashed", label="Ideal Coverage")
 
             # Plot empirical coverage difference
             ax.plot(width_rep, diff_est, color=color, alpha=1.0, label="Coverage Difference")
@@ -149,23 +151,19 @@ def coverage(
             )
 
             # Plot ideal coverage line (y = x)
-            ax.plot([0, 1], [0, 1], color="skyblue", linewidth=2.0, label="Ideal Coverage")
+            ax.plot([0, 1], [0, 1], color="black", linestyle="dashed", label="Ideal Coverage")
 
             # Plot empirical coverage
             ax.plot(width_rep, coverage_est, color=color, alpha=1.0, label="Empirical Coverage")
 
-        # Set axis limits
-        ax.set_xlim(0, 1)
-        ax.set_ylim(0, 1)
-
         # Add legend to first subplot
         if i == 0:
             ax.legend(fontsize=legend_fontsize, loc=legend_location)
 
     prettify_subplots(plot_data["axes"], num_subplots=plot_data["num_variables"], tick_fontsize=tick_fontsize)
 
     # Add labels, titles, and set font sizes
-    ylabel = "Observed coverage difference" if difference else "Observed coverage"
+    ylabel = "Empirical coverage difference" if difference else "Empirical coverage"
     add_titles_and_labels(
         axes=plot_data["axes"],
         num_row=plot_data["num_row"],

bayesflow/diagnostics/plots/loss.py

@@ -168,7 +168,7 @@ def loss(
         num_col=1,
         title=["Loss Trajectory"],
         xlabel="Training epoch #",
-        ylabel="Value",
+        ylabel="Loss",
         title_fontsize=title_fontsize,
         label_fontsize=label_fontsize,
     )

bayesflow/diagnostics/plots/recovery.py

@@ -87,8 +87,10 @@ def recovery(
         The number of rows for the subplots. Dynamically determined if None.
     num_col           : int, optional, default: None
         The number of columns for the subplots. Dynamically determined if None.
-    xlabel :
-    ylabel :
+    xlabel            : str, optional, default: "Ground truth"
+        The label shown on the x-axis.
+    ylabel            : str, optional, default: "Estimate"
+        The label shown on the y-axis.
     markersize        : float, optional, default: None
         The marker size in points.
 

examples/Multimodal_Data.ipynb

@@ -524,7 +524,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "id": "2415fd0b-f5d6-4fc9-83d7-8952e6270186",
    "metadata": {},
    "outputs": [