Enhance loss plot, including viridisifying the val loss

Author: jerrymhuang

bayesflow/diagnostics/plots/loss.py

@@ -7,21 +7,26 @@
 
 import keras.src.callbacks
 
-from ...utils.plot_utils import make_figure, add_titles_and_labels
+from matplotlib.colors import Normalize
+from ...utils.plot_utils import make_figure, add_titles_and_labels, gradient_line, gradient_legend
 
 
 def loss(
  history: keras.callbacks.History,
  train_key: str = "loss",
  val_key: str = "val_loss",
- moving_average: bool = False,
+ moving_average: bool = True,
  per_training_step: bool = False,
- ma_window_fraction: float = 0.01,
+ moving_average_span: int = 10,
  figsize: Sequence[float] = None,
  train_color: str = "#132a70",
- val_color: str = "black",
+ val_color: str = None,
+ val_colormap: str = 'viridis',
  lw_train: float = 2.0,
  lw_val: float = 3.0,
+ val_marker_type: str = "o",
+ val_marker_size: int = 34,
+ grid_alpha: float = 0.2,
  legend_fontsize: int = 14,
  label_fontsize: int = 14,
  title_fontsize: int = 16,
@@ -39,7 +44,7 @@ def loss(
  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 a moving average line of the train_losses.
+ A flag for adding an exponential moving average line of the train_losses.
  per_training_step : bool, optional, default: False
  A flag for making loss trajectory detailed (to training steps) rather than per epoch.
  ma_window_fraction : int, optional, default: 0.01
@@ -99,27 +104,51 @@ 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[:, i], color=train_color, lw=lw_train, alpha=0.9, label="Training")
- if moving_average and train_losses.columns[i] == "Loss":
- moving_average_window = int(train_losses.shape[0] * ma_window_fraction)
- smoothed_loss = train_losses.iloc[:, i].rolling(window=moving_average_window).mean()
+ ax.plot(train_step_index, train_losses.iloc[:, 0], color=train_color, lw=lw_train, alpha=0.2, label="Training")
+ if moving_average:
+ smoothed_loss = train_losses.iloc[:, 0].ewm(span=moving_average_span, adjust=True).mean()
  ax.plot(train_step_index, smoothed_loss, color="grey", lw=lw_train, label="Training (Moving Average)")
 
  # Plot optional val curve
  if val_losses is not None:
- if i < val_losses.shape[1]:
- ax.plot(
- val_step_index,
- val_losses.iloc[:, i],
- linestyle="--",
- marker="o",
- color=val_color,
- lw=lw_val,
- label="Validation",
- )
+ if val_color is not None:
+ ax.plot(
+ val_step_index,
+ val_losses.iloc[:, 0],
+ linestyle="--",
+ marker=val_marker_type,
+ color=val_color,
+ lw=lw_val,
+ label="Validation",
+ )
+ else:
+ # Create line segments between each epoch
+ points = np.array([val_step_index, val_losses.iloc[:,0]]).T.reshape(-1, 1, 2)
+ segments = np.concatenate([points[:-1], points[1:]], axis=1)
+
+ # Normalize color based on loss values
+ lc = gradient_line(
+ val_step_index,
+ val_losses.iloc[:,0],
+ c=val_step_index,
+ cmap=val_colormap,
+ lw=lw_val,
+ ax=ax
+ )
+ scatter = ax.scatter(
+ val_step_index,
+ val_losses.iloc[:,0],
+ c=val_step_index,
+ cmap=val_colormap,
+ marker=val_marker_type,
+ s=val_marker_size,
+ zorder=10,
+ edgecolors='none',
+ label='Validation'
+ )
 
  sns.despine(ax=ax)
- ax.grid(alpha=0.5)
+ ax.grid(alpha=grid_alpha)
 
  # Only add legend if there is a validation curve
  if val_losses is not None or moving_average:

bayesflow/utils/plot_utils.py

@@ -4,6 +4,11 @@
 import matplotlib.pyplot as plt
 import seaborn as sns
 
+from matplotlib.collections import LineCollection
+from matplotlib.colors import Normalize
+from matplotlib.patches import Rectangle
+from matplotlib.legend_handler import HandlerPatch
+
 from .validators import check_estimates_prior_shapes
 from .dict_utils import dicts_to_arrays
 
@@ -260,3 +265,69 @@ def make_quadratic(ax: plt.Axes, x_data: np.ndarray, y_data: np.ndarray):
  alpha=0.9,
  linestyle="dashed",
  )
+
+
+def gradient_line(x, y, c=None, cmap='viridis', lw=2, ax=None):
+ """
+ Plot a 1D line with color gradient determined by `c` (same shape as x and y).
+ """
+ if ax is None:
+ ax = plt.gca()
+
+ # Default color value = y
+ if c is None:
+ c = y
+
+ # Create segments for LineCollection
+ points = np.array([x, y]).T.reshape(-1, 1, 2)
+ segments = np.concatenate([points[:-1], points[1:]], axis=1)
+
+ norm = Normalize(np.min(c), np.max(c))
+ lc = LineCollection(segments, array=c, cmap=cmap, norm=norm, linewidth=lw)
+
+ ax.add_collection(lc)
+ ax.set_xlim(np.min(x), np.max(x))
+ ax.set_ylim(np.min(y), np.max(y))
+ return lc
+
+
+def gradient_legend(ax, label, cmap, norm, loc='upper right'):
+ """
+ Adds a single gradient swatch to the legend of the given Axes.
+
+ Parameters:
+ - ax: matplotlib Axes
+ - label: str, label to display in the legend
+ - cmap: matplotlib colormap
+ - norm: matplotlib Normalize object
+ - loc: legend location (default 'upper right')
+ """
+
+ # Custom dummy handle to represent the gradient
+ class _GradientSwatch(Rectangle): pass
+
+ # Custom legend handler that draws a horizontal gradient
+ class _HandlerGradient(HandlerPatch):
+ def create_artists(self, legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans):
+ gradient = np.linspace(0, 1, 256).reshape(1, -1)
+ im = ax.imshow(
+ gradient,
+ aspect='auto',
+ extent=[xdescent, xdescent + width, ydescent, ydescent + height],
+ transform=trans,
+ cmap=cmap,
+ norm=norm
+ )
+ return [im]
+
+ # Add to existing legend entries
+ handles, labels = ax.get_legend_handles_labels()
+ handles.append(_GradientSwatch((0, 0), 1, 1))
+ labels.append(label)
+
+ ax.legend(
+ handles=handles,
+ labels=labels,
+ loc=loc,
+ handler_map={_GradientSwatch: _HandlerGradient()}
+ )