visualizer.py

import numpy as np
import pythreejs as p3
from IPython.display import display
import matplotlib.pyplot as plt
from matplotlib import cm
import ipywidgets as widgets
from neural_network import SimpleNeuralNetwork, create_sample_network


class NeuralNetworkVisualizer:
    def __init__(self, network):
        self.network = network
        self.scene = None
        self.renderer = None
        self.camera = None
        self.controls = None
        
    def create_neuron_geometry(self, position, neuron_type='hidden', activation=0.5):
        """Create a sphere geometry for a neuron"""
        # Color based on neuron type
        if neuron_type == 'input':
            color = 0x4CAF50  # Green
        elif neuron_type == 'output':
            color = 0xF44336  # Red
        else:
            # Color based on activation level
            intensity = max(0.2, min(1.0, activation))
            color_value = int(255 * intensity)
            color = (color_value << 16) | (color_value << 8) | 255  # Blue gradient
        
        # Create sphere geometry
        geometry = p3.SphereGeometry(radius=0.2, widthSegments=8, heightSegments=6)
        material = p3.MeshLambertMaterial(color=color, transparent=True, opacity=0.8)
        sphere = p3.Mesh(geometry=geometry, material=material)
        sphere.position = position
        
        return sphere
    
    def create_connection_geometry(self, start_pos, end_pos, weight, max_weight=1.0):
        """Create a line geometry for connections between neurons"""
        # Normalize weight for visualization
        normalized_weight = abs(weight) / max_weight if max_weight > 0 else 0.1
        
        # Color based on weight strength
        if weight > 0:
            color = 0xFF6B6B  # Red for positive weights
        else:
            color = 0x4ECDC4  # Cyan for negative weights
        
        # Line opacity based on weight strength
        opacity = max(0.1, min(1.0, normalized_weight))
        
        # Create line geometry
        geometry = p3.BufferGeometry(
            attributes={
                'position': p3.BufferAttribute(
                    np.array([start_pos, end_pos], dtype=np.float32).flatten(),
                    normalized=False
                )
            }
        )
        
        material = p3.LineBasicMaterial(
            color=color,
            transparent=True,
            opacity=opacity,
            linewidth=max(1, int(normalized_weight * 5))
        )
        
        line = p3.Line(geometry=geometry, material=material)
        return line
    
    def create_layer_label(self, position, text, layer_type):
        """Create text label for layers"""
        # This is simplified - in a full implementation you might use TextGeometry
        # For now, we'll create a simple colored box as a placeholder
        if layer_type == 'input':
            color = 0x4CAF50
        elif layer_type == 'output':
            color = 0xF44336
        else:
            color = 0x2196F3
        
        geometry = p3.BoxGeometry(0.5, 0.1, 0.1)
        material = p3.MeshLambertMaterial(color=color)
        label = p3.Mesh(geometry=geometry, material=material)
        label.position = [position[0], position[1] - 2, position[2]]
        
        return label
    
    def create_visualization(self):
        """Create the complete 3D visualization of the neural network"""
        # Get network structure
        layer_positions = self.network.get_layer_positions()
        connections = self.network.get_connections()
        
        # Create scene
        self.scene = p3.Scene()
        
        # Add lighting
        ambient_light = p3.AmbientLight(color=0x404040, intensity=0.4)
        directional_light = p3.DirectionalLight(color=0xffffff, intensity=0.6)
        directional_light.position = [10, 10, 10]
        self.scene.add([ambient_light, directional_light])
        
        # Create neurons
        all_objects = []
        
        for layer_idx, layer_data in enumerate(layer_positions):
            layer_type = layer_data['type']
            positions = layer_data['positions']
            
            # Create layer label
            center_pos = np.mean(positions, axis=0)
            label = self.create_layer_label(center_pos, f"{layer_type.title()}", layer_type)
            all_objects.append(label)
            
            # Create neurons for this layer
            for neuron_idx, pos in enumerate(positions):
                # Simulate some activation for visualization
                activation = np.random.random() * 0.8 + 0.2
                neuron = self.create_neuron_geometry(pos, layer_type, activation)
                all_objects.append(neuron)
        
        # Create connections
        max_weight = 0
        all_connections = []
        for layer_connections in connections:
            for conn in layer_connections:
                max_weight = max(max_weight, abs(conn['weight']))
                all_connections.append(conn)
        
        for conn in all_connections:
            line = self.create_connection_geometry(
                conn['start'], conn['end'], conn['weight'], max_weight
            )
            all_objects.append(line)
        
        # Add all objects to scene
        self.scene.add(all_objects)
        
        # Create camera
        self.camera = p3.PerspectiveCamera(
            fov=75,
            aspect=1.0,
            near=0.1,
            far=1000,
            position=[10, 5, 10]
        )
        
        # Create controls
        self.controls = p3.OrbitControls(controlling=self.camera)
        
        # Create renderer
        self.renderer = p3.Renderer(
            scene=self.scene,
            camera=self.camera,
            controls=[self.controls],
            width=800,
            height=600,
            background='white'
        )
        
        return self.renderer
    
    def create_network_info_widget(self):
        """Create an information widget showing network details"""
        layer_info = self.network.layer_info
        
        info_text = "Neural Network Architecture:\n"
        info_text += "=" * 30 + "\n"
        
        for i, (layer_type, size, layer_obj) in enumerate(layer_info):
            info_text += f"Layer {i}: {layer_type.title()} ({size} neurons)\n"
            if layer_obj is not None:
                params = sum(p.numel() for p in layer_obj.parameters())
                info_text += f"  Parameters: {params}\n"
        
        total_params = sum(p.numel() for p in self.network.parameters())
        info_text += f"\nTotal Parameters: {total_params}"
        
        return widgets.Textarea(
            value=info_text,
            description='Network Info:',
            layout=widgets.Layout(width='400px', height='200px')
        )
    
    def show(self):
        """Display the complete visualization with controls"""
        # Create the 3D visualization
        renderer = self.create_visualization()
        
        # Create info widget
        info_widget = self.create_network_info_widget()
        
        # Create control buttons
        reset_button = widgets.Button(description="Reset View")
        randomize_button = widgets.Button(description="Randomize Activations")
        
        def reset_view(b):
            self.camera.position = [10, 5, 10]
            self.controls.reset()
        
        def randomize_activations(b):
            # This would update neuron colors - simplified for this demo
            print("Randomizing neuron activations...")
        
        reset_button.on_click(reset_view)
        randomize_button.on_click(randomize_activations)
        
        # Layout widgets
        controls_box = widgets.HBox([reset_button, randomize_button])
        main_display = widgets.VBox([
            widgets.HTML("<h2>Neural Network 3D Visualization</h2>"),
            renderer,
            controls_box,
            info_widget
        ])
        
        return main_display


def create_visualization_demo():
    """Create a demonstration of the neural network visualization"""
    # Create sample network
    network = create_sample_network()
    
    # Create visualizer
    visualizer = NeuralNetworkVisualizer(network)
    
    # Return the visualization widget
    return visualizer.show()