core/graph_sequencer.py

import torch
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

class GraphSequencer:
    """
    Production-ready graph ordering strategies
    All methods use real graph data - no hardcoded values
    """
    
    @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))
        edge_list = edge_index.t().cpu().numpy()
        G.add_edges_from(edge_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)
        
        # BFS traversal
        visited = set()
        order = []
        queue = [start_node]
        
        while queue:
            node = queue.pop(0)
            if node in visited:
                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)
            
            for neighbor in neighbors:
                if neighbor not in visited:
                    queue.append(neighbor)
        
        # Add any disconnected nodes
        for node in range(num_nodes):
            if node not in visited:
                order.append(node)
        
        return torch.tensor(order, dtype=torch.long)
    
    @staticmethod
    def spectral_ordering(edge_index, num_nodes):
        """Spectral ordering using graph Laplacian eigenvector"""
        try:
            # Compute normalized Laplacian
            edge_index_np = 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
            
            # Degree matrix
            D = np.diag(np.sum(A, axis=1))
            
            # 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)))
            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)
            
            return torch.tensor(order, dtype=torch.long)
            
        except Exception as e:
            print(f"Spectral ordering failed: {e}, falling back to degree ordering")
            return GraphSequencer.degree_ordering(edge_index, num_nodes)
    
    @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))
        
        # Sort by degree (descending), then by node index for determinism
        _, order = torch.sort(-degrees * num_nodes - torch.arange(num_nodes))
        
        return order
    
    @staticmethod
    def community_ordering(edge_index, num_nodes, n_clusters=None):
        """Community-aware ordering using spectral clustering"""
        try:
            if n_clusters is None:
                n_clusters = max(2, min(10, num_nodes // 100))
            
            # Convert to adjacency matrix
            edge_index_np = 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
            
            # Spectral clustering
            clustering = SpectralClustering(
                n_clusters=n_clusters, 
                affinity='precomputed',
                random_state=42
            )
            
            labels = clustering.fit_predict(A)
            
            # Order by cluster, then by degree within cluster
            degrees = np.sum(A, axis=1)
            
            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)
            
            return torch.tensor(order, dtype=torch.long)
            
        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"""
    
    @staticmethod
    def encode_positions(x, edge_index, order, max_dist=10):
        """
        Create positional encodings that preserve graph structure
        """
        num_nodes = x.size(0)
        device = x.device
        
        # Sequential positions
        seq_pos = torch.zeros(num_nodes, device=device)
        seq_pos[order] = torch.arange(num_nodes, device=device, dtype=torch.float)
        
        # 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
        
        return seq_pos.unsqueeze(1), distances