utils/visualization.py

"""Visualization utilities for physics-informed Bayesian optimization."""

from typing import Callable, Dict, List, Optional, Tuple

import torch
from torch import Tensor
import numpy as np


def plot_convergence(
    campaign_df,
    maximize: bool = True,
    title: str = "Optimization Convergence",
    figsize: Tuple[int, int] = (10, 6),
):
    """Plot the optimization convergence curve.

    Args:
        campaign_df: DataFrame from OptimizationCampaign.to_dataframe().
        maximize: Whether the objective is being maximized.
        title: Plot title.
        figsize: Figure size.

    Returns:
        matplotlib Figure.
    """
    import matplotlib.pyplot as plt

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    objectives = campaign_df["objective"].values

    # Left: all observations
    ax1.plot(range(len(objectives)), objectives, "o-", alpha=0.6, markersize=4)
    ax1.set_xlabel("Experiment Number")
    ax1.set_ylabel("Objective")
    ax1.set_title("All Observations")
    ax1.grid(True, alpha=0.3)

    # Right: best-so-far
    if maximize:
        best_so_far = np.maximum.accumulate(objectives)
    else:
        best_so_far = np.minimum.accumulate(objectives)

    ax2.plot(range(len(best_so_far)), best_so_far, "s-", color="green", markersize=4)
    ax2.set_xlabel("Experiment Number")
    ax2.set_ylabel("Best Objective")
    ax2.set_title("Best So Far")
    ax2.grid(True, alpha=0.3)

    fig.suptitle(title, fontsize=14)
    plt.tight_layout()
    return fig


def plot_surrogate_1d(
    surrogate,
    bounds: Tuple[float, float],
    X_observed: Optional[Tensor] = None,
    y_observed: Optional[Tensor] = None,
    physics_fn: Optional[Callable] = None,
    true_fn: Optional[Callable] = None,
    n_grid: int = 200,
    title: str = "Surrogate Model",
    figsize: Tuple[int, int] = (10, 6),
):
    """Plot a 1D surrogate model with confidence intervals.

    Args:
        surrogate: A SurrogateModel instance.
        bounds: (lower, upper) for the 1D input.
        X_observed: Observed inputs (n, 1).
        y_observed: Observed outputs (n, 1).
        physics_fn: Optional physics model for comparison.
        true_fn: Optional true function for comparison.
        n_grid: Number of grid points.
        title: Plot title.
        figsize: Figure size.

    Returns:
        matplotlib Figure.
    """
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=figsize)

    X_grid = torch.linspace(bounds[0], bounds[1], n_grid).unsqueeze(-1).to(torch.float64)
    mean, var = surrogate.predict(X_grid)
    std = var.sqrt()

    x_np = X_grid.squeeze().numpy()
    mean_np = mean.squeeze().detach().numpy()
    std_np = std.squeeze().detach().numpy()

    # Surrogate prediction
    ax.plot(x_np, mean_np, "b-", label="Surrogate Mean", linewidth=2)
    ax.fill_between(
        x_np,
        mean_np - 2 * std_np,
        mean_np + 2 * std_np,
        alpha=0.2,
        color="blue",
        label="95% CI",
    )

    # Physics model
    if physics_fn is not None:
        with torch.no_grad():
            physics_pred = physics_fn(X_grid).squeeze().numpy()
        ax.plot(x_np, physics_pred, "r--", label="Physics Model", linewidth=1.5)

    # True function
    if true_fn is not None:
        with torch.no_grad():
            true_pred = true_fn(X_grid).squeeze().numpy()
        ax.plot(x_np, true_pred, "k-", label="True Function", linewidth=1.5, alpha=0.7)

    # Observations
    if X_observed is not None and y_observed is not None:
        ax.scatter(
            X_observed.squeeze().numpy(),
            y_observed.squeeze().numpy(),
            c="red",
            s=50,
            zorder=5,
            label="Observations",
            edgecolors="black",
        )

    ax.set_xlabel("Input")
    ax.set_ylabel("Output")
    ax.set_title(title)
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    return fig


def plot_surrogate_2d(
    surrogate,
    bounds: Tensor,
    param_names: Tuple[str, str] = ("x1", "x2"),
    X_observed: Optional[Tensor] = None,
    n_grid: int = 50,
    title: str = "Surrogate Model (2D)",
    figsize: Tuple[int, int] = (12, 5),
):
    """Plot 2D surrogate model as contour plots (mean and uncertainty).

    Args:
        surrogate: A SurrogateModel instance.
        bounds: Tensor of shape (2, 2) with [lower, upper] bounds.
        param_names: Names of the two parameters.
        X_observed: Observed inputs (n, 2).
        n_grid: Grid resolution per dimension.
        title: Plot title.
        figsize: Figure size.

    Returns:
        matplotlib Figure.
    """
    import matplotlib.pyplot as plt

    x1 = torch.linspace(float(bounds[0, 0]), float(bounds[1, 0]), n_grid)
    x2 = torch.linspace(float(bounds[0, 1]), float(bounds[1, 1]), n_grid)
    X1, X2 = torch.meshgrid(x1, x2, indexing="ij")
    X_grid = torch.stack([X1.flatten(), X2.flatten()], dim=-1).to(torch.float64)

    mean, var = surrogate.predict(X_grid)
    mean_2d = mean.squeeze().reshape(n_grid, n_grid).detach().numpy()
    std_2d = var.sqrt().squeeze().reshape(n_grid, n_grid).detach().numpy()

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    # Mean
    c1 = ax1.contourf(
        X1.numpy(), X2.numpy(), mean_2d, levels=20, cmap="viridis"
    )
    plt.colorbar(c1, ax=ax1)
    ax1.set_xlabel(param_names[0])
    ax1.set_ylabel(param_names[1])
    ax1.set_title("Predicted Mean")

    # Uncertainty
    c2 = ax2.contourf(
        X1.numpy(), X2.numpy(), std_2d, levels=20, cmap="plasma"
    )
    plt.colorbar(c2, ax=ax2)
    ax2.set_xlabel(param_names[0])
    ax2.set_ylabel(param_names[1])
    ax2.set_title("Predicted Std Dev")

    # Overlay observations
    if X_observed is not None:
        for ax in [ax1, ax2]:
            ax.scatter(
                X_observed[:, 0].numpy(),
                X_observed[:, 1].numpy(),
                c="red",
                s=30,
                edgecolors="white",
                zorder=5,
            )

    fig.suptitle(title, fontsize=14)
    plt.tight_layout()
    return fig