visualizeGD.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.widgets import TextBox, Button, Slider, RadioButtons

FUNCTIONS = {
    'Himmelblau (Non-Convex)': {
        'func': lambda w: (w[0]**2 + w[1] - 11)**2 + (w[0] + w[1]**2 - 7)**2,
        'bounds': (-6, 6),
        'default_lr': '0.01, 0.045, 0.002',
        'z_log_levels': True,
        'log_loss': True
    },
    'Convex Bowl (Easy)': {
        'func': lambda w: w[0]**2 + w[1]**2,
        'bounds': (-5, 5),
        'default_lr': '0.1, 0.8, 1.05',
        'z_log_levels': False,
        'log_loss': False
    },
    'Rosenbrock (Hard Valley)': {
        'func': lambda w: (1 - w[0])**2 + 100 * (w[1] - w[0]**2)**2,
        'bounds': (-2, 2),
        'default_lr': '0.001, 0.003',
        'z_log_levels': True,
        'log_loss': True
    }
}

def numerical_gradient(f, w, h=1e-5):
    """Vectorized central difference quotient."""
    N = len(w)
    perturbation = np.eye(N) * h 
    w_plus = w[:, None] + perturbation
    w_minus = w[:, None] - perturbation
    return (f(w_plus) - f(w_minus)) / (2 * h)

def gradient_descent(f, start_w, lr, num_steps=100):
    """Runs GD, returning the path and the loss at each step."""
    w = np.array(start_w, dtype=float)
    path = [w.copy()]
    losses = [f(w)]
    
    for _ in range(num_steps):
        grad = numerical_gradient(f, w)
        w = w - lr * grad
        loss = f(w)
        
        # Stop early if exploding gradient (divergence)
        if loss > 1e6 or np.any(np.isnan(w)): 
            break
            
        path.append(w.copy())
        losses.append(loss)
        
    return np.array(path), np.array(losses)

class InteractiveGD:
    def __init__(self):
        self.fig = plt.figure(figsize=(20, 9))
        self.fig.canvas.manager.set_window_title('Interactive Gradient Descent')
        self.fig.suptitle("Interactive Gradient Descent Visualizer", fontsize=18, fontweight='bold')
        
        self.ax1 = self.fig.add_axes([0.03, 0.35, 0.28, 0.55], projection='3d')
        self.ax2 = self.fig.add_axes([0.35, 0.35, 0.28, 0.55])
        self.ax3 = self.fig.add_axes([0.68, 0.35, 0.28, 0.55])
        
        self.current_func_name = list(FUNCTIONS.keys())[0]
        self.runs = []
        self.colors =['#00FFFF', '#FF00FF', '#FFA500', '#FF0000']
        self.drawn_path_objects =[] 
        self.z_offset = 10
        
        self.init_ui()
        self.draw_background()
        
        # Event hooks
        self.fig.canvas.mpl_connect('button_press_event', self.onclick)
        plt.show()

    def init_ui(self):
        """Sets up the UI widgets exactly where we want them."""
        # 3D Instructions
        self.fig.text(0.01, 0.93, " 3D MOUSE CONTROLS:\n . Left-Click + Drag: Rotate\n . Right-Click + Drag: Zoom\n . Scroll Wheel: Zoom", 
                      fontsize=10, bbox=dict(facecolor='white', alpha=0.8, edgecolor='gray'))

        # Function Selector
        self.fig.text(0.03, 0.25, "1. Choose Objective Function", fontsize=11, fontweight='bold')
        ax_radio = self.fig.add_axes([0.03, 0.05, 0.18, 0.18])
        self.radio = RadioButtons(ax_radio, list(FUNCTIONS.keys()))
        self.radio.on_clicked(self.change_function)

        # LR Text Box
        self.fig.text(0.25, 0.15, "2. Set Learning Rates (comma separated)", fontsize=11, fontweight='bold')
        ax_lr = self.fig.add_axes([0.25, 0.08, 0.2, 0.05])
        self.lr_box = TextBox(ax_lr, '', initial=FUNCTIONS[self.current_func_name]['default_lr'])
        
        # Step Animation Slider
        self.fig.text(0.50, 0.15, "4. Scrub to Step Through Time", fontsize=11, fontweight='bold')
        ax_slider = self.fig.add_axes([0.50, 0.08, 0.3, 0.05])
        self.slider = Slider(ax_slider, 'Step', 0, 100, valinit=100, valstep=1)
        self.slider.on_changed(self.update_paths)

        # Clear Button
        ax_clear = self.fig.add_axes([0.85, 0.08, 0.1, 0.05])
        self.btn_clear = Button(ax_clear, 'Clear Paths')
        self.btn_clear.on_clicked(self.clear_paths)

    def change_function(self, label):
        """Triggered when user selects a new math function."""
        self.current_func_name = label
        self.runs =[]
        self.lr_box.set_val(FUNCTIONS[label]['default_lr'])
        self.draw_background()
        self.update_paths()

    def draw_background(self):
        """Calculates and draws the static math surfaces."""
        self.ax1.clear()
        self.ax2.clear()
        self.ax3.clear()
        
        func_info = FUNCTIONS[self.current_func_name]
        f = func_info['func']
        bounds = func_info['bounds']
        
        # Make grid
        x_vals = np.linspace(bounds[0], bounds[1], 60)
        y_vals = np.linspace(bounds[0], bounds[1], 60)
        X, Y = np.meshgrid(x_vals, y_vals)
        Z = f(np.array([X, Y]))
        self.z_offset = Z.max() * 0.02 # dynamic offset to prevent line clipping
        
        # 3D Surface
        self.ax1.plot_surface(X, Y, Z, cmap='viridis', alpha=0.6, edgecolor='none')
        self.ax1.set_title(f"3D Loss Surface: {self.current_func_name}", fontsize=14)
        self.ax1.set_xlabel('Weight 1')
        self.ax1.set_ylabel('Weight 2')
        self.ax1.view_init(elev=40, azim=-45)

        # 2D Contour
        if func_info['z_log_levels']:
            levels = np.logspace(0, np.log10(Z.max() + 1), 30)
        else:
            levels = np.linspace(0, Z.max(), 30)
            
        self.ax2.contour(X, Y, Z, levels=levels, cmap='viridis', alpha=0.8)
        self.ax2.set_title("3. Click ANYWHERE Here To Start!", fontsize=14, color='darkred', fontweight='bold')
        self.ax2.set_xlabel('Weight 1')
        self.ax2.set_ylabel('Weight 2')
        
        # Convergence Graph Base
        self.ax3.set_title("Convergence (Loss vs Iteration)", fontsize=14)
        self.ax3.set_xlabel("Iterations")
        self.ax3.set_ylabel("Loss")
        if func_info['log_loss']:
            self.ax3.set_yscale('log')
        else:
            self.ax3.set_yscale('linear')
        self.ax3.grid(True, alpha=0.3)
        
        self.fig.canvas.draw_idle()

    def clear_paths(self, event=None):
        self.runs =[]
        self.update_paths()

    def onclick(self, event):
        """Calculates new Gradient Descent paths when contour map is clicked."""
        if event.inaxes != self.ax2: return
        
        try:
            lrs =[float(x.strip()) for x in self.lr_box.text.split(',')]
        except ValueError:
            print("Invalid learning rate format. Use commas.")
            return

        start_w = [event.xdata, event.ydata]
        f = FUNCTIONS[self.current_func_name]['func']
        
        for lr in lrs:
            path, losses = gradient_descent(f, start_w, lr=lr)
            self.runs.append({'path': path, 'losses': losses, 'lr': lr, 'start': start_w})
            
        # Keep only the latest 4 lines memory
        self.runs = self.runs[-4:]
        
        # Reset slider to 100 to show the new paths instantly
        if self.slider.val != 100:
            self.slider.set_val(100) 
        else:
            self.update_paths() 

    def update_paths(self, val=None):
        """Redraws the animated lines according to the Step Slider."""
        # Erase old lines
        for obj in self.drawn_path_objects:
            try: obj.remove()
            except: pass
        self.drawn_path_objects.clear()
        
        k = int(self.slider.val) # Current Step
        f = FUNCTIONS[self.current_func_name]['func']
        
        # Lock graph boundaries to prevent jumping while animating
        if self.runs:
            max_len = max(len(r['losses']) for r in self.runs)
            max_loss = max(max(r['losses']) for r in self.runs)
            min_loss = min(min(r['losses']) for r in self.runs)
            self.ax3.set_xlim(0, max_len)
            if FUNCTIONS[self.current_func_name]['log_loss']:
                self.ax3.set_ylim(max(min_loss, 1e-5), max_loss * 1.5)
            else:
                self.ax3.set_ylim(min_loss * 0.9, max_loss * 1.1)
        
        # Draw sliced paths up to step k
        for i, run in enumerate(self.runs):
            path, losses = run['path'], run['losses']
            color = self.colors[i % len(self.colors)]
            lbl = f"LR: {run['lr']} (Start: {run['start'][0]:.1f}, {run['start'][1]:.1f})"
            
            # Slice array up to current frame for animation
            idx = min(k + 1, len(path)) 
            plot_path = path[:idx]
            plot_losses = losses[:idx]
            
            # 3D
            z_path = f(plot_path.T) + self.z_offset
            l3d, = self.ax1.plot(plot_path[:, 0], plot_path[:, 1], z_path, color=color, 
                                 marker='.', markersize=4, linewidth=2, label=lbl)
            # 2D
            l2d, = self.ax2.plot(plot_path[:, 0], plot_path[:, 1], color=color, 
                                 marker='.', markersize=4, linewidth=2)
            start_pt, = self.ax2.plot(path[0, 0], path[0, 1], color=color, 
                                      marker='X', markersize=10, markeredgecolor='black')
            # Convergence
            lconv, = self.ax3.plot(range(idx), plot_losses, color=color, linewidth=2, label=lbl)
            
            self.drawn_path_objects.extend([l3d, l2d, start_pt, lconv])
            
        # 4. Redraw legends
        if self.runs:
            leg1 = self.ax1.legend(loc='upper right', fontsize=8)
            leg3 = self.ax3.legend(loc='upper right', fontsize=8)
            self.drawn_path_objects.extend([leg1, leg3])
            
        self.fig.canvas.draw_idle()

if __name__ == '__main__':
    InteractiveGD()