Utils Plt

plt

FUNCTION	DESCRIPTION
enable_mpl_params	Enable publication-ready matplotlib parameters for plotting.
set_figsize	Set figure dimensions to avoid scaling in LaTeX.
setup_axes	Configure axes with publication-ready formatting.
save_figure	Save figure in multiple formats with publication settings.
create_colormap	Create custom colormap for consistent publication figures.
set_color_palette	Set publication-appropriate color palette.
label_subplots	Add panel labels (a, b, c, etc.) to subplots.
plot_kde	Plot free energy from KDE calculation.
plot_figure	Create a publication-ready figure with optimal settings.
plot_training_metrics	Plot training metrics from parametric model training.
plot_loss_curves	Plot training and validation loss curves.
plot_metric_comparison	Compare a specific metric across multiple model runs or configurations.

Functions

enable_mpl_params(config_dict: dict | None = None) -> None

Enable publication-ready matplotlib parameters for plotting.

PARAMETER	DESCRIPTION
config_dict	Dictionary with matplotlib settings, if None uses the predefined publication settings. TYPE: dict, optional, by default None DEFAULT: None

Source code in spectre/utils/plt.py

def enable_mpl_params(config_dict: dict | None = None) -> None:
    """
    Enable publication-ready matplotlib parameters for plotting.

    Parameters
    ----------
    config_dict : dict, optional, by default None
        Dictionary with matplotlib settings, if None uses the predefined
        publication settings.
    """
    config_dict = mpl_dict if config_dict is None else config_dict
    try:
        plt.rcParams.update(config_dict)
    except Exception as e:
        if "text.usetex" in config_dict and config_dict["text.usetex"]:
            import warnings

            warnings.warn(f"LaTeX rendering failed: {e}. Disabling LaTeX.", UserWarning)
            config_dict = config_dict.copy()
            config_dict["text.usetex"] = False
            plt.rcParams.update(config_dict)
        else:
            raise

set_figsize(width: float | str = 490.0, fraction: float = 1.0, subplots: tuple[int, int] = (1, 1), aspect_ratio: float | str = 'golden') -> tuple[float, float]

Set figure dimensions to avoid scaling in LaTeX.

PARAMETER	DESCRIPTION
width	Figure width in points or preset ("textwidth", "columnwidth"). TYPE: float | str, optional, by default 490.0 DEFAULT: 490.0
fraction	Fraction of the total width. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
subplots	Number of rows and columns of subplots. TYPE: tuple[int, int], optional, by default (1, 1) DEFAULT: (1, 1)
aspect_ratio	Height to width ratio or preset ("golden", "square"). TYPE: float | str, optional, by default "golden" DEFAULT: 'golden'

RETURNS	DESCRIPTION
tuple[float, float]	Figure dimensions in inches (width, height).

Source code in spectre/utils/plt.py

def set_figsize(
    width: float | str = 490.0,
    fraction: float = 1.0,
    subplots: tuple[int, int] = (1, 1),
    aspect_ratio: float | str = "golden",
) -> tuple[float, float]:
    """
    Set figure dimensions to avoid scaling in LaTeX.

    Parameters
    ----------
    width : float | str, optional, by default 490.0
        Figure width in points or preset ("textwidth", "columnwidth").

    fraction : float, optional, by default 1.0
        Fraction of the total width.

    subplots : tuple[int, int], optional, by default (1, 1)
        Number of rows and columns of subplots.

    aspect_ratio : float | str, optional, by default "golden"
        Height to width ratio or preset ("golden", "square").

    Returns
    -------
    tuple[float, float]
        Figure dimensions in inches (width, height).
    """
    # Handle preset widths.
    width_presets = {
        "textwidth": 426.79135,  # Standard LaTeX textwidth.
        "columnwidth": 213.39566,  # Two-column format.
    }

    if isinstance(width, str):
        if width in width_presets:
            width = width_presets[width]
        else:
            raise ValueError(
                f"Unknown width preset '{width}'. Available: {list(width_presets.keys())}."
            )

    fig_width_pt = width * fraction
    inches_per_pt = 1 / 72.27  # Convert from pt to inches.
    fig_width_in = fig_width_pt * inches_per_pt

    # Handle aspect ratios.
    if isinstance(aspect_ratio, str):
        if aspect_ratio == "golden":
            ratio = (5**0.5 - 1) / 2  # Golden ratio is approximately 0.618.
        elif aspect_ratio == "square":
            ratio = 1.0
        else:
            raise ValueError(
                f"Unknown aspect ratio '{aspect_ratio}'. Use 'golden', 'square', or a float."
            )
    else:
        ratio = aspect_ratio

    fig_height_in = fig_width_in * ratio * (subplots[0] / subplots[1])

    return fig_width_in, fig_height_in

setup_axes(ax: plt.Axes, xlabel: str = '', ylabel: str = '', title: str = '', grid: bool = True, minor_ticks: bool = True) -> None

Configure axes with publication-ready formatting.

PARAMETER	DESCRIPTION
ax	Matplotlib axes object to configure. TYPE: Axes
xlabel	X-axis label. TYPE: str, optional, by default "" DEFAULT: ''
ylabel	Y-axis label. TYPE: str, optional, by default "" DEFAULT: ''
title	Plot title. TYPE: str, optional, by default "" DEFAULT: ''
grid	Enable grid. TYPE: bool, optional, by default True DEFAULT: True
minor_ticks	Enable minor ticks. TYPE: bool, optional, by default True DEFAULT: True

Source code in spectre/utils/plt.py

def setup_axes(
    ax: plt.Axes,
    xlabel: str = "",
    ylabel: str = "",
    title: str = "",
    grid: bool = True,
    minor_ticks: bool = True,
) -> None:
    """
    Configure axes with publication-ready formatting.

    Parameters
    ----------
    ax : plt.Axes
        Matplotlib axes object to configure.

    xlabel : str, optional, by default ""
        X-axis label.

    ylabel : str, optional, by default ""
        Y-axis label.

    title : str, optional, by default ""
        Plot title.

    grid : bool, optional, by default True
        Enable grid.

    minor_ticks : bool, optional, by default True
        Enable minor ticks.
    """
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)

    if grid:
        ax.grid(True, alpha=0.3, linestyle="-", linewidth=0.5)
        ax.set_axisbelow(True)

    if minor_ticks:
        ax.minorticks_on()

save_figure(fig: plt.Figure, filename: str, formats: str | list[str] = 'pdf', dpi: int = 1000, bbox_inches: str = 'tight', pad_inches: float = 0.1) -> None

Save figure in multiple formats with publication settings.

PARAMETER	DESCRIPTION
fig	Matplotlib figure object to save. TYPE: Figure
filename	Base filename (without extension). TYPE: str
formats	Output format(s). TYPE: str | list[str], optional, by default "pdf" DEFAULT: 'pdf'
dpi	Resolution for raster formats. TYPE: int, optional, by default 1000 DEFAULT: 1000
bbox_inches	Bounding box specification. TYPE: str, optional, by default "tight" DEFAULT: 'tight'
pad_inches	Padding around figure. TYPE: float, optional, by default 0.1 DEFAULT: 0.1

Source code in spectre/utils/plt.py

def save_figure(
    fig: plt.Figure,
    filename: str,
    formats: str | list[str] = "pdf",
    dpi: int = 1000,
    bbox_inches: str = "tight",
    pad_inches: float = 0.1,
) -> None:
    """
    Save figure in multiple formats with publication settings.

    Parameters
    ----------
    fig : plt.Figure
        Matplotlib figure object to save.

    filename : str
        Base filename (without extension).

    formats : str | list[str], optional, by default "pdf"
        Output format(s).

    dpi : int, optional, by default 1000
        Resolution for raster formats.

    bbox_inches : str, optional, by default "tight"
        Bounding box specification.

    pad_inches : float, optional, by default 0.1
        Padding around figure.
    """
    if isinstance(formats, str):
        formats = [formats]

    for fmt in formats:
        output_file = f"{filename}.{fmt}"
        fig.savefig(
            output_file,
            format=fmt,
            dpi=dpi,
            bbox_inches=bbox_inches,
            pad_inches=pad_inches,
        )

create_colormap(colors: list[str], name: str = 'custom', n_colors: int = 256) -> mpl.colors.LinearSegmentedColormap

Create custom colormap for consistent publication figures.

PARAMETER	DESCRIPTION
colors	List of color values (hex, named colors, etc.). TYPE: list[str]
name	Colormap name. TYPE: str, optional, by default "custom" DEFAULT: 'custom'
n_colors	Number of color levels. TYPE: int, optional, by default 256 DEFAULT: 256

RETURNS	DESCRIPTION
LinearSegmentedColormap	Custom colormap object.

Examples:

>>> cmap = create_colormap(["#1f77b4", "#ff7f0e", "#2ca02c"])
>>> plt.contourf(X, Y, Z, cmap=cmap)

Source code in spectre/utils/plt.py

def create_colormap(
    colors: list[str],
    name: str = "custom",
    n_colors: int = 256,
) -> mpl.colors.LinearSegmentedColormap:
    """
    Create custom colormap for consistent publication figures.

    Parameters
    ----------
    colors : list[str]
        List of color values (hex, named colors, etc.).

    name : str, optional, by default "custom"
        Colormap name.

    n_colors : int, optional, by default 256
        Number of color levels.

    Returns
    -------
    mpl.colors.LinearSegmentedColormap
        Custom colormap object.

    Examples
    --------
    >>> cmap = create_colormap(["#1f77b4", "#ff7f0e", "#2ca02c"])
    >>> plt.contourf(X, Y, Z, cmap=cmap)
    """
    return mpl.colors.LinearSegmentedColormap.from_list(name, colors, N=n_colors)

get_colormap(name: str, n_colors: int, set_over_color: str | None = None) -> mpl.colors.LinearSegmentedColormap

A set of custom colormaps.

PARAMETER	DESCRIPTION
name	Colormap name. Available: "kanso" (and "kanso_r"): https://github.com/webhooked/kanso.nvim TYPE: str
n_colors	Number of color levels. TYPE: int
set_over_color	Color to set over the maximum color value of a plot. TYPE: str | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
LinearSegmentedColormap	Custom colormap object.

Source code in spectre/utils/plt.py

def get_colormap(
    name: str, n_colors: int, set_over_color: str | None = None
) -> mpl.colors.LinearSegmentedColormap:
    """
    A set of custom colormaps.

    Parameters
    ----------
    name : str
        Colormap name.

        Available:

        - "kanso" (and "kanso_r"): https://github.com/webhooked/kanso.nvim

    n_colors : int
        Number of color levels.

    set_over_color : str | None, by default None
        Color to set over the maximum color value of a plot.

    Returns
    -------
    mpl.colors.LinearSegmentedColormap
        Custom colormap object.
    """
    avail_cm = ["kanso", "kanso_r"]
    if name not in avail_cm:
        raise ValueError(f"Unknown colormap: {name}. Select from {avail_cm}.")

    if name == "kanso":
        cm = create_colormap(
            ["#23262D", "navy", "#8BA4B0", "lightblue"], n_colors=n_colors
        )
    elif name == "kanso_r":
        cm = create_colormap(
            ["lightblue", "#8BA4B0", "navy", "#23262D"], n_colors=n_colors
        )

    if set_over_color:
        cm.set_over(set_over_color)

    return cm

set_color_palette(palette: Literal['scientific', 'colorblind', 'print'] = 'scientific') -> list[str]

Set publication-appropriate color palette.

PARAMETER	DESCRIPTION
palette	Color palette type. TYPE: str, by default "scientific" DEFAULT: 'scientific'

RETURNS	DESCRIPTION
list[str]	List of color hex codes.

Examples:

>>> colors = set_color_palette("colorblind")
>>> for i, color in enumerate(colors):
...     plt.plot(x, y + i, color=color, label=f"Series {i}")

Source code in spectre/utils/plt.py

def set_color_palette(
    palette: Literal["scientific", "colorblind", "print"] = "scientific",
) -> list[str]:
    """
    Set publication-appropriate color palette.

    Parameters
    ----------
    palette : str, by default "scientific"
        Color palette type.

    Returns
    -------
    list[str]
        List of color hex codes.

    Examples
    --------
    >>> colors = set_color_palette("colorblind")
    >>> for i, color in enumerate(colors):
    ...     plt.plot(x, y + i, color=color, label=f"Series {i}")
    """
    palettes = {
        "scientific": [
            "#1f77b4",
            "#ff7f0e",
            "#2ca02c",
            "#d62728",
            "#9467bd",
            "#8c564b",
            "#e377c2",
            "#7f7f7f",
            "#bcbd22",
            "#17becf",
        ],
        "colorblind": [
            "#0173b2",
            "#de8f05",
            "#029e73",
            "#cc78bc",
            "#ca9161",
            "#fbafe4",
            "#949494",
            "#ece133",
            "#56b4e9",
            "#009e73",
        ],
        "print": [
            "#000000",
            "#525252",
            "#969696",
            "#cccccc",
            "#f7f7f7",
            "#3182bd",
            "#6baed6",
            "#9ecae1",
            "#c6dbef",
            "#e6f0ff",
        ],
    }

    if palette not in palettes:
        raise ValueError(
            f"Unknown palette '{palette}'. Available: {list(palettes.keys())}."
        )

    colors = palettes[palette]
    plt.rcParams["axes.prop_cycle"] = plt.cycler("color", colors)

    return colors

label_subplots(axes: plt.Axes | np.ndarray, labels: list[str] | None = None, loc: tuple[float, float] = (-0.15, 1.05), fontsize: int = 8, fontweight: str = 'bold') -> None

Add panel labels (a, b, c, etc.) to subplots.

PARAMETER	DESCRIPTION
axes	Single axes or array of axes from plt.subplots(). TYPE: Axes | ndarray
labels	Custom labels. If None, uses lowercase letters (a, b, c, ...). TYPE: list[str] | None, optional, by default None DEFAULT: None
loc	Label position in axes coordinates (x, y). Negative x places label to the left of the plot area, y > 1.0 places it above. TYPE: tuple[float, float], optional, by default (-0.15, 1.05) DEFAULT: (-0.15, 1.05)
fontsize	Font size for panel labels. TYPE: int, optional, by default 10 DEFAULT: 8
fontweight	Font weight for panel labels ("normal", "bold", "light", "heavy"). TYPE: str, optional, by default "bold" DEFAULT: 'bold'

Examples:

>>> # Automatic lowercase letters
>>> fig, axes = plt.subplots(1, 3)
>>> label_subplots(axes)

>>> # Custom labels with parentheses
>>> fig, axes = plt.subplots(2, 2)
>>> label_subplots(axes, labels=["(a)", "(b)", "(c)", "(d)"])

>>> # Position labels inside plots
>>> fig, axes = plt.subplots(1, 2)
>>> label_subplots(axes, loc=(0.05, 0.95), fontsize=12)

Source code in spectre/utils/plt.py

def label_subplots(
    axes: plt.Axes | np.ndarray,
    labels: list[str] | None = None,
    loc: tuple[float, float] = (-0.15, 1.05),
    fontsize: int = 8,
    fontweight: str = "bold",
) -> None:
    """
    Add panel labels (a, b, c, etc.) to subplots.

    Parameters
    ----------
    axes : plt.Axes | np.ndarray
        Single axes or array of axes from `plt.subplots()`.

    labels : list[str] | None, optional, by default None
        Custom labels. If None, uses lowercase letters (a, b, c, ...).

    loc : tuple[float, float], optional, by default (-0.15, 1.05)
        Label position in axes coordinates (x, y). Negative x places label to
        the left of the plot area, y > 1.0 places it above.

    fontsize : int, optional, by default 10
        Font size for panel labels.

    fontweight : str, optional, by default "bold"
        Font weight for panel labels ("normal", "bold", "light", "heavy").

    Examples
    --------
    >>> # Automatic lowercase letters
    >>> fig, axes = plt.subplots(1, 3)
    >>> label_subplots(axes)

    >>> # Custom labels with parentheses
    >>> fig, axes = plt.subplots(2, 2)
    >>> label_subplots(axes, labels=["(a)", "(b)", "(c)", "(d)"])

    >>> # Position labels inside plots
    >>> fig, axes = plt.subplots(1, 2)
    >>> label_subplots(axes, loc=(0.05, 0.95), fontsize=12)
    """
    # Convert single axes to array for uniform handling
    if isinstance(axes, plt.Axes):
        axes_flat = [axes]
    else:
        axes_flat = axes.flatten()

    # Generate default labels if not provided
    if labels is None:
        labels = list(string.ascii_lowercase[: len(axes_flat)])

    if len(labels) < len(axes_flat):
        raise ValueError(
            f"Not enough labels provided: {len(labels)} labels for {len(axes_flat)} subplots."
        )

    # Add labels to each subplot
    for ax, label in zip(axes_flat, labels):
        ax.text(
            loc[0],
            loc[1],
            label,
            transform=ax.transAxes,
            fontsize=fontsize,
            fontweight=fontweight,
            verticalalignment="top",
            horizontalalignment="right" if loc[0] < 0 else "left",
        )

plot_kde(energy: np.ndarray, extent: tuple[float, float] | tuple[float, float, float, float] | None = None, error: np.ndarray | None = None, ax: plt.Axes | None = None, plot_type: Literal['imshow', 'contourf', 'contour', 'line'] = 'imshow', cmap: str = 'bwr', levels: int | list[float] | None = None, colorbar: bool = True, colorbar_label: str = 'Free Energy', xlabel: str | None = None, ylabel: str | None = None, title: str | None = None, vmin: float | None = None, vmax: float | None = None, show_error: bool = False, error_alpha: float = 0.3) -> tuple[plt.Figure, plt.Axes]

Plot free energy from KDE calculation.

Supports both 1D and 2D free energy landscapes with multiple visualization options including heatmaps, contours, and line plots.

PARAMETER	DESCRIPTION
energy	Free energy array from KDE.fit(). Shape (n_bins,) for 1D or (n_x_bins, n_y_bins) for 2D. TYPE: ndarray
extent	Axis limits from KDE.extent. For 1D: (x_min, x_max). For 2D: (x_min, x_max, y_min, y_max). If None, uses array indices. TYPE: tuple[float, float] | tuple[float, float, float, float] | None DEFAULT: None
error	Error estimate from KDE.fit() with n_splits > 1. Same shape as energy. TYPE: np.ndarray | None, optional, by default None DEFAULT: None
ax	Axes to plot on. If None, creates new figure with plt.subplots(). TYPE: plt.Axes | None, optional, by default None DEFAULT: None
plot_type	Visualization type. For 1D data, only "line" is valid. For 2D: - "imshow": Heatmap with interpolation - "contourf": Filled contours - "contour": Contour lines only - "line": Line plot (1D only) TYPE: str, by default "imshow" DEFAULT: 'imshow'
cmap	Colormap name for 2D plots. TYPE: str, optional, by default "viridis" DEFAULT: 'bwr'
levels	Contour levels. If int, creates that many evenly spaced levels. If list, uses specified levels. If None, uses 10 levels for contours. TYPE: int | list[float] | None, optional, by default None DEFAULT: None
colorbar	Add colorbar for 2D plots. TYPE: bool, optional, by default True DEFAULT: True
colorbar_label	Colorbar label text. TYPE: str, optional, by default "Free Energy" DEFAULT: 'Free Energy'
xlabel	X-axis label. If None, no label is set. TYPE: str | None, optional, by default None DEFAULT: None
ylabel	Y-axis label. If None, no label is set. TYPE: str | None, optional, by default None DEFAULT: None
title	Plot title. If None, no title is set. TYPE: str | None, optional, by default None DEFAULT: None
vmin	Minimum value for color scale. If None, uses data minimum. TYPE: float | None, optional, by default None DEFAULT: None
vmax	Maximum value for color scale. If None, uses data maximum. TYPE: float | None, optional, by default None DEFAULT: None
show_error	Show error bands for 1D plots or error as transparency for 2D. TYPE: bool, optional, by default False DEFAULT: False
error_alpha	Transparency for error visualization. TYPE: float, optional, by default 0.3 DEFAULT: 0.3

RETURNS	DESCRIPTION
tuple[Figure, Axes]	Figure and axes objects.

Examples:

>>> from spectre.utils.kde import KDE
>>> from spectre.utils.plt import plot_kde, save_figure
>>>
>>> # 1D free energy
>>> kde = KDE()
>>> energy, error = kde.fit(X_1d, n_splits=5)
>>> fig, ax = plot_kde(
...     energy,
...     extent=kde.extent,
...     error=error,
...     plot_type="line",
...     show_error=True,
...     xlabel="Reaction Coordinate",
... )
>>> save_figure(fig, "free_energy_1d")
>>>
>>> # 2D free energy with contours
>>> kde = KDE()
>>> energy, _ = kde.fit(X_2d, bins=(100, 100))
>>> fig, ax = plot_kde(
...     energy,
...     extent=kde.extent,
...     plot_type="contourf",
...     levels=20,
...     xlabel="CV1",
...     ylabel="CV2",
...     title="Free Energy Landscape",
... )
>>>
>>> # Multi-panel comparison
>>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
>>> plot_kde(energy1, kde1.extent, ax=axes[0], plot_type="imshow")
>>> plot_kde(energy2, kde2.extent, ax=axes[1], plot_type="contourf")

Source code in spectre/utils/plt.py

def plot_kde(
    energy: np.ndarray,
    extent: tuple[float, float] | tuple[float, float, float, float] | None = None,
    error: np.ndarray | None = None,
    ax: plt.Axes | None = None,
    plot_type: Literal["imshow", "contourf", "contour", "line"] = "imshow",
    cmap: str = "bwr",
    levels: int | list[float] | None = None,
    colorbar: bool = True,
    colorbar_label: str = "Free Energy",
    xlabel: str | None = None,
    ylabel: str | None = None,
    title: str | None = None,
    vmin: float | None = None,
    vmax: float | None = None,
    show_error: bool = False,
    error_alpha: float = 0.3,
) -> tuple[plt.Figure, plt.Axes]:
    """
    Plot free energy from KDE calculation.

    Supports both 1D and 2D free energy landscapes with multiple visualization
    options including heatmaps, contours, and line plots.

    Parameters
    ----------
    energy : np.ndarray
        Free energy array from `KDE.fit()`. Shape `(n_bins,)` for 1D or
        `(n_x_bins, n_y_bins)` for 2D.

    extent : tuple[float, float] | tuple[float, float, float, float] | None
        Axis limits from `KDE.extent`. For 1D: `(x_min, x_max)`. For 2D:
        `(x_min, x_max, y_min, y_max)`. If None, uses array indices.

    error : np.ndarray | None, optional, by default None
        Error estimate from `KDE.fit()` with `n_splits > 1`. Same shape as
        `energy`.

    ax : plt.Axes | None, optional, by default None
        Axes to plot on. If None, creates new figure with `plt.subplots()`.

    plot_type : str, by default "imshow"
        Visualization type. For 1D data, only "line" is valid. For 2D:
        - "imshow": Heatmap with interpolation
        - "contourf": Filled contours
        - "contour": Contour lines only
        - "line": Line plot (1D only)

    cmap : str, optional, by default "viridis"
        Colormap name for 2D plots.

    levels : int | list[float] | None, optional, by default None
        Contour levels. If int, creates that many evenly spaced levels. If list,
        uses specified levels. If None, uses 10 levels for contours.

    colorbar : bool, optional, by default True
        Add colorbar for 2D plots.

    colorbar_label : str, optional, by default "Free Energy"
        Colorbar label text.

    xlabel : str | None, optional, by default None
        X-axis label. If None, no label is set.

    ylabel : str | None, optional, by default None
        Y-axis label. If None, no label is set.

    title : str | None, optional, by default None
        Plot title. If None, no title is set.

    vmin : float | None, optional, by default None
        Minimum value for color scale. If None, uses data minimum.

    vmax : float | None, optional, by default None
        Maximum value for color scale. If None, uses data maximum.

    show_error : bool, optional, by default False
        Show error bands for 1D plots or error as transparency for 2D.

    error_alpha : float, optional, by default 0.3
        Transparency for error visualization.

    Returns
    -------
    tuple[plt.Figure, plt.Axes]
        Figure and axes objects.

    Examples
    --------
    >>> from spectre.utils.kde import KDE
    >>> from spectre.utils.plt import plot_kde, save_figure
    >>>
    >>> # 1D free energy
    >>> kde = KDE()
    >>> energy, error = kde.fit(X_1d, n_splits=5)
    >>> fig, ax = plot_kde(
    ...     energy,
    ...     extent=kde.extent,
    ...     error=error,
    ...     plot_type="line",
    ...     show_error=True,
    ...     xlabel="Reaction Coordinate",
    ... )
    >>> save_figure(fig, "free_energy_1d")
    >>>
    >>> # 2D free energy with contours
    >>> kde = KDE()
    >>> energy, _ = kde.fit(X_2d, bins=(100, 100))
    >>> fig, ax = plot_kde(
    ...     energy,
    ...     extent=kde.extent,
    ...     plot_type="contourf",
    ...     levels=20,
    ...     xlabel="CV1",
    ...     ylabel="CV2",
    ...     title="Free Energy Landscape",
    ... )
    >>>
    >>> # Multi-panel comparison
    >>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
    >>> plot_kde(energy1, kde1.extent, ax=axes[0], plot_type="imshow")
    >>> plot_kde(energy2, kde2.extent, ax=axes[1], plot_type="contourf")
    """
    # Determine dimensionality
    if energy.ndim == 1:
        is_1d = True
    elif energy.ndim == 2:
        is_1d = False
    else:
        raise ValueError(f"Expected energy to be 1D or 2D array, got {energy.ndim}D.")

    # Validate plot type for dimensionality
    if is_1d and plot_type != "line":
        raise ValueError(f"For 1D data, plot_type must be 'line', got '{plot_type}'.")
    if not is_1d and plot_type == "line":
        raise ValueError(f"For 2D data, plot_type cannot be 'line', got '{plot_type}'.")

    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.get_figure()

    # 1D plotting
    if is_1d:
        if extent is not None and len(extent) == 2:
            x = np.linspace(extent[0], extent[1], len(energy))
        else:
            x = np.arange(len(energy))

        ax.plot(x, energy, linewidth=1.5)

        if error is not None and show_error:
            ax.fill_between(
                x,
                energy - error,
                energy + error,
                alpha=error_alpha,
                label="Standard Error",
            )
            ax.legend(frameon=False)

        if xlabel:
            ax.set_xlabel(xlabel)
        ax.set_ylabel(colorbar_label if ylabel is None else ylabel)

    # 2D plotting
    else:
        # Transpose for correct orientation (imshow convention)
        energy_T = energy.T

        if vmin is None:
            vmin = energy_T.min()
        if vmax is None:
            vmax = energy_T.max()

        if plot_type == "imshow":
            im = ax.imshow(
                energy_T,
                origin="lower",
                extent=extent,
                aspect="auto",
                cmap=cmap,
                vmin=vmin,
                vmax=vmax,
                interpolation="bilinear",
            )

        elif plot_type in ["contourf", "contour"]:
            if extent is not None and len(extent) == 4:
                x = np.linspace(extent[0], extent[1], energy.shape[0])
                y = np.linspace(extent[2], extent[3], energy.shape[1])  # ty: ignore[index-out-of-bounds]
            else:
                x = np.arange(energy.shape[0])
                y = np.arange(energy.shape[1])

            X, Y = np.meshgrid(x, y)

            if levels is None:
                levels = 10

            if plot_type == "contourf":
                im = ax.contourf(
                    X, Y, energy_T, levels=levels, cmap=cmap, vmin=vmin, vmax=vmax
                )
            else:  # contour
                im = ax.contour(
                    X, Y, energy_T, levels=levels, cmap=cmap, vmin=vmin, vmax=vmax
                )
                ax.clabel(im, inline=True, fontsize=8)

        # Add colorbar
        if colorbar:
            cbar = plt.colorbar(im, ax=ax)
            cbar.set_label(colorbar_label)

        # Show error as overlay (optional)
        if error is not None and show_error:
            error_T = error.T
            ax.contour(
                X if plot_type != "imshow" else None,
                Y if plot_type != "imshow" else None,
                error_T,
                levels=3,
                colors="white",
                linewidths=0.5,
                alpha=error_alpha,
                extent=extent if plot_type == "imshow" else None,
            )

        if xlabel:
            ax.set_xlabel(xlabel)
        if ylabel:
            ax.set_ylabel(ylabel)

    if title:
        ax.set_title(title)

    ax.grid(False)  # Disable grid for free energy plots

    return fig, ax

plot_figure(width: float | str = 490.0, fraction: float = 1.0, subplots: tuple[int, int] = (1, 1), aspect_ratio: float | str = 'golden', enable_latex: bool = True, label_panels: bool = False, panel_labels: list[str] | None = None, label_loc: tuple[float, float] = (-0.15, 1.05), constrained_layout: bool = True, spacing: dict | None = None) -> tuple[plt.Figure, plt.Axes | np.ndarray]

Create a publication-ready figure with optimal settings.

PARAMETER	DESCRIPTION
width	Figure width in points or preset ("textwidth", "columnwidth"). TYPE: float | str, optional, by default 490.0 DEFAULT: 490.0
fraction	Fraction of total width. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
subplots	Number of subplots (rows, cols). TYPE: tuple[int, int], optional, by default (1, 1) DEFAULT: (1, 1)
aspect_ratio	Aspect ratio or preset ("golden", "square"). TYPE: float | str, optional, by default "golden" DEFAULT: 'golden'
enable_latex	Enable LaTeX rendering. TYPE: bool, optional, by default True DEFAULT: True
label_panels	Automatically add panel labels (a, b, c, etc.) to subplots. TYPE: bool, optional, by default False DEFAULT: False
panel_labels	Custom panel labels. If None and label_panels=True, uses lowercase letters. TYPE: list[str] | None, optional, by default None DEFAULT: None
label_loc	Panel label position in axes coordinates (x, y). TYPE: tuple[float, float], optional, by default (-0.15, 1.05) DEFAULT: (-0.15, 1.05)
constrained_layout	Use constrained layout for better spacing. If False and spacing is provided, uses subplots_adjust instead. TYPE: bool, optional, by default True DEFAULT: True
spacing	Manual spacing control via fig.subplots_adjust(). Keys: "wspace", "hspace", "left", "right", "top", "bottom". Only used if constrained_layout=False. TYPE: dict | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
tuple[Figure, Axes | ndarray]	Figure and axes objects. For single subplot returns plt.Axes, otherwise returns np.ndarray of axes.

Examples:

>>> # Single plot
>>> fig, ax = plot_figure(width="columnwidth", fraction=0.8)
>>> ax.plot(x, y)
>>> save_figure(fig, "single_plot")

>>> # Multi-panel figure with automatic labeling
>>> fig, axes = plot_figure(
...     width="textwidth",
...     subplots=(1, 3),
...     label_panels=True,
... )
>>> for i, ax in enumerate(axes):
...     ax.plot(x, y[i])
...     setup_axes(ax, xlabel="X", ylabel="Y")
>>> save_figure(fig, "multi_panel")

>>> # Custom spacing and labels
>>> fig, axes = plot_figure(
...     width="columnwidth",
...     subplots=(2, 2),
...     label_panels=True,
...     panel_labels=["(a)", "(b)", "(c)", "(d)"],
...     constrained_layout=False,
...     spacing={"wspace": 0.3, "hspace": 0.4},
... )

Source code in spectre/utils/plt.py

def plot_figure(
    width: float | str = 490.0,
    fraction: float = 1.0,
    subplots: tuple[int, int] = (1, 1),
    aspect_ratio: float | str = "golden",
    enable_latex: bool = True,
    label_panels: bool = False,
    panel_labels: list[str] | None = None,
    label_loc: tuple[float, float] = (-0.15, 1.05),
    constrained_layout: bool = True,
    spacing: dict | None = None,
) -> tuple[plt.Figure, plt.Axes | np.ndarray]:
    """
    Create a publication-ready figure with optimal settings.

    Parameters
    ----------
    width : float | str, optional, by default 490.0
        Figure width in points or preset ("textwidth", "columnwidth").

    fraction : float, optional, by default 1.0
        Fraction of total width.

    subplots : tuple[int, int], optional, by default (1, 1)
        Number of subplots (rows, cols).

    aspect_ratio : float | str, optional, by default "golden"
        Aspect ratio or preset ("golden", "square").

    enable_latex : bool, optional, by default True
        Enable LaTeX rendering.

    label_panels : bool, optional, by default False
        Automatically add panel labels (a, b, c, etc.) to subplots.

    panel_labels : list[str] | None, optional, by default None
        Custom panel labels. If None and `label_panels=True`, uses lowercase
        letters.

    label_loc : tuple[float, float], optional, by default (-0.15, 1.05)
        Panel label position in axes coordinates (x, y).

    constrained_layout : bool, optional, by default True
        Use constrained layout for better spacing. If False and `spacing` is
        provided, uses `subplots_adjust` instead.

    spacing : dict | None, optional, by default None
        Manual spacing control via `fig.subplots_adjust()`. Keys: "wspace",
        "hspace", "left", "right", "top", "bottom". Only used if
        `constrained_layout=False`.

    Returns
    -------
    tuple[plt.Figure, plt.Axes | np.ndarray]
        Figure and axes objects. For single subplot returns plt.Axes, otherwise
        returns np.ndarray of axes.

    Examples
    --------
    >>> # Single plot
    >>> fig, ax = plot_figure(width="columnwidth", fraction=0.8)
    >>> ax.plot(x, y)
    >>> save_figure(fig, "single_plot")

    >>> # Multi-panel figure with automatic labeling
    >>> fig, axes = plot_figure(
    ...     width="textwidth",
    ...     subplots=(1, 3),
    ...     label_panels=True,
    ... )
    >>> for i, ax in enumerate(axes):
    ...     ax.plot(x, y[i])
    ...     setup_axes(ax, xlabel="X", ylabel="Y")
    >>> save_figure(fig, "multi_panel")

    >>> # Custom spacing and labels
    >>> fig, axes = plot_figure(
    ...     width="columnwidth",
    ...     subplots=(2, 2),
    ...     label_panels=True,
    ...     panel_labels=["(a)", "(b)", "(c)", "(d)"],
    ...     constrained_layout=False,
    ...     spacing={"wspace": 0.3, "hspace": 0.4},
    ... )
    """
    if enable_latex:
        enable_mpl_params()

    figsize = set_figsize(width, fraction, subplots, aspect_ratio)

    if subplots == (1, 1):
        fig, ax = plt.subplots(figsize=figsize, constrained_layout=constrained_layout)

        if label_panels:
            label_subplots(ax, labels=panel_labels, loc=label_loc)

        return fig, ax
    else:
        fig, axes = plt.subplots(
            *subplots, figsize=figsize, constrained_layout=constrained_layout
        )

        # Apply manual spacing if requested
        if not constrained_layout and spacing is not None:
            fig.subplots_adjust(**spacing)

        if label_panels:
            label_subplots(axes, labels=panel_labels, loc=label_loc)

        return fig, axes

plot_training_metrics(metrics, metric_names: list[str] | str | None = None, ax: plt.Axes | None = None, figsize: tuple[float, float] | None = None, title: str = 'Training Metrics', smooth: float | None = None, log_scale: bool = False) -> tuple[plt.Figure, plt.Axes]

Plot training metrics from parametric model training.

PARAMETER	DESCRIPTION
metrics	Metrics source. Can be: - MetricsCallback instance from training - Dictionary of metric names and their values over epochs TYPE: MetricsCallback | dict[str, list[float]]
metric_names	Specific metric names to plot. If None, plots all metrics except 'epoch'. Can be a single metric name string or list of metric names. Supports wildcard patterns (e.g., "train_" matches all metrics starting with "train_", "_loss" matches all metrics ending with "_loss"). TYPE: list[str] | str | None, optional, by default None DEFAULT: None
ax	Axes to plot on. If None, creates new figure with plt.subplots(). TYPE: plt.Axes | None, optional, by default None DEFAULT: None
figsize	Figure size in inches. Only used if ax is None. TYPE: tuple[float, float] | None, optional, by default None DEFAULT: None
title	Plot title. TYPE: str, optional, by default "Training Metrics" DEFAULT: 'Training Metrics'
smooth	Smoothing factor for exponential moving average (0 < smooth < 1). If None, no smoothing is applied. TYPE: float | None, optional, by default None DEFAULT: None
log_scale	Whether to use log scale for y-axis. TYPE: bool, optional, by default False DEFAULT: False

RETURNS	DESCRIPTION
tuple[Figure, Axes]	Figure and axes objects.

Examples:

>>> from spectre.utils.callbacks import MetricsCallback
>>> from spectre.utils.plt import plot_training_metrics
>>>
>>> # Create callback for training
>>> metrics_cb = MetricsCallback(track_lr=True, track_epoch_time=True)
>>> trainer = Trainer(callbacks=[metrics_cb])
>>> trainer.fit(model, datamodule)
>>>
>>> # Plot all metrics
>>> fig, ax = plot_training_metrics(metrics_cb, smooth=0.9)
>>>
>>> # Plot specific metrics
>>> fig, ax = plot_training_metrics(
...     metrics_cb,
...     metric_names=["train_loss", "val_loss"],
... )
>>>
>>> # Plot metrics matching a pattern
>>> fig, ax = plot_training_metrics(
...     metrics_cb,
...     metric_names="train_*",  # Matches metrics starting with "train_"
... )
>>>
>>> # Multi-panel comparison
>>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
>>> plot_training_metrics(metrics_cb1, ax=axes[0], title="Model 1")
>>> plot_training_metrics(metrics_cb2, ax=axes[1], title="Model 2")
>>>
>>> # Also works with dict
>>> metrics_dict = {
...     "train_loss": [0.5, 0.3, 0.2],
...     "val_loss": [0.6, 0.4, 0.3],
... }
>>> fig, ax = plot_training_metrics(metrics_dict)

Source code in spectre/utils/plt.py

def plot_training_metrics(
    metrics,
    metric_names: list[str] | str | None = None,
    ax: plt.Axes | None = None,
    figsize: tuple[float, float] | None = None,
    title: str = "Training Metrics",
    smooth: float | None = None,
    log_scale: bool = False,
) -> tuple[plt.Figure, plt.Axes]:
    """
    Plot training metrics from parametric model training.

    Parameters
    ----------
    metrics : MetricsCallback | dict[str, list[float]]
        Metrics source. Can be:
        - MetricsCallback instance from training
        - Dictionary of metric names and their values over epochs

    metric_names : list[str] | str | None, optional, by default None
        Specific metric names to plot. If None, plots all metrics except
        'epoch'. Can be a single metric name string or list of metric names.
        Supports wildcard patterns (e.g., "train_*" matches all metrics
        starting with "train_", "*_loss" matches all metrics ending with
        "_loss").

    ax : plt.Axes | None, optional, by default None
        Axes to plot on. If None, creates new figure with `plt.subplots()`.

    figsize : tuple[float, float] | None, optional, by default None
        Figure size in inches. Only used if `ax` is None.

    title : str, optional, by default "Training Metrics"
        Plot title.

    smooth : float | None, optional, by default None
        Smoothing factor for exponential moving average (0 < smooth < 1).
        If None, no smoothing is applied.

    log_scale : bool, optional, by default False
        Whether to use log scale for y-axis.

    Returns
    -------
    tuple[plt.Figure, plt.Axes]
        Figure and axes objects.

    Examples
    --------
    >>> from spectre.utils.callbacks import MetricsCallback
    >>> from spectre.utils.plt import plot_training_metrics
    >>>
    >>> # Create callback for training
    >>> metrics_cb = MetricsCallback(track_lr=True, track_epoch_time=True)
    >>> trainer = Trainer(callbacks=[metrics_cb])
    >>> trainer.fit(model, datamodule)
    >>>
    >>> # Plot all metrics
    >>> fig, ax = plot_training_metrics(metrics_cb, smooth=0.9)
    >>>
    >>> # Plot specific metrics
    >>> fig, ax = plot_training_metrics(
    ...     metrics_cb,
    ...     metric_names=["train_loss", "val_loss"],
    ... )
    >>>
    >>> # Plot metrics matching a pattern
    >>> fig, ax = plot_training_metrics(
    ...     metrics_cb,
    ...     metric_names="train_*",  # Matches metrics starting with "train_"
    ... )
    >>>
    >>> # Multi-panel comparison
    >>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
    >>> plot_training_metrics(metrics_cb1, ax=axes[0], title="Model 1")
    >>> plot_training_metrics(metrics_cb2, ax=axes[1], title="Model 2")
    >>>
    >>> # Also works with dict
    >>> metrics_dict = {
    ...     "train_loss": [0.5, 0.3, 0.2],
    ...     "val_loss": [0.6, 0.4, 0.3],
    ... }
    >>> fig, ax = plot_training_metrics(metrics_dict)
    """
    # Extract metrics dict from MetricsCallback if needed
    if hasattr(metrics, "metrics"):
        metrics_dict = dict(metrics.metrics)
    else:
        metrics_dict = metrics

    # Filter metrics to plot
    if metric_names is not None:
        if isinstance(metric_names, str):
            metric_names = [metric_names]

        # Collect all matching metrics (supports wildcards)
        matched_keys = set()
        for pattern in metric_names:
            if "*" in pattern or "?" in pattern or "[" in pattern:
                # Pattern matching
                matched_keys.update(fnmatch.filter(metrics_dict.keys(), pattern))
            else:
                # Exact match
                if pattern in metrics_dict:
                    matched_keys.add(pattern)

        plot_metrics = {k: v for k, v in metrics_dict.items() if k in matched_keys}
    else:
        # Plot all except 'epoch'
        plot_metrics = {k: v for k, v in metrics_dict.items() if k != "epoch"}

    if not plot_metrics:
        raise ValueError("No metrics to plot. Check metric names or callback data.")

    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.get_figure()

    # Use epoch data if available, otherwise use indices
    if "epoch" in metrics_dict:
        epochs = metrics_dict["epoch"]
    else:
        epochs = list(range(len(next(iter(plot_metrics.values())))))

    for name, values in plot_metrics.items():
        if smooth is not None:
            # Apply exponential moving average
            smoothed = _exponential_moving_average(values, smooth)
            ax.plot(epochs, smoothed, label=name, linewidth=1.5)
            ax.plot(epochs, values, alpha=0.3, linewidth=0.5)
        else:
            ax.plot(epochs, values, label=name, linewidth=1.5)

    ax.set_xlabel("Epoch")
    ax.set_ylabel("Metric Value")
    ax.set_title(title)
    ax.legend(frameon=False)
    ax.grid(True, alpha=0.3)

    if log_scale:
        ax.set_yscale("log")

    # Only apply tight_layout if we created the figure
    if figsize is not None:
        plt.tight_layout()

    return fig, ax

plot_loss_curves(metrics, train_key: str = 'train_loss', val_key: str = 'val_loss', ax: plt.Axes | None = None, figsize: tuple[float, float] | None = None, title: str = 'Training and Validation Loss', smooth: float | None = 0.9, log_scale: bool = False) -> tuple[plt.Figure, plt.Axes]

Plot training and validation loss curves.

PARAMETER	DESCRIPTION
metrics	Metrics source. Can be: - MetricsCallback instance from training - Dictionary containing loss values TYPE: MetricsCallback | dict[str, list[float]]
train_key	Key name for training loss in metrics dict. TYPE: str, optional, by default "train_loss" DEFAULT: 'train_loss'
val_key	Key name for validation loss in metrics dict. TYPE: str, optional, by default "val_loss" DEFAULT: 'val_loss'
ax	Axes to plot on. If None, creates new figure with plt.subplots(). TYPE: plt.Axes | None, optional, by default None DEFAULT: None
figsize	Figure size in inches. Only used if ax is None. TYPE: tuple[float, float] | None, optional, by default None DEFAULT: None
title	Plot title. TYPE: str, optional, by default "Training and Validation Loss" DEFAULT: 'Training and Validation Loss'
smooth	Smoothing factor for exponential moving average (0 < smooth < 1). If None, no smoothing is applied. TYPE: float | None, optional, by default 0.9 DEFAULT: 0.9
log_scale	Whether to use log scale for y-axis. TYPE: bool, optional, by default False DEFAULT: False

RETURNS	DESCRIPTION
tuple[Figure, Axes]	Figure and axes objects.

Examples:

>>> from spectre.utils.callbacks import MetricsCallback
>>> from spectre.utils.plt import plot_loss_curves
>>>
>>> # After training with MetricsCallback
>>> metrics_cb = MetricsCallback(track_lr=True)
>>> trainer = Trainer(callbacks=[metrics_cb])
>>> trainer.fit(model, datamodule)
>>>
>>> # Plot loss curves
>>> fig, ax = plot_loss_curves(metrics_cb)
>>>
>>> # Multi-panel with different models
>>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
>>> plot_loss_curves(metrics_cb1, ax=axes[0], title="Model 1")
>>> plot_loss_curves(metrics_cb2, ax=axes[1], title="Model 2")
>>>
>>> # Custom metric names
>>> fig, ax = plot_loss_curves(
...     metrics_cb,
...     train_key="loss",
...     val_key="val_loss",
... )

Source code in spectre/utils/plt.py

def plot_loss_curves(
    metrics,
    train_key: str = "train_loss",
    val_key: str = "val_loss",
    ax: plt.Axes | None = None,
    figsize: tuple[float, float] | None = None,
    title: str = "Training and Validation Loss",
    smooth: float | None = 0.9,
    log_scale: bool = False,
) -> tuple[plt.Figure, plt.Axes]:
    """
    Plot training and validation loss curves.

    Parameters
    ----------
    metrics : MetricsCallback | dict[str, list[float]]
        Metrics source. Can be:
        - MetricsCallback instance from training
        - Dictionary containing loss values

    train_key : str, optional, by default "train_loss"
        Key name for training loss in metrics dict.

    val_key : str, optional, by default "val_loss"
        Key name for validation loss in metrics dict.

    ax : plt.Axes | None, optional, by default None
        Axes to plot on. If None, creates new figure with `plt.subplots()`.

    figsize : tuple[float, float] | None, optional, by default None
        Figure size in inches. Only used if `ax` is None.

    title : str, optional, by default "Training and Validation Loss"
        Plot title.

    smooth : float | None, optional, by default 0.9
        Smoothing factor for exponential moving average (0 < smooth < 1).
        If None, no smoothing is applied.

    log_scale : bool, optional, by default False
        Whether to use log scale for y-axis.

    Returns
    -------
    tuple[plt.Figure, plt.Axes]
        Figure and axes objects.

    Examples
    --------
    >>> from spectre.utils.callbacks import MetricsCallback
    >>> from spectre.utils.plt import plot_loss_curves
    >>>
    >>> # After training with MetricsCallback
    >>> metrics_cb = MetricsCallback(track_lr=True)
    >>> trainer = Trainer(callbacks=[metrics_cb])
    >>> trainer.fit(model, datamodule)
    >>>
    >>> # Plot loss curves
    >>> fig, ax = plot_loss_curves(metrics_cb)
    >>>
    >>> # Multi-panel with different models
    >>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
    >>> plot_loss_curves(metrics_cb1, ax=axes[0], title="Model 1")
    >>> plot_loss_curves(metrics_cb2, ax=axes[1], title="Model 2")
    >>>
    >>> # Custom metric names
    >>> fig, ax = plot_loss_curves(
    ...     metrics_cb,
    ...     train_key="loss",
    ...     val_key="val_loss",
    ... )
    """
    # Extract metrics dict from MetricsCallback if needed
    if hasattr(metrics, "metrics"):
        metrics_dict = dict(metrics.metrics)
    else:
        metrics_dict = metrics

    if train_key not in metrics_dict:
        raise ValueError(f"Training loss key '{train_key}' not found in metrics.")

    train_loss = metrics_dict[train_key]
    val_loss = metrics_dict.get(val_key)

    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.get_figure()

    # Use epoch data if available, otherwise use indices
    if "epoch" in metrics_dict:
        epochs = metrics_dict["epoch"]
    else:
        epochs = list(range(len(train_loss)))

    # Plot training loss
    if smooth is not None:
        smoothed_train = _exponential_moving_average(train_loss, smooth)
        ax.plot(epochs, smoothed_train, label="Train Loss (smoothed)", linewidth=1.5)
        ax.plot(epochs, train_loss, alpha=0.3, linewidth=0.5, color="C0")
    else:
        ax.plot(epochs, train_loss, label="Train Loss", linewidth=1.5)

    # Plot validation loss
    if val_loss is not None:
        if smooth is not None:
            smoothed_val = _exponential_moving_average(val_loss, smooth)
            ax.plot(epochs, smoothed_val, label="Val Loss (smoothed)", linewidth=1.5)
            ax.plot(epochs, val_loss, alpha=0.3, linewidth=0.5, color="C1")
        else:
            ax.plot(epochs, val_loss, label="Val Loss", linewidth=1.5)

    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss")
    ax.set_title(title)
    ax.legend(frameon=False)
    ax.grid(True, alpha=0.3)

    if log_scale:
        ax.set_yscale("log")

    # Only apply tight_layout if we created the figure
    if figsize is not None:
        plt.tight_layout()

    return fig, ax

plot_metric_comparison(callbacks_dict: dict[str, Any], metric_name: str, ax: plt.Axes | None = None, figsize: tuple[float, float] | None = None, title: str | None = None, smooth: float | None = None, log_scale: bool = False) -> tuple[plt.Figure, plt.Axes]

Compare a specific metric across multiple model runs or configurations.

PARAMETER	DESCRIPTION
callbacks_dict	Dictionary mapping run names to their metrics sources. Each value can be a MetricsCallback or a metrics dict. Format: {"run1": metrics_cb1, "run2": metrics_cb2} or {"run1": {"train_loss": [...]}, "run2": {...}}. TYPE: dict[str, MetricsCallback | dict[str, list[float]]]
metric_name	Name of the metric to compare (e.g., "train_loss", "val_loss"). TYPE: str
ax	Axes to plot on. If None, creates new figure with plt.subplots(). TYPE: plt.Axes | None, optional, by default None DEFAULT: None
figsize	Figure size in inches. Only used if ax is None. TYPE: tuple[float, float] | None, optional, by default None DEFAULT: None
title	Plot title. If None, uses f"Comparison: {metric_name}". TYPE: str | None, optional, by default None DEFAULT: None
smooth	Smoothing factor for exponential moving average (0 < smooth < 1). TYPE: float | None, optional, by default None DEFAULT: None
log_scale	Whether to use log scale for y-axis. TYPE: bool, optional, by default False DEFAULT: False

RETURNS	DESCRIPTION
tuple[Figure, Axes]	Figure and axes objects.

Examples:

>>> from spectre.utils.callbacks import MetricsCallback
>>> from spectre.utils.plt import plot_metric_comparison
>>>
>>> # Compare different learning rates
>>> metrics_cb1 = MetricsCallback()
>>> metrics_cb2 = MetricsCallback()
>>> # ... train with different configs ...
>>>
>>> runs = {"lr=0.001": metrics_cb1, "lr=0.01": metrics_cb2}
>>> fig, ax = plot_metric_comparison(runs, "train_loss", smooth=0.8)
>>>
>>> # Multi-panel comparison across metrics
>>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
>>> plot_metric_comparison(
...     runs,
...     "train_loss",
...     ax=axes[0],
...     title="Training Loss",
... )
>>> plot_metric_comparison(
...     runs,
...     "val_loss",
...     ax=axes[1],
...     title="Validation Loss",
... )
>>>
>>> # Also works with dict format
>>> runs = {
...     "lr=0.001": {"train_loss": [0.5, 0.3, 0.2]},
...     "lr=0.01": {"train_loss": [0.4, 0.25, 0.15]},
... }
>>> fig, ax = plot_metric_comparison(runs, "train_loss")

Source code in spectre/utils/plt.py

def plot_metric_comparison(
    callbacks_dict: dict[str, Any],
    metric_name: str,
    ax: plt.Axes | None = None,
    figsize: tuple[float, float] | None = None,
    title: str | None = None,
    smooth: float | None = None,
    log_scale: bool = False,
) -> tuple[plt.Figure, plt.Axes]:
    """
    Compare a specific metric across multiple model runs or configurations.

    Parameters
    ----------
    callbacks_dict : dict[str, MetricsCallback | dict[str, list[float]]]
        Dictionary mapping run names to their metrics sources.
        Each value can be a MetricsCallback or a metrics dict.
        Format: `{"run1": metrics_cb1, "run2": metrics_cb2}` or
        `{"run1": {"train_loss": [...]}, "run2": {...}}`.

    metric_name : str
        Name of the metric to compare (e.g., "train_loss", "val_loss").

    ax : plt.Axes | None, optional, by default None
        Axes to plot on. If None, creates new figure with `plt.subplots()`.

    figsize : tuple[float, float] | None, optional, by default None
        Figure size in inches. Only used if `ax` is None.

    title : str | None, optional, by default None
        Plot title. If None, uses `f"Comparison: {metric_name}"`.

    smooth : float | None, optional, by default None
        Smoothing factor for exponential moving average (0 < smooth < 1).

    log_scale : bool, optional, by default False
        Whether to use log scale for y-axis.

    Returns
    -------
    tuple[plt.Figure, plt.Axes]
        Figure and axes objects.

    Examples
    --------
    >>> from spectre.utils.callbacks import MetricsCallback
    >>> from spectre.utils.plt import plot_metric_comparison
    >>>
    >>> # Compare different learning rates
    >>> metrics_cb1 = MetricsCallback()
    >>> metrics_cb2 = MetricsCallback()
    >>> # ... train with different configs ...
    >>>
    >>> runs = {"lr=0.001": metrics_cb1, "lr=0.01": metrics_cb2}
    >>> fig, ax = plot_metric_comparison(runs, "train_loss", smooth=0.8)
    >>>
    >>> # Multi-panel comparison across metrics
    >>> fig, axes = plot_figure(subplots=(1, 2), label_panels=True)
    >>> plot_metric_comparison(
    ...     runs,
    ...     "train_loss",
    ...     ax=axes[0],
    ...     title="Training Loss",
    ... )
    >>> plot_metric_comparison(
    ...     runs,
    ...     "val_loss",
    ...     ax=axes[1],
    ...     title="Validation Loss",
    ... )
    >>>
    >>> # Also works with dict format
    >>> runs = {
    ...     "lr=0.001": {"train_loss": [0.5, 0.3, 0.2]},
    ...     "lr=0.01": {"train_loss": [0.4, 0.25, 0.15]},
    ... }
    >>> fig, ax = plot_metric_comparison(runs, "train_loss")
    """
    # Create figure if needed
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.get_figure()

    for run_name, callback_or_dict in callbacks_dict.items():
        # Extract metrics dict from MetricsCallback if needed
        if hasattr(callback_or_dict, "metrics"):
            metrics = dict(callback_or_dict.metrics)
        else:
            metrics = callback_or_dict

        if metric_name not in metrics:
            continue

        values = metrics[metric_name]

        # Use epoch data if available, otherwise use indices
        if "epoch" in metrics:
            epochs = metrics["epoch"]
        else:
            epochs = list(range(len(values)))

        if smooth is not None:
            smoothed = _exponential_moving_average(values, smooth)
            ax.plot(epochs, smoothed, label=run_name, linewidth=1.5)
            ax.plot(epochs, values, alpha=0.2, linewidth=0.5)
        else:
            ax.plot(epochs, values, label=run_name, linewidth=1.5)

    ax.set_xlabel("Epoch")
    ax.set_ylabel(metric_name.replace("_", " ").title())
    ax.set_title(title or f"Comparison: {metric_name}")
    ax.legend(frameon=False)
    ax.grid(True, alpha=0.3)

    if log_scale:
        ax.set_yscale("log")

    # Only apply tight_layout if we created the figure
    if figsize is not None:
        plt.tight_layout()

    return fig, ax