Upload utils/visualization.py with huggingface_hub

utils/visualization.py

@@ -0,0 +1,208 @@
+"""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