kmeans.py

import numpy as np
import matplotlib.pyplot as plt

class KMeansVisualizer:
    def __init__(self):
        # Initial data
        self.initial_data = np.array([
            [20, 2],  # A
            [22, 1],  # B
            [78, 11], # C
            [85, 9],  # D
            [80, 10], # E
            [18, 3]   # F
        ])
        self.initial_names = ['A', 'B', 'C', 'D', 'E', 'F']
        self.client_data = self.initial_data.copy()
        self.client_names = self.initial_names.copy()
        self.colors = ['#0088FE', '#FF8042', '#00C49F', '#FFBB28']
        self.k = 2
        self.centroids = None
        self.labels = None
        self.iteration = 0
        self.converged = False
        self.selecting_centroids = False
        self.selected_centroids = []
        self.adding_points = False
        self.xlim = (0, 100)
        self.ylim = (0, 12)
        self.step_history = []  # Store states for back/forward navigation
        self.current_step = -1  # Current position in history
        
    def init_centroids(self):
        # Randomly choose k points as initial centroids
        indices = np.random.choice(len(self.client_data), self.k, replace=False)
        self.centroids = self.client_data[indices]
        
    def compute_distances(self):
        distances = np.sqrt(((self.client_data - self.centroids[:, np.newaxis])**2).sum(axis=2))
        self.labels = np.argmin(distances, axis=0)
        
    def update_centroids(self):
        old_centroids = self.centroids.copy()
        for i in range(self.k):
            points = self.client_data[self.labels == i]
            if len(points) > 0:
                self.centroids[i] = points.mean(axis=0)
        return np.allclose(old_centroids, self.centroids)

    def calculate_ssd(self):
        if self.centroids is None or self.labels is None:
            return float('inf') # Or some other appropriate value if not clustered yet

        ssd = 0
        for i in range(self.k):
            points_in_cluster = self.client_data[self.labels == i]
            if len(points_in_cluster) > 0:
                distance = np.sum((points_in_cluster - self.centroids[i])**2)
                ssd += distance
        return ssd

    def plot(self, fig, ax):
        ax.clear()
        ax.grid(True)  # Always show grid regardless of data state
        
        if len(self.client_data) == 0:
            ax.set_xlabel('Purchases')
            ax.set_ylabel('Visits')
            ax.set_title('K-means Clustering (No Data)')
            ax.set_xlim(0, 100)
            ax.set_ylim(0, 12)
            self.xlim = ax.get_xlim()
            self.ylim = ax.get_ylim()
            return
        
        # Calculate axis limits with padding
        x_min, x_max = self.client_data[:, 0].min(), self.client_data[:, 0].max()
        y_min, y_max = self.client_data[:, 1].min(), self.client_data[:, 1].max()
        
        # Add 10% padding to the limits
        x_padding = (x_max - x_min) * 0.1
        y_padding = (y_max - y_min) * 0.1
        
        # Set axis limits with padding
        ax.set_xlim(x_min - x_padding, x_max + x_padding)
        ax.set_ylim(y_min - y_padding, y_max + y_padding)
        self.xlim = ax.get_xlim()
        self.ylim = ax.get_ylim()
        
        # Always plot all points
        if self.centroids is None:
            ax.scatter(self.client_data[:, 0], self.client_data[:, 1], 
                      c='gray', label='Unassigned', picker=True)
            for i, name in enumerate(self.client_names):
                ax.annotate(name, (self.client_data[i, 0], self.client_data[i, 1]))
        else:
            # Plot clustered points
            for i in range(self.k):
                mask = self.labels == i
                ax.scatter(self.client_data[mask, 0], self.client_data[mask, 1], 
                          c=self.colors[i], label=f'Cluster {i+1}')
                
            # Plot centroids
            ax.scatter(self.centroids[:, 0], self.centroids[:, 1], 
                      c=self.colors[:self.k], marker='*', s=200, label='Centroids')
            
            # Add client labels
            for i, name in enumerate(self.client_names):
                ax.annotate(name, (self.client_data[i, 0], self.client_data[i, 1]))
        
        # Show selected centroids during manual selection
        if self.selecting_centroids and self.selected_centroids:
            points = np.array(self.selected_centroids)
            ax.scatter(points[:, 0], points[:, 1], 
                      c='red', marker='*', s=200, label='Selected Centroids')

        ax.set_xlabel('Purchases')
        ax.set_ylabel('Visits')
        ax.set_title(f'K-means Clustering (Iteration {self.iteration})')
        ax.legend()
        
    def handle_click(self, event):
        if event.inaxes != self.ax:
            return
        
        if self.adding_points:
            # Add new client point
            new_point = np.array([[event.xdata, event.ydata]])
            self.client_data = np.append(self.client_data, new_point, axis=0)
            self.client_names.append(f'{len(self.client_names) + 1}')
            self.plot(self.fig, self.ax)
            plt.draw()
        
        elif self.selecting_centroids:
            # Add centroid point
            self.selected_centroids.append([event.xdata, event.ydata])
            if len(self.selected_centroids) == self.k:
                self.centroids = np.array(self.selected_centroids)
                self.selecting_centroids = False
                self.selected_centroids = []
                self.compute_distances()
                self.plot(self.fig, self.ax)
            plt.draw()

    def clear_data(self):
        # Reset to initial data
        self.client_data = self.initial_data.copy()
        self.client_names = self.initial_names.copy()
        self.centroids = None
        self.labels = None
        self.iteration = 0
        self.converged = False
        self.selected_centroids = []
        self.selecting_centroids = False
        self.adding_points = False
        self.step_history = []
        self.current_step = -1

    def initialize_custom_clients(self, points, names):
        self.client_data = points
        self.client_names = names
        self.centroids = None
        self.labels = None
        self.iteration = 0
        self.converged = False

    def save_state(self):
        """Save current state to history"""
        state = {
            'centroids': self.centroids.copy() if self.centroids is not None else None,
            'labels': self.labels.copy() if self.labels is not None else None,
            'iteration': self.iteration,
            'converged': self.converged
        }
        # Remove future states if we're not at the end
        self.step_history = self.step_history[:self.current_step + 1]
        self.step_history.append(state)
        self.current_step = len(self.step_history) - 1

    def load_state(self, step_index):
        """Load state from history"""
        if 0 <= step_index < len(self.step_history):
            state = self.step_history[step_index]
            self.centroids = state['centroids'].copy() if state['centroids'] is not None else None
            self.labels = state['labels'].copy() if state['labels'] is not None else None
            self.iteration = state['iteration']
            self.converged = state['converged']
            self.current_step = step_index
            return True
        return False

    def step(self):
        if not self.converged and len(self.client_data) > 0:
            if self.iteration == 0:
                # Check if we have enough data points for k clusters
                if len(self.client_data) < self.k:
                    print(f"Need at least {self.k} points to create {self.k} clusters")
                    return
                self.init_centroids()
            
            self.compute_distances()
            self.converged = self.update_centroids()
            self.iteration += 1
            
            if self.iteration > 10:
                self.converged = True
            
            self.save_state()  # Save state after each step
    
    def run(self):
        self.fig, self.ax = plt.subplots(figsize=(10, 6))
        
        # Initialize the plot with unassigned points
        self.plot(self.fig, self.ax)
        plt.draw()
        
        # Setup keyboard event handling
        def on_key(event):
            if event.key == ' ':  # Space bar
                self.step()
            elif event.key == 'q':
                plt.close()
            elif event.key == 'a':  # Add points mode
                self.adding_points = not self.adding_points
                self.selecting_centroids = False
                print("Adding points mode:", "ON" if self.adding_points else "OFF")
            elif event.key == 's':  # Select centroids mode
                self.selecting_centroids = not self.selecting_centroids
                self.adding_points = False
                self.selected_centroids = []
                print("Select centroids mode:", "ON" if self.selecting_centroids else "OFF")
            elif event.key == 'c':  # Clear all data
                self.clear_data()
            elif event.key == 'left':  # Backward step
                if self.current_step > 0:
                    self.load_state(self.current_step - 1)
                    print("Stepped back to iteration:", self.iteration)
            elif event.key == 'right':  # Forward step
                if self.current_step < len(self.step_history) - 1:
                    self.load_state(self.current_step + 1)
                    print("Stepped forward to iteration:", self.iteration)
        
        self.fig.canvas.mpl_connect('key_press_event', on_key)
        self.fig.canvas.mpl_connect('button_press_event', self.handle_click)
        
        print("Controls:")
        print("- Press SPACE to perform next iteration")
        print("- Press A to toggle adding points mode")
        print("- Press S to toggle centroid selection mode")
        print("- Press C to clear all data")
        print("- Press Q to quit")
        print("- Press LEFT/RIGHT arrow keys to navigate steps")
        print("\nClick on the plot to add points or select centroids when in respective modes")
        
        plt.show()

    def add_start_point(self, x, y):
        """Handle manual selection of starting centroids"""
        self.selected_centroids.append([x, y])
        if len(self.selected_centroids) == self.k:
            self.centroids = np.array(self.selected_centroids)
            self.selected_centroids = []
            self.compute_distances()
            return True
        return False

    def initialize_clients(self):
        """Reset to initial state"""
        self.clear_data()

if __name__ == "__main__":
    visualizer = KMeansVisualizer()
    visualizer.run()