Refactor loss plotting logic and remove unused parameters.

Simplified the loss plotting code by consolidating duplication and aligning the handling of smoothing logic. Removed unused arguments like markers and colormap, reducing potential confusion in the API. Updated comments and improved code readability for maintainability.

Author: LarsKue

bayesflow/diagnostics/plots/loss.py

@@ -7,7 +7,7 @@
 
 import keras.src.callbacks
 
-from ...utils.plot_utils import make_figure, add_titles_and_labels, add_gradient_plot
+from ...utils.plot_utils import make_figure, add_titles_and_labels
 
 
 def loss(
@@ -17,13 +17,9 @@ def loss(
     smoothing_factor: float = 0.8,
     figsize: Sequence[float] = None,
     train_color: str = "#132a70",
-    val_color: str = None,
-    val_colormap: str = "viridis",
+    val_color: str = "black",
     lw_train: float = 2.0,
-    lw_val: float = 3.0,
-    marker: bool = True,
-    val_marker_type: str = ".",
-    val_marker_size: int = 34,
+    lw_val: float = 2.0,
     grid_alpha: float = 0.2,
     legend_fontsize: int = 14,
     label_fontsize: int = 14,
@@ -41,29 +37,19 @@ def loss(
         The training loss key to look for in the history
     val_key     : str, optional, default: "val_loss"
         The validation loss key to look for in the history
-    moving_average     : bool, optional, default: False
-        A flag for adding an exponential moving average line of the train_losses.
-    moving_average_alpha : int, optional, default: 0.8
-        Smoothing factor for the moving average.
+    smoothing_factor : float, optional, default: 0.8
+        If greater than zero, smooth the loss curves by applying an exponential moving average.
     figsize            : tuple or None, optional, default: None
         The figure size passed to the ``matplotlib`` constructor.
         Inferred if ``None``
     train_color        : str, optional, default: '#8f2727'
         The color for the train loss trajectory
     val_color          : str, optional, default: None
         The color for the optional validation loss trajectory
-    val_colormap       : str, optional, default: "viridis"
-        The colormap for the optional validation loss trajectory
     lw_train           : int, optional, default: 2
         The linewidth for the training loss curve
     lw_val             : int, optional, default: 3
         The linewidth for the validation loss curve
-    marker  : bool, optional, default: False
-        A flag for whether marker should be added in the validation loss trajectory
-    val_marker_type     : str, optional, default: o
-        The marker type for the validation loss curve
-    val_marker_size     : int, optional, default: 34
-        The marker size for the validation loss curve
     grid_alpha          : float, optional, default: 0.2
         The transparency of the background grid
     legend_fontsize    : int, optional, default: 14
@@ -108,68 +94,60 @@ def loss(
 
     # Loop through loss entries and populate plot
     for i, ax in enumerate(axes.flat):
-        # Plot train curve
-        ax.plot(train_step_index, train_losses.iloc[:, 0], color=train_color, lw=lw_train, alpha=0.05, label="Training")
-        if moving_average:
-            smoothed_train_loss = train_losses.iloc[:, 0].ewm(alpha=moving_average_alpha, adjust=True).mean()
-            ax.plot(train_step_index, smoothed_train_loss, color="grey", lw=lw_train, label="Training (Moving Average)")
+        if smoothing_factor > 0:
+            # plot unsmoothed train loss
+            ax.plot(
+                train_step_index, train_losses.iloc[:, 0], color=train_color, lw=lw_train, alpha=0.3, label="Training"
+            )
+
+            # plot smoothed train loss
+            smoothed_train_loss = train_losses.iloc[:, 0].ewm(alpha=1.0 - smoothing_factor, adjust=True).mean()
+            ax.plot(
+                train_step_index,
+                smoothed_train_loss,
+                color=train_color,
+                lw=lw_train,
+                alpha=0.8,
+                label="Training (Moving Average)",
+            )
+        else:
+            # plot unsmoothed train loss
+            ax.plot(
+                train_step_index, train_losses.iloc[:, 0], color=train_color, lw=lw_train, alpha=0.8, label="Training"
+            )
 
         # Plot optional val curve
         if val_losses is not None:
             if val_color is not None:
-                ax.plot(
-                    val_step_index,
-                    val_losses.iloc[:, 0],
-                    linestyle="--",
-                    marker=val_marker_type if marker else None,
-                    color=val_color,
-                    lw=lw_val,
-                    alpha=0.2,
-                    label="Validation",
-                )
-                if moving_average:
-                    smoothed_val_loss = val_losses.iloc[:, 0].ewm(alpha=moving_average_alpha, adjust=True).mean()
+                if smoothing_factor > 0:
+                    # plot unsmoothed val loss
+                    ax.plot(
+                        val_step_index, val_losses.iloc[:, 0], color=val_color, lw=lw_val, alpha=0.3, label="Validation"
+                    )
+
+                    # plot smoothed val loss
+                    smoothed_val_loss = val_losses.iloc[:, 0].ewm(alpha=1.0 - smoothing_factor, adjust=True).mean()
                     ax.plot(
                         val_step_index,
                         smoothed_val_loss,
                         color=val_color,
                         lw=lw_val,
+                        alpha=0.8,
                         label="Validation (Moving Average)",
                     )
-            else:
-                # Make gradient lines
-                add_gradient_plot(
-                    val_step_index,
-                    val_losses.iloc[:, 0],
-                    ax,
-                    val_colormap,
-                    lw_val,
-                    marker,
-                    val_marker_type,
-                    val_marker_size,
-                    alpha=0.05,
-                    label="Validation",
-                )
-                if moving_average:
-                    smoothed_val_loss = val_losses.iloc[:, 0].ewm(alpha=moving_average_alpha, adjust=True).mean()
-                    add_gradient_plot(
-                        val_step_index,
-                        smoothed_val_loss,
-                        ax,
-                        val_colormap,
-                        lw_val,
-                        marker,
-                        val_marker_type,
-                        val_marker_size,
-                        alpha=1,
-                        label="Validation (Moving Average)",
+                else:
+                    # plot unsmoothed val loss
+                    ax.plot(
+                        val_step_index, val_losses.iloc[:, 0], color=val_color, lw=lw_val, alpha=0.8, label="Validation"
                     )
 
         sns.despine(ax=ax)
         ax.grid(alpha=grid_alpha)
 
-        # Only add legend if there is a validation curve
-        if val_losses is not None or moving_average:
+        ax.set_xlim(train_step_index[0], train_step_index[-1])
+
+        # Only add the legend if there are multiple curves
+        if val_losses is not None or smoothing_factor > 0:
             ax.legend(fontsize=legend_fontsize)
 
     # Add labels, titles, and set font sizes