Create visualization.py

utils/visualization.py

@@ -0,0 +1,153 @@
+import plotly.graph_objects as go
+import plotly.express as px
+import networkx as nx
+import torch
+import numpy as np
+
+class GraphVisualizer:
+    """Graph visualization utilities"""
+    
+    @staticmethod
+    def create_graph_plot(data, max_nodes=500):
+        """Create interactive graph visualization"""
+        try:
+            # Limit nodes for performance
+            num_nodes = min(data.num_nodes, max_nodes)
+            
+            # Create NetworkX graph
+            G = nx.Graph()
+            edge_list = data.edge_index.t().cpu().numpy()
+            
+            # Filter edges to include only first max_nodes
+            edge_list = edge_list[
+                (edge_list[:, 0] < num_nodes) & (edge_list[:, 1] < num_nodes)
+            ]
+            
+            if len(edge_list) > 0:
+                G.add_edges_from(edge_list)
+            
+            # 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)
+            else:
+                pos = nx.spring_layout(G, k=1, iterations=50)
+            
+            # Node colors
+            if hasattr(data, 'y') and data.y is not None:
+                node_colors = data.y.cpu().numpy()[:num_nodes]
+            else:
+                node_colors = [0] * num_nodes
+            
+            # Create edge traces
+            edge_x, edge_y = [], []
+            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])
+            
+            # 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]
+            
+            fig = go.Figure()
+            
+            # Add edges
+            if edge_x:
+                fig.add_trace(go.Scatter(
+                    x=edge_x, y=edge_y,
+                    line=dict(width=0.5, color='#888'),
+                    hoverinfo='none',
+                    mode='lines',
+                    name='Edges'
+                ))
+            
+            # Add nodes
+            fig.add_trace(go.Scatter(
+                x=node_x, y=node_y,
+                mode='markers',
+                hoverinfo='text',
+                text=[f'Node {i}' for i in range(len(node_x))],
+                marker=dict(
+                    size=8,
+                    color=node_colors[:len(node_x)],
+                    colorscale='Viridis',
+                    line=dict(width=1)
+                ),
+                name='Nodes'
+            ))
+            
+            fig.update_layout(
+                title=f'Graph Visualization ({num_nodes} nodes)',
+                showlegend=False,
+                hovermode='closest',
+                margin=dict(b=20, l=5, r=5, t=40),
+                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                plot_bgcolor='white'
+            )
+            
+            return fig
+            
+        except Exception as e:
+            # Return error plot
+            fig = go.Figure()
+            fig.add_annotation(
+                text=f"Visualization error: {str(e)}",
+                x=0.5, y=0.5,
+                xref="paper", yref="paper",
+                showarrow=False
+            )
+            return fig
+    
+    @staticmethod
+    def create_metrics_plot(metrics):
+        """Create metrics visualization"""
+        try:
+            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 metric_names:
+                fig = go.Figure(data=[
+                    go.Bar(
+                        x=metric_names,
+                        y=metric_values,
+                        marker_color='lightblue'
+                    )
+                ])
+                
+                fig.update_layout(
+                    title='Model Performance Metrics',
+                    xaxis_title='Metric',
+                    yaxis_title='Value',
+                    yaxis=dict(range=[0, 1])
+                )
+            else:
+                fig = go.Figure()
+                fig.add_annotation(
+                    text="No metrics to display",
+                    x=0.5, y=0.5,
+                    xref="paper", yref="paper",
+                    showarrow=False
+                )
+                
+            return fig
+            
+        except Exception as e:
+            fig = go.Figure()
+            fig.add_annotation(
+                text=f"Metrics plot error: {str(e)}",
+                x=0.5, y=0.5,
+                xref="paper", yref="paper",
+                showarrow=False
+            )
+            return fig