add fontsize arguments

Author: LuSchumacher

bayesflow/diagnostics.py

@@ -1140,10 +1140,12 @@ def plot_confusion_matrix(
     pred_models,
     model_names=None,
     fig_size=(5, 5),
+    label_fontsize=16,
     title_fontsize=18,
+    value_fontsize=10,
     tick_fontsize=12,
-    xtick_rotation=None,
-    ytick_rotation=None,
+    xtick_rotation=0,
+    ytick_rotation=90,
     normalize=True,
     cmap=None,
     title=True,
@@ -1160,13 +1162,17 @@ def plot_confusion_matrix(
         The model names for nice plot titles. Inferred if None.
     fig_size       : tuple or None, optional, default: (5, 5)
         The figure size passed to the ``matplotlib`` constructor. Inferred if ``None``
+    label_fontsize    : int, optional, default: 16
+        The font size of the y-label and y-label texts
     title_fontsize : int, optional, default: 18
         The font size of the title text.
+    value_fontsize  : int, optional, default: 10
+        The font size of the text annotations and the colorbar tick labels.
     tick_fontsize  : int, optional, default: 12
         The font size of the axis label and model name texts.
-    xtick_rotation: int, optional, default: None
+    xtick_rotation: int, optional, default: 0
         Rotation of x-axis tick labels (helps with long model names).
-    ytick_rotation: int, optional, default: None
+    ytick_rotation: int, optional, default: 90
         Rotation of y-axis tick labels (helps with long model names).
     normalize      : bool, optional, default: True
         A flag for normalization of the confusion matrix.
@@ -1201,8 +1207,10 @@ def plot_confusion_matrix(
     # Initialize figure
     fig, ax = plt.subplots(1, 1, figsize=fig_size)
 
-    im = ax.imshow(cm, interpolation="nearest", cmap=cmap)
-    ax.figure.colorbar(im, ax=ax, shrink=0.7)
+    im = ax.imshow(cm, interpolation="nearest", cmap=cmap, vmin=0.0, vmax=1.0)
+    cbar = ax.figure.colorbar(im, ax=ax, shrink=0.75)
+
+    cbar.ax.tick_params(labelsize=value_fontsize)
 
     ax.set(xticks=np.arange(cm.shape[1]), yticks=np.arange(cm.shape[0]))
     ax.set_xticklabels(model_names, fontsize=tick_fontsize)
@@ -1211,16 +1219,16 @@ def plot_confusion_matrix(
     ax.set_yticklabels(model_names, fontsize=tick_fontsize)
     if ytick_rotation:
         plt.yticks(rotation=ytick_rotation)
-    ax.set_xlabel("Predicted model", fontsize=tick_fontsize)
-    ax.set_ylabel("True model", fontsize=tick_fontsize)
+    ax.set_xlabel("Predicted model", fontsize=label_fontsize)
+    ax.set_ylabel("True model", fontsize=label_fontsize)
 
     # Loop over data dimensions and create text annotations
     fmt = ".2f" if normalize else "d"
     thresh = cm.max() / 2.0
     for i in range(cm.shape[0]):
         for j in range(cm.shape[1]):
             ax.text(
-                j, i, format(cm[i, j], fmt), ha="center", va="center", color="white" if cm[i, j] > thresh else "black"
+                j, i, format(cm[i, j], fmt), fontsize=value_fontsize, ha="center", va="center", color="white" if cm[i, j] > thresh else "black"
             )
     if title:
         ax.set_title("Confusion Matrix", fontsize=title_fontsize)