Update utils/visualization.py

utils/visualization.py

@@ -1,15 +1,18 @@
 import plotly.graph_objects as go
 import plotly.express as px
+import plotly.figure_factory as ff
+from plotly.subplots import make_subplots
 import networkx as nx
 import torch
 import numpy as np
+import pandas as pd
 
 class GraphVisualizer:
-    """Graph visualization utilities"""
+    """Advanced graph visualization utilities"""
     
     @staticmethod
-    def create_graph_plot(data, max_nodes=500):
-        """Create interactive graph visualization"""
+    def create_graph_plot(data, max_nodes=500, layout_algorithm='spring'):
+        """Create interactive graph visualization with multiple layout options"""
         try:
             # Limit nodes for performance
             num_nodes = min(data.num_nodes, max_nodes)
@@ -29,30 +32,71 @@ class GraphVisualizer:
             # Add isolated nodes
             G.add_nodes_from(range(num_nodes))
             
-            # Layout
-            if len(G.nodes()) > 100:
-                pos = nx.spring_layout(G, k=0.5, iterations=20)
+            # Choose layout algorithm
+            if layout_algorithm == 'spring':
+                if len(G.nodes()) > 100:
+                    pos = nx.spring_layout(G, k=0.5, iterations=20)
+                else:
+                    pos = nx.spring_layout(G, k=1, iterations=50)
+            elif layout_algorithm == 'circular':
+                pos = nx.circular_layout(G)
+            elif layout_algorithm == 'kamada_kawai':
+                try:
+                    pos = nx.kamada_kawai_layout(G)
+                except:
+                    pos = nx.spring_layout(G)
+            elif layout_algorithm == 'spectral':
+                try:
+                    pos = nx.spectral_layout(G)
+                except:
+                    pos = nx.spring_layout(G)
             else:
-                pos = nx.spring_layout(G, k=1, iterations=50)
+                pos = nx.spring_layout(G)
             
-            # Node colors
+            # Node colors and sizes
             if hasattr(data, 'y') and data.y is not None:
                 node_colors = data.y.cpu().numpy()[:num_nodes]
+                unique_labels = np.unique(node_colors)
+                color_map = px.colors.qualitative.Set3[:len(unique_labels)]
             else:
                 node_colors = [0] * num_nodes
+                color_map = ['lightblue']
+            
+            # Node sizes based on degree
+            node_sizes = []
+            for node in G.nodes():
+                degree = G.degree(node)
+                node_sizes.append(max(5, min(20, 5 + degree)))
             
             # Create edge traces
             edge_x, edge_y = [], []
+            edge_info = []
+            
             for edge in G.edges():
                 if edge[0] in pos and edge[1] in pos:
                     x0, y0 = pos[edge[0]]
                     x1, y1 = pos[edge[1]]
                     edge_x.extend([x0, x1, None])
                     edge_y.extend([y0, y1, None])
+                    edge_info.append(f"Edge: {edge[0]} - {edge[1]}")
             
             # Create node traces
-            node_x = [pos[node][0] for node in G.nodes() if node in pos]
-            node_y = [pos[node][1] for node in G.nodes() if node in pos]
+            node_x = []
+            node_y = []
+            node_text = []
+            node_info = []
+            
+            for node in G.nodes():
+                if node in pos:
+                    x, y = pos[node]
+                    node_x.append(x)
+                    node_y.append(y)
+                    
+                    # Node info
+                    degree = G.degree(node)
+                    label = node_colors[node] if node < len(node_colors) else 0
+                    node_text.append(f"Node {node}")
+                    node_info.append(f"Node: {node}<br>Degree: {degree}<br>Label: {label}")
             
             fig = go.Figure()
             
@@ -60,10 +104,11 @@ class GraphVisualizer:
             if edge_x:
                 fig.add_trace(go.Scatter(
                     x=edge_x, y=edge_y,
-                    line=dict(width=0.5, color='#888'),
+                    line=dict(width=0.8, color='rgba(125,125,125,0.5)'),
                     hoverinfo='none',
                     mode='lines',
-                    name='Edges'
+                    name='Edges',
+                    showlegend=False
                 ))
             
             # Add nodes
@@ -71,24 +116,43 @@ class GraphVisualizer:
                 x=node_x, y=node_y,
                 mode='markers',
                 hoverinfo='text',
-                text=[f'Node {i}' for i in range(len(node_x))],
+                hovertext=node_info,
+                text=node_text,
                 marker=dict(
-                    size=8,
+                    size=node_sizes,
                     color=node_colors[:len(node_x)],
                     colorscale='Viridis',
-                    line=dict(width=1)
+                    line=dict(width=2, color='white'),
+                    opacity=0.8
                 ),
-                name='Nodes'
+                name='Nodes',
+                showlegend=False
             ))
             
             fig.update_layout(
-                title=f'Graph Visualization ({num_nodes} nodes)',
+                title=dict(
+                    text=f'Graph Visualization ({num_nodes} nodes, {len(edge_list)} edges)',
+                    x=0.5,
+                    font=dict(size=16)
+                ),
                 showlegend=False,
                 hovermode='closest',
                 margin=dict(b=20, l=5, r=5, t=40),
+                annotations=[
+                    dict(
+                        text=f"Layout: {layout_algorithm.title()}",
+                        showarrow=False,
+                        xref="paper", yref="paper",
+                        x=0.005, y=-0.002,
+                        xanchor='left', yanchor='bottom',
+                        font=dict(color="gray", size=10)
+                    )
+                ],
                 xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                 yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-                plot_bgcolor='white'
+                plot_bgcolor='white',
+                width=800,
+                height=600
             )
             
             return fig
@@ -100,45 +164,96 @@ class GraphVisualizer:
                 text=f"Visualization error: {str(e)}",
                 x=0.5, y=0.5,
                 xref="paper", yref="paper",
-                showarrow=False
+                showarrow=False,
+                font=dict(size=14, color="red")
+            )
+            fig.update_layout(
+                title="Graph Visualization Error",
+                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                plot_bgcolor='white'
             )
             return fig
     
     @staticmethod
     def create_metrics_plot(metrics):
-        """Create metrics visualization"""
+        """Create comprehensive metrics visualization"""
         try:
+            # Filter numeric metrics
             metric_names = []
             metric_values = []
             
             for key, value in metrics.items():
-                if isinstance(value, (int, float)) and key != 'error':
-                    metric_names.append(key.replace('_', ' ').title())
-                    metric_values.append(value)
+                if isinstance(value, (int, float)) and key not in ['error', 'loss']:
+                    if not (np.isnan(value) or np.isinf(value)):
+                        metric_names.append(key.replace('_', ' ').title())
+                        metric_values.append(value)
             
             if metric_names:
-                fig = go.Figure(data=[
+                # Create subplots
+                fig = make_subplots(
+                    rows=1, cols=2,
+                    subplot_titles=('Performance Metrics', 'Metric Comparison'),
+                    specs=[[{"type": "bar"}, {"type": "scatter"}]]
+                )
+                
+                # Bar chart
+                colors = px.colors.qualitative.Set3[:len(metric_names)]
+                
+                fig.add_trace(
                     go.Bar(
                         x=metric_names,
                         y=metric_values,
-                        marker_color='lightblue'
-                    )
-                ])
+                        marker_color=colors,
+                        text=[f'{v:.3f}' for v in metric_values],
+                        textposition='auto',
+                        name='Metrics'
+                    ),
+                    row=1, col=1
+                )
+                
+                # Radar chart data
+                fig.add_trace(
+                    go.Scatterpolar(
+                        r=metric_values,
+                        theta=metric_names,
+                        fill='toself',
+                        name='Performance',
+                        line=dict(color='blue')
+                    ),
+                    row=1, col=2
+                )
                 
                 fig.update_layout(
-                    title='Model Performance Metrics',
-                    xaxis_title='Metric',
-                    yaxis_title='Value',
-                    yaxis=dict(range=[0, 1])
+                    title=dict(
+                        text='Model Performance Dashboard',
+                        x=0.5,
+                        font=dict(size=18)
+                    ),
+                    showlegend=False,
+                    height=400
                 )
+                
+                # Update bar chart
+                fig.update_xaxes(title_text="Metrics", row=1, col=1)
+                fig.update_yaxes(title_text="Score", range=[0, 1], row=1, col=1)
+                
+                # Update polar chart
+                fig.update_polars(
+                    radialaxis=dict(range=[0, 1], showticklabels=True),
+                    row=1, col=2
+                )
+                
             else:
                 fig = go.Figure()
                 fig.add_annotation(
-                    text="No metrics to display",
+                    text="No valid metrics to display",
                     x=0.5, y=0.5,
                     xref="paper", yref="paper",
-                    showarrow=False
+                    showarrow=False,
+                    font=dict(size=14)
                 )
+                fig.update_layout(title="Metrics Dashboard")
                 
             return fig
             
@@ -148,6 +263,232 @@ class GraphVisualizer:
                 text=f"Metrics plot error: {str(e)}",
                 x=0.5, y=0.5,
                 xref="paper", yref="paper",
-                showarrow=False
+                showarrow=False,
+                font=dict(size=14, color="red")
+            )
+            fig.update_layout(title="Metrics Error")
+            return fig
+    
+    @staticmethod
+    def create_training_history_plot(history):
+        """Create training history visualization"""
+        try:
+            epochs = list(range(len(history['train_loss'])))
+            
+            # Create subplots
+            fig = make_subplots(
+                rows=2, cols=2,
+                subplot_titles=('Training Loss', 'Training Accuracy', 'Learning Rate', 'Loss Comparison'),
+                specs=[[{"secondary_y": False}, {"secondary_y": False}],
+                       [{"secondary_y": False}, {"secondary_y": False}]]
+            )
+            
+            # Training loss
+            fig.add_trace(
+                go.Scatter(
+                    x=epochs, y=history['train_loss'],
+                    mode='lines', name='Train Loss',
+                    line=dict(color='blue', width=2)
+                ),
+                row=1, col=1
+            )
+            
+            if 'val_loss' in history:
+                fig.add_trace(
+                    go.Scatter(
+                        x=epochs, y=history['val_loss'],
+                        mode='lines', name='Val Loss',
+                        line=dict(color='red', width=2)
+                    ),
+                    row=1, col=1
+                )
+            
+            # Training accuracy
+            fig.add_trace(
+                go.Scatter(
+                    x=epochs, y=history['train_acc'],
+                    mode='lines', name='Train Acc',
+                    line=dict(color='green', width=2)
+                ),
+                row=1, col=2
+            )
+            
+            if 'val_acc' in history:
+                fig.add_trace(
+                    go.Scatter(
+                        x=epochs, y=history['val_acc'],
+                        mode='lines', name='Val Acc',
+                        line=dict(color='orange', width=2)
+                    ),
+                    row=1, col=2
+                )
+            
+            # Learning rate
+            if 'lr' in history:
+                fig.add_trace(
+                    go.Scatter(
+                        x=epochs, y=history['lr'],
+                        mode='lines', name='Learning Rate',
+                        line=dict(color='purple', width=2)
+                    ),
+                    row=2, col=1
+                )
+            
+            # Loss comparison
+            if 'train_loss' in history and 'val_loss' in history:
+                fig.add_trace(
+                    go.Scatter(
+                        x=history['train_loss'], y=history['val_loss'],
+                        mode='markers', name='Train vs Val Loss',
+                        marker=dict(color=epochs, colorscale='Viridis', size=8),
+                        text=[f'Epoch {i}' for i in epochs],
+                        hovertemplate='Train Loss: %{x:.4f}<br>Val Loss: %{y:.4f}<br>%{text}'
+                    ),
+                    row=2, col=2
+                )
+                
+                # Add diagonal line
+                min_loss = min(min(history['train_loss']), min(history['val_loss']))
+                max_loss = max(max(history['train_loss']), max(history['val_loss']))
+                fig.add_trace(
+                    go.Scatter(
+                        x=[min_loss, max_loss], y=[min_loss, max_loss],
+                        mode='lines', name='Perfect Fit',
+                        line=dict(color='gray', dash='dash'),
+                        showlegend=False
+                    ),
+                    row=2, col=2
+                )
+            
+            fig.update_layout(
+                title=dict(
+                    text='Training History Dashboard',
+                    x=0.5,
+                    font=dict(size=18)
+                ),
+                height=600,
+                showlegend=True
+            )
+            
+            # Update axes
+            fig.update_xaxes(title_text="Epoch", row=1, col=1)
+            fig.update_xaxes(title_text="Epoch", row=1, col=2)
+            fig.update_xaxes(title_text="Epoch", row=2, col=1)
+            fig.update_xaxes(title_text="Train Loss", row=2, col=2)
+            
+            fig.update_yaxes(title_text="Loss", row=1, col=1)
+            fig.update_yaxes(title_text="Accuracy", row=1, col=2)
+            fig.update_yaxes(title_text="Learning Rate", type="log", row=2, col=1)
+            fig.update_yaxes(title_text="Val Loss", row=2, col=2)
+            
+            return fig
+            
+        except Exception as e:
+            fig = go.Figure()
+            fig.add_annotation(
+                text=f"Training history plot error: {str(e)}",
+                x=0.5, y=0.5,
+                xref="paper", yref="paper",
+                showarrow=False,
+                font=dict(size=14, color="red")
+            )
+            return fig
+    
+    @staticmethod
+    def create_dataset_stats_plot(dataset_info):
+        """Create dataset statistics visualization"""
+        try:
+            # Prepare data
+            stats_data = []
+            for key, value in dataset_info.items():
+                if isinstance(value, (int, float)) and not np.isnan(value):
+                    stats_data.append({
+                        'Metric': key.replace('_', ' ').title(),
+                        'Value': value
+                    })
+            
+            if not stats_data:
+                raise ValueError("No valid statistics to display")
+            
+            df = pd.DataFrame(stats_data)
+            
+            # Create subplots
+            fig = make_subplots(
+                rows=1, cols=2,
+                subplot_titles=('Dataset Overview', 'Graph Size Distribution'),
+                specs=[[{"type": "bar"}, {"type": "box"}]]
+            )
+            
+            # Bar chart of statistics
+            fig.add_trace(
+                go.Bar(
+                    x=df['Metric'],
+                    y=df['Value'],
+                    marker_color=px.colors.qualitative.Pastel1,
+                    text=df['Value'],
+                    texttemplate='%{text:,.0f}',
+                    textposition='auto'
+                ),
+                row=1, col=1
+            )
+            
+            # Box plot for size distribution (if multiple graphs)
+            if dataset_info.get('num_graphs', 1) > 1:
+                # Simulate distribution based on min/max/avg
+                avg_nodes = dataset_info.get('avg_nodes', 100)
+                min_nodes = dataset_info.get('min_nodes', avg_nodes * 0.5)
+                max_nodes = dataset_info.get('max_nodes', avg_nodes * 1.5)
+                
+                # Generate synthetic distribution
+                np.random.seed(42)
+                node_dist = np.random.normal(avg_nodes, (max_nodes - min_nodes) / 4, 100)
+                node_dist = np.clip(node_dist, min_nodes, max_nodes)
+                
+                fig.add_trace(
+                    go.Box(
+                        y=node_dist,
+                        name='Node Count',
+                        marker_color='lightblue'
+                    ),
+                    row=1, col=2
+                )
+            else:
+                # Single graph - show as point
+                fig.add_trace(
+                    go.Scatter(
+                        x=['Nodes'],
+                        y=[dataset_info.get('avg_nodes', 0)],
+                        mode='markers',
+                        marker=dict(size=20, color='blue'),
+                        name='Node Count'
+                    ),
+                    row=1, col=2
+                )
+            
+            fig.update_layout(
+                title=dict(
+                    text='Dataset Statistics Dashboard',
+                    x=0.5,
+                    font=dict(size=16)
+                ),
+                height=400,
+                showlegend=False
+            )
+            
+            # Update axes
+            fig.update_xaxes(title_text="Metrics", tickangle=45, row=1, col=1)
+            fig.update_yaxes(title_text="Count", row=1, col=1)
+            fig.update_yaxes(title_text="Number of Nodes", row=1, col=2)
+            
+            return fig
+            
+        except Exception as e:
+            fig = go.Figure()
+            fig.add_annotation(
+                text=f"Dataset stats error: {str(e)}",
+                x=0.5, y=0.5,
+                xref="paper", yref="paper",
+                showarrow=False,
+                font=dict(size=14, color="red")
             )
             return fig