vibe_ml_lab/optimization_lab.py

"""Optimizer-race and loss-breakdown demos."""

from __future__ import annotations

from functools import lru_cache

import numpy as np
from matplotlib import pyplot as plt

from .utils import figure_to_image

LOSS_LABELS = ["airplane", "car", "cat", "ship"]


def _objective(point: np.ndarray) -> float:
    x_value, y_value = float(point[0]), float(point[1])
    return (
        0.18 * (x_value**2)
        + 0.75 * (y_value**2)
        + 0.35 * np.sin(1.35 * x_value) * np.cos(0.85 * y_value)
        + 0.08 * x_value * y_value
    )


def _gradient(point: np.ndarray) -> np.ndarray:
    x_value, y_value = float(point[0]), float(point[1])
    grad_x = 0.36 * x_value + 0.4725 * np.cos(1.35 * x_value) * np.cos(0.85 * y_value) + 0.08 * y_value
    grad_y = 1.5 * y_value - 0.2975 * np.sin(1.35 * x_value) * np.sin(0.85 * y_value) + 0.08 * x_value
    return np.array([grad_x, grad_y], dtype=np.float32)


def _run_optimizer(
    method: str,
    learning_rate: float,
    beta: float,
    steps: int,
    start: np.ndarray,
    noise_bank: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    path = [start.astype(np.float32)]
    losses = [_objective(start)]
    velocity = np.zeros(2, dtype=np.float32)
    point = start.astype(np.float32).copy()

    for step_id in range(steps):
        grad = _gradient(point) + noise_bank[step_id]
        if method == "momentum":
            velocity = beta * velocity - learning_rate * grad
            point = point + velocity
        else:
            point = point - learning_rate * grad
        path.append(point.copy())
        losses.append(_objective(point))

    return np.stack(path, axis=0), np.asarray(losses, dtype=np.float32)


def _plot_optimizer_paths(sgd_path: np.ndarray, momentum_path: np.ndarray) -> np.ndarray:
    xs = np.linspace(-3.2, 3.2, 240)
    ys = np.linspace(-3.2, 3.2, 240)
    grid_x, grid_y = np.meshgrid(xs, ys)
    values = (
        0.18 * np.square(grid_x)
        + 0.75 * np.square(grid_y)
        + 0.35 * np.sin(1.35 * grid_x) * np.cos(0.85 * grid_y)
        + 0.08 * grid_x * grid_y
    )

    fig, ax = plt.subplots(figsize=(6.0, 5.4))
    contour = ax.contourf(grid_x, grid_y, values, levels=35, cmap="YlOrRd")
    ax.contour(grid_x, grid_y, values, levels=18, colors="white", linewidths=0.45, alpha=0.55)
    ax.plot(sgd_path[:, 0], sgd_path[:, 1], marker="o", markersize=3.2, linewidth=2.0, color="#355070", label="SGD")
    ax.plot(momentum_path[:, 0], momentum_path[:, 1], marker="o", markersize=3.2, linewidth=2.0, color="#2a9d8f", label="Momentum")
    ax.scatter([sgd_path[0, 0]], [sgd_path[0, 1]], color="#ff7f7f", s=55, zorder=5, label="Start")
    ax.set_title("Optimizer paths on the same loss surface", fontsize=14, fontweight="bold")
    ax.set_xlabel("w1")
    ax.set_ylabel("w2")
    ax.legend(frameon=False, loc="upper right")
    fig.colorbar(contour, ax=ax, shrink=0.82, label="Loss value")
    fig.tight_layout()
    return figure_to_image(fig)


def _plot_loss_curves(sgd_losses: np.ndarray, momentum_losses: np.ndarray) -> np.ndarray:
    fig, ax = plt.subplots(figsize=(6.0, 4.0))
    steps = np.arange(len(sgd_losses))
    ax.plot(steps, sgd_losses, color="#355070", linewidth=2.4, label="SGD")
    ax.plot(steps, momentum_losses, color="#2a9d8f", linewidth=2.4, label="Momentum")
    ax.set_title("Loss vs step", fontsize=14, fontweight="bold")
    ax.set_xlabel("Step")
    ax.set_ylabel("Loss")
    ax.grid(alpha=0.24)
    ax.legend(frameon=False)
    fig.tight_layout()
    return figure_to_image(fig)


@lru_cache(maxsize=40)
def run_optimizer_demo(
    learning_rate: float,
    beta: float,
    steps: int,
    noise_scale: float,
    start_x: float,
    start_y: float,
) -> dict[str, object]:
    """Compare SGD and momentum on the same toy objective."""
    rng = np.random.default_rng(23)
    start = np.array([start_x, start_y], dtype=np.float32)
    noise_bank = rng.normal(0.0, noise_scale, size=(steps, 2)).astype(np.float32)

    sgd_path, sgd_losses = _run_optimizer("sgd", learning_rate, beta, steps, start, noise_bank)
    momentum_path, momentum_losses = _run_optimizer("momentum", learning_rate, beta, steps, start, noise_bank)

    sgd_improvement = float(sgd_losses[0] - sgd_losses[-1])
    momentum_improvement = float(momentum_losses[0] - momentum_losses[-1])
    summary = (
        "### Optimizer Race Snapshot\n"
        f"- Start point: `({start_x:.2f}, {start_y:.2f})`\n"
        f"- Learning rate `{learning_rate:.3f}`, momentum beta `{beta:.2f}`, gradient noise `{noise_scale:.2f}`\n"
        f"- SGD final loss: **{sgd_losses[-1]:.4f}** (improvement `{sgd_improvement:.4f}`)\n"
        f"- Momentum final loss: **{momentum_losses[-1]:.4f}** (improvement `{momentum_improvement:.4f}`)"
    )

    return {
        "summary": summary,
        "path_plot": _plot_optimizer_paths(sgd_path, momentum_path),
        "curve_plot": _plot_loss_curves(sgd_losses, momentum_losses),
    }


def _plot_loss_breakdown(scores: np.ndarray, probabilities: np.ndarray, target_index: int) -> np.ndarray:
    fig, axes = plt.subplots(1, 2, figsize=(8.2, 3.8))
    score_colors = ["#ff7f7f" if index == target_index else "#9bf6ff" for index in range(len(scores))]
    prob_colors = ["#2a9d8f" if index == target_index else "#ffd166" for index in range(len(scores))]

    axes[0].bar(LOSS_LABELS, scores, color=score_colors, edgecolor="#355070")
    axes[0].set_title("Raw class scores", fontsize=13, fontweight="bold")
    axes[0].grid(axis="y", alpha=0.22)

    axes[1].bar(LOSS_LABELS, probabilities, color=prob_colors, edgecolor="#355070")
    axes[1].set_title("Softmax probabilities", fontsize=13, fontweight="bold")
    axes[1].set_ylim(0.0, 1.0)
    axes[1].grid(axis="y", alpha=0.22)

    fig.tight_layout()
    return figure_to_image(fig)


def run_loss_demo(
    score_airplane: float,
    score_car: float,
    score_cat: float,
    score_ship: float,
    target_name: str,
) -> dict[str, object]:
    """Show a step-by-step comparison of three different losses."""
    scores = np.array([score_airplane, score_car, score_cat, score_ship], dtype=np.float32)
    target_index = LOSS_LABELS.index(target_name)
    target = np.zeros(len(LOSS_LABELS), dtype=np.float32)
    target[target_index] = 1.0

    correct_score = float(scores[target_index])
    margins = np.maximum(0.0, scores - correct_score + 1.0)
    margins[target_index] = 0.0
    svm_loss = float(np.sum(margins))

    shifted = scores - scores.max()
    exp_values = np.exp(shifted)
    probabilities = exp_values / exp_values.sum()
    cross_entropy = float(-np.log(probabilities[target_index] + 1e-8))
    squared_errors = np.square(probabilities - target)
    mse_loss = float(np.mean(squared_errors))
    prediction = LOSS_LABELS[int(np.argmax(scores))]

    summary = (
        "### Loss Breakdown Snapshot\n"
        f"- Predicted class from raw scores: **{prediction}**\n"
        f"- Target class: **{target_name}**\n"
        f"- Multiclass SVM loss: `{svm_loss:.4f}` from `sum(max(0, s_j - s_y + 1))`\n"
        f"- Softmax cross-entropy: `{cross_entropy:.4f}` from `-log(p_target)`\n"
        f"- Probability MSE baseline: `{mse_loss:.4f}` from `mean((p - one_hot)^2)`"
    )

    table = []
    for index, class_name in enumerate(LOSS_LABELS):
        table.append(
            [
                class_name,
                round(float(scores[index]), 4),
                round(float(margins[index]), 4),
                round(float(probabilities[index]), 4),
                round(float(squared_errors[index]), 4),
            ]
        )

    return {
        "summary": summary,
        "plot": _plot_loss_breakdown(scores, probabilities, target_index),
        "table": table,
    }


def run_optimizer_demo_ui(
    learning_rate: float,
    beta: float,
    steps: int,
    noise_scale: float,
    start_x: float,
    start_y: float,
) -> tuple[str, np.ndarray, np.ndarray]:
    """UI adapter for the optimizer demo."""
    result = run_optimizer_demo(
        learning_rate=float(learning_rate),
        beta=float(beta),
        steps=int(steps),
        noise_scale=float(noise_scale),
        start_x=float(start_x),
        start_y=float(start_y),
    )
    return result["summary"], result["path_plot"], result["curve_plot"]


def run_loss_demo_ui(
    score_airplane: float,
    score_car: float,
    score_cat: float,
    score_ship: float,
    target_name: str,
) -> tuple[str, np.ndarray, list[list[object]]]:
    """UI adapter for the loss demo."""
    result = run_loss_demo(
        score_airplane=float(score_airplane),
        score_car=float(score_car),
        score_cat=float(score_cat),
        score_ship=float(score_ship),
        target_name=str(target_name),
    )
    return result["summary"], result["plot"], result["table"]