"""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"]