Aesthetisize plot_calibration_curves

Author: stefanradev93

bayesflow/diagnostics.py

@@ -1015,6 +1015,7 @@ def plot_calibration_curves(
     legend_fontsize=14,
     title_fontsize=18,
     tick_fontsize=12,
+    epsilon=0.02,
     fig_size=None,
     color="#8f2727",
 ):
@@ -1040,6 +1041,8 @@ def plot_calibration_curves(
         The font size of the title text. Only relevant if `stacked=False`
     tick_fontsize     : int, optional, default: 12
         The font size of the axis ticklabels
+    epsilon           : float, optional, default: 0.02
+        A small amount to pad the [0, 1]-bounded axes from both side.
     fig_size          : tuple or None, optional, default: None
         The figure size passed to the ``matplotlib`` constructor. Inferred if ``None``
     color             : str, optional, default: '#8f2727'
@@ -1073,26 +1076,31 @@ def plot_calibration_curves(
         ax = axarr
     for j in range(num_models):
         # Plot calibration curve
-        ax[j].plot(probs_pred[j], probs_true[j], color=color)
-
-        # Plot AB line
-        ax[j].plot(ax[j].get_xlim(), ax[j].get_xlim(), "--", color="darkgrey")
+        ax[j].plot(probs_pred[j], probs_true[j], "o-", color=color)
 
         # Plot PMP distribution over bins
         uniform_bins = np.linspace(0.0, 1.0, num_bins + 1)
         norm_weights = np.ones_like(pred_models) / len(pred_models)
         ax[j].hist(pred_models[:, j], bins=uniform_bins, weights=norm_weights[:, j], color="grey", alpha=0.3)
 
+        # Plot AB line
+        ax[j].plot((0, 1), (0, 1), "--", color="black", alpha=0.9)
+
         # Tweak plot
+        ax[j].tick_params(axis="both", which="major", labelsize=tick_fontsize)
+        ax[j].tick_params(axis="both", which="minor", labelsize=tick_fontsize)
+        ax[j].set_title(model_names[j], fontsize=title_fontsize)
         ax[j].spines["right"].set_visible(False)
         ax[j].spines["top"].set_visible(False)
-        ax[j].set_xlim([0, 1])
-        ax[j].set_ylim([0, 1])
+        ax[j].set_xlim([0 - epsilon, 1 + epsilon])
+        ax[j].set_ylim([0 - epsilon, 1 + epsilon])
         ax[j].set_xlabel("Predicted probability", fontsize=label_fontsize)
         ax[j].set_ylabel("True probability", fontsize=label_fontsize)
-        ax[j].set_xticks([0.2, 0.4, 0.6, 0.8, 1.0])
-        ax[j].set_yticks([0.2, 0.4, 0.6, 0.8, 1.0])
+        ax[j].set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
+        ax[j].set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
         ax[j].grid(alpha=0.5)
+
+        # Add ECE label
         ax[j].text(
             0.1,
             0.9,
@@ -1102,11 +1110,7 @@ def plot_calibration_curves(
             transform=ax[j].transAxes,
             size=legend_fontsize,
         )
-        ax[j].tick_params(axis="both", which="major", labelsize=tick_fontsize)
-        ax[j].tick_params(axis="both", which="minor", labelsize=tick_fontsize)
 
-        # Set title
-        ax[j].set_title(model_names[j], fontsize=title_fontsize)
     fig.tight_layout()
     return fig