Update core/graph_sequencer.py

core/graph_sequencer.py

@@ -3,28 +3,39 @@ import numpy as np
 import networkx as nx
 from scipy.sparse.linalg import eigsh
 from sklearn.cluster import SpectralClustering
-from torch_geometric.utils import to_networkx, get_laplacian
-import torch_geometric.utils as pyg_utils
+import warnings
+warnings.filterwarnings('ignore')
 
 class GraphSequencer:
     """
     Production-ready graph ordering strategies
-    All methods use real graph data - no hardcoded values
+    Device-safe implementation with performance optimizations
     """
     
     @staticmethod
     def bfs_ordering(edge_index, num_nodes, start_node=None):
-        """Breadth-first search ordering"""
-        # Convert to NetworkX for BFS
-        G = nx.Graph()
-        G.add_nodes_from(range(num_nodes))
+        """Breadth-first search ordering - optimized version"""
+        device = edge_index.device
+        
+        if num_nodes <= 1:
+            return torch.arange(num_nodes, device=device)
+        
+        # Convert to adjacency list efficiently
+        adj_list = [[] for _ in range(num_nodes)]
         edge_list = edge_index.t().cpu().numpy()
-        G.add_edges_from(edge_list)
+        
+        for src, dst in edge_list:
+            if src < num_nodes and dst < num_nodes:
+                adj_list[src].append(dst)
+                adj_list[dst].append(src)
+        
+        # Remove duplicates and sort for determinism
+        adj_list = [sorted(list(set(neighbors))) for neighbors in adj_list]
         
         # Start from highest degree node if not specified
         if start_node is None:
-            degrees = dict(G.degree())
-            start_node = max(degrees, key=degrees.get)
+            degrees = [len(neighbors) for neighbors in adj_list]
+            start_node = np.argmax(degrees) if degrees else 0
         
         # BFS traversal
         visited = set()
@@ -33,15 +44,15 @@ class GraphSequencer:
         
         while queue:
             node = queue.pop(0)
-            if node in visited:
+            if node in visited or node >= num_nodes:
                 continue
                 
             visited.add(node)
             order.append(node)
             
             # Add neighbors by degree (deterministic)
-            neighbors = list(G.neighbors(node))
-            neighbors.sort(key=lambda n: G.degree(n), reverse=True)
+            neighbors = adj_list[node]
+            neighbors.sort(key=lambda n: (len(adj_list[n]), n), reverse=True)
             
             for neighbor in neighbors:
                 if neighbor not in visited:
@@ -52,35 +63,64 @@ class GraphSequencer:
             if node not in visited:
                 order.append(node)
         
-        return torch.tensor(order, dtype=torch.long)
+        return torch.tensor(order, dtype=torch.long, device=device)
     
     @staticmethod
     def spectral_ordering(edge_index, num_nodes):
-        """Spectral ordering using graph Laplacian eigenvector"""
+        """Spectral ordering using graph Laplacian eigenvector - robust version"""
+        device = edge_index.device
+        
+        if num_nodes <= 2:
+            return torch.arange(num_nodes, device=device)
+        
         try:
-            # Compute normalized Laplacian
-            edge_index_np = edge_index.cpu().numpy()
+            # Move to CPU for scipy operations
+            edge_index_cpu = edge_index.cpu().numpy()
             
             # Create adjacency matrix
             A = np.zeros((num_nodes, num_nodes))
-            A[edge_index_np[0], edge_index_np[1]] = 1
-            A[edge_index_np[1], edge_index_np[0]] = 1  # Undirected
+            valid_edges = (edge_index_cpu[0] < num_nodes) & (edge_index_cpu[1] < num_nodes)
+            valid_edge_index = edge_index_cpu[:, valid_edges]
+            
+            A[valid_edge_index[0], valid_edge_index[1]] = 1
+            A[valid_edge_index[1], valid_edge_index[0]] = 1  # Undirected
             
             # Degree matrix
-            D = np.diag(np.sum(A, axis=1))
+            degrees = np.sum(A, axis=1)
+            
+            # Handle disconnected components
+            if np.any(degrees == 0):
+                # Add self-loops to isolated nodes
+                isolated = degrees == 0
+                A[isolated, isolated] = 1
+                degrees = np.sum(A, axis=1)
+            
+            D = np.diag(degrees)
             
             # Normalized Laplacian: L = D^(-1/2) * (D - A) * D^(-1/2)
-            D_sqrt_inv = np.diag(1.0 / np.sqrt(np.maximum(np.diag(D), 1e-12)))
+            degrees_sqrt_inv = np.where(degrees > 0, 1.0 / np.sqrt(degrees), 0)
+            D_sqrt_inv = np.diag(degrees_sqrt_inv)
             L = D_sqrt_inv @ (D - A) @ D_sqrt_inv
             
-            # Compute second smallest eigenvector (Fiedler vector)
-            eigenvals, eigenvecs = eigsh(L, k=min(10, num_nodes-1), which='SM')
-            fiedler_vector = eigenvecs[:, 1]  # Second smallest
-            
-            # Order by Fiedler vector values
-            order = np.argsort(fiedler_vector)
+            # Compute eigenvectors
+            k = min(10, num_nodes - 1)
+            try:
+                eigenvals, eigenvecs = eigsh(L, k=k, which='SM', sigma=0.0)
+                
+                # Use second smallest eigenvector (Fiedler vector)
+                if eigenvecs.shape[1] > 1:
+                    fiedler_vector = eigenvecs[:, 1]
+                else:
+                    fiedler_vector = eigenvecs[:, 0]
+                
+                # Order by Fiedler vector values
+                order = np.argsort(fiedler_vector)
+                
+            except Exception:
+                # Fallback to degree ordering
+                order = np.argsort(-degrees)
             
-            return torch.tensor(order, dtype=torch.long)
+            return torch.tensor(order, dtype=torch.long, device=device)
             
         except Exception as e:
             print(f"Spectral ordering failed: {e}, falling back to degree ordering")
@@ -88,35 +128,61 @@ class GraphSequencer:
     
     @staticmethod
     def degree_ordering(edge_index, num_nodes):
-        """Order nodes by degree (high to low)"""
-        # Count degrees
-        degrees = torch.zeros(num_nodes, dtype=torch.long)
-        degrees.index_add_(0, edge_index[0], torch.ones(edge_index.shape[1], dtype=torch.long))
-        degrees.index_add_(0, edge_index[1], torch.ones(edge_index.shape[1], dtype=torch.long))
+        """Order nodes by degree (high to low) - optimized version"""
+        device = edge_index.device
+        
+        # Count degrees efficiently
+        degrees = torch.zeros(num_nodes, dtype=torch.long, device=device)
+        
+        if edge_index.shape[1] > 0:
+            # Ensure valid indices
+            valid_mask = (edge_index[0] < num_nodes) & (edge_index[1] < num_nodes)
+            valid_edges = edge_index[:, valid_mask]
+            
+            if valid_edges.shape[1] > 0:
+                degrees.index_add_(0, valid_edges[0], torch.ones(valid_edges.shape[1], dtype=torch.long, device=device))
+                degrees.index_add_(0, valid_edges[1], torch.ones(valid_edges.shape[1], dtype=torch.long, device=device))
         
         # Sort by degree (descending), then by node index for determinism
-        _, order = torch.sort(-degrees * num_nodes - torch.arange(num_nodes))
+        node_indices = torch.arange(num_nodes, device=device)
+        _, order = torch.sort(-degrees * num_nodes - node_indices)
         
         return order
     
     @staticmethod
     def community_ordering(edge_index, num_nodes, n_clusters=None):
-        """Community-aware ordering using spectral clustering"""
+        """Community-aware ordering - robust version"""
+        device = edge_index.device
+        
+        if num_nodes <= 3:
+            return GraphSequencer.degree_ordering(edge_index, num_nodes)
+        
         try:
             if n_clusters is None:
-                n_clusters = max(2, min(10, num_nodes // 100))
+                n_clusters = max(2, min(10, int(np.sqrt(num_nodes))))
             
-            # Convert to adjacency matrix
-            edge_index_np = edge_index.cpu().numpy()
+            n_clusters = min(n_clusters, num_nodes)
+            
+            # Convert to adjacency matrix on CPU
+            edge_index_cpu = edge_index.cpu().numpy()
             A = np.zeros((num_nodes, num_nodes))
-            A[edge_index_np[0], edge_index_np[1]] = 1
-            A[edge_index_np[1], edge_index_np[0]] = 1
+            
+            valid_edges = (edge_index_cpu[0] < num_nodes) & (edge_index_cpu[1] < num_nodes)
+            valid_edge_index = edge_index_cpu[:, valid_edges]
+            
+            if valid_edge_index.shape[1] > 0:
+                A[valid_edge_index[0], valid_edge_index[1]] = 1
+                A[valid_edge_index[1], valid_edge_index[0]] = 1
+            
+            # Add small diagonal for stability
+            A += np.eye(num_nodes) * 0.01
             
             # Spectral clustering
             clustering = SpectralClustering(
                 n_clusters=n_clusters, 
                 affinity='precomputed',
-                random_state=42
+                random_state=42,
+                assign_labels='discretize'
             )
             
             labels = clustering.fit_predict(A)
@@ -127,36 +193,30 @@ class GraphSequencer:
             order = []
             for cluster in range(n_clusters):
                 cluster_nodes = np.where(labels == cluster)[0]
-                cluster_degrees = degrees[cluster_nodes]
-                cluster_order = cluster_nodes[np.argsort(-cluster_degrees)]
-                order.extend(cluster_order)
+                if len(cluster_nodes) > 0:
+                    cluster_degrees = degrees[cluster_nodes]
+                    cluster_order = cluster_nodes[np.argsort(-cluster_degrees)]
+                    order.extend(cluster_order)
             
-            return torch.tensor(order, dtype=torch.long)
+            # Add any missed nodes
+            for i in range(num_nodes):
+                if i not in order:
+                    order.append(i)
+            
+            return torch.tensor(order, dtype=torch.long, device=device)
             
         except Exception as e:
             print(f"Community ordering failed: {e}, falling back to BFS ordering")
             return GraphSequencer.bfs_ordering(edge_index, num_nodes)
-    
-    @staticmethod
-    def multi_view_ordering(edge_index, num_nodes):
-        """Generate multiple orderings for different perspectives"""
-        orderings = {}
-        
-        # Primary orderings
-        orderings['bfs'] = GraphSequencer.bfs_ordering(edge_index, num_nodes)
-        orderings['degree'] = GraphSequencer.degree_ordering(edge_index, num_nodes)
-        orderings['spectral'] = GraphSequencer.spectral_ordering(edge_index, num_nodes)
-        orderings['community'] = GraphSequencer.community_ordering(edge_index, num_nodes)
-        
-        return orderings
 
 class PositionalEncoder:
-    """Graph-aware positional encoding"""
+    """Graph-aware positional encoding - optimized version"""
     
     @staticmethod
     def encode_positions(x, edge_index, order, max_dist=10):
         """
         Create positional encodings that preserve graph structure
+        Optimized for training stability
         """
         num_nodes = x.size(0)
         device = x.device
@@ -164,30 +224,53 @@ class PositionalEncoder:
         # Sequential positions
         seq_pos = torch.zeros(num_nodes, device=device)
         seq_pos[order] = torch.arange(num_nodes, device=device, dtype=torch.float)
+        seq_pos = seq_pos / max(num_nodes, 1)
+        
+        # Enhanced distance encoding
+        distances = torch.zeros((num_nodes, max_dist), device=device)
         
-        # Graph distances (local neighborhood)
-        G = nx.Graph()
-        G.add_edges_from(edge_index.t().cpu().numpy())
-        
-        # Compute shortest path distances
-        distances = torch.full((num_nodes, max_dist), float('inf'), device=device)
-        
-        for i, node in enumerate(order):
-            # Get distances to previous nodes in sequence
-            start_idx = max(0, i - max_dist)
-            for j in range(start_idx, i):
-                prev_node = order[j].item()
-                try:
-                    dist = nx.shortest_path_length(G, source=node.item(), target=prev_node)
-                    distances[node, j - start_idx] = min(dist, max_dist - 1)
-                except nx.NetworkXNoPath:
-                    distances[node, j - start_idx] = max_dist - 1
-        
-        # Replace infinities with max distance
-        distances[distances == float('inf')] = max_dist - 1
-        
-        # Normalize
-        seq_pos = seq_pos / num_nodes
-        distances = distances / max_dist
+        if edge_index.shape[1] > 0:
+            # Create adjacency matrix efficiently
+            adj = torch.zeros(num_nodes, num_nodes, device=device, dtype=torch.bool)
+            
+            # Filter valid edges
+            valid_mask = (edge_index[0] < num_nodes) & (edge_index[1] < num_nodes)
+            if valid_mask.any():
+                valid_edges = edge_index[:, valid_mask]
+                adj[valid_edges[0], valid_edges[1]] = True
+                adj[valid_edges[1], valid_edges[0]] = True  # Undirected
+            
+            # Compute 2-hop neighbors for richer encoding
+            adj2 = torch.matmul(adj.float(), adj.float()) > 0
+            
+            # Fill distance features
+            for i, node in enumerate(order):
+                node_idx = node.item() if isinstance(node, torch.Tensor) else node
+                
+                if node_idx < num_nodes:
+                    # Get 1-hop and 2-hop neighbors
+                    neighbors_1hop = torch.where(adj[node_idx])[0]
+                    neighbors_2hop = torch.where(adj2[node_idx] & ~adj[node_idx])[0]
+                    
+                    # Fill distance features based on order position
+                    start_idx = max(0, i - max_dist)
+                    for j in range(start_idx, i):
+                        if j - start_idx < max_dist:
+                            prev_node = order[j]
+                            prev_idx = prev_node.item() if isinstance(prev_node, torch.Tensor) else prev_node
+                            
+                            if prev_idx < num_nodes:
+                                # Multi-scale distance encoding
+                                if prev_idx in neighbors_1hop:
+                                    distances[node_idx, j - start_idx] = 0.9  # Direct neighbor
+                                elif prev_idx in neighbors_2hop:
+                                    distances[node_idx, j - start_idx] = 0.6  # 2-hop neighbor
+                                else:
+                                    distances[node_idx, j - start_idx] = 0.3  # Distant
+        else:
+            # No edges - use position-based encoding
+            for i in range(num_nodes):
+                for j in range(max_dist):
+                    distances[i, j] = (max_dist - j) / max_dist
         
         return seq_pos.unsqueeze(1), distances