CA/models/spatiotemporal_gcn/sampler.py

#!/usr/bin/env python3
"""
GraphSAINT-style Sampler for PyTorch Geometric.

Mini-batch sampler for large graphs (140k+ nodes) using neighbor sampling.
Compatible with base PyG installation (no torch-sparse or pyg-lib required).

Usage:
    from models.spatiotemporal_gcn.sampler import GraphSAINTSampler
    
    sampler = GraphSAINTSampler(
        data=data,
        batch_size=256,
        num_neighbors=[256, 128, 64]
    )
"""

import torch
from torch.utils.data import DataLoader, Dataset
from torch_geometric.data import Data
from torch_geometric.utils import subgraph


class GraphSAINTDataset(Dataset):
    """Dataset that samples node indices for mini-batching."""
    
    def __init__(self, num_nodes: int, num_steps: int = 10):
        self.num_nodes = num_nodes
        self.num_steps = num_steps
    
    def __len__(self):
        return self.num_steps
    
    def __getitem__(self, idx):
        return torch.randint(0, self.num_nodes, (1,))


class GraphSAINTSampler:
    """
    GraphSAINT-style mini-batch sampler for large graphs.
    
    Implements neighbor sampling to create subgraphs that fit in GPU memory.
    For each batch, samples seed nodes and their multi-hop neighbors.
    
    Args:
        data: Full graph with edge_index and node features.
        batch_size: Seed nodes per batch (default: 256).
        num_neighbors: Neighbors per layer [layer0, layer1, ...].
            Default: [256, 128, 64] for 3-layer GCN.
        num_steps: Batches per epoch (default: 10).
    """
    
    def __init__(
        self,
        data: Data,
        batch_size: int = 256,
        num_neighbors: list = None,
        num_steps: int = 10,
    ):
        if num_neighbors is None:
            num_neighbors = [256, 128, 64]
        
        self.data = data
        self.batch_size = batch_size
        self.num_neighbors = num_neighbors
        self.num_steps = num_steps
        self.num_nodes = data.num_nodes
        self.edge_index = data.edge_index
        
        if data.num_nodes > 100000:
            print(f"Sampler for large graph: {data.num_nodes:,} nodes")
            print(f"  Batch size: {batch_size}")
            print(f"  Layer depths: {num_neighbors}")
    
    def _sample_neighbors(self, seed_nodes: torch.Tensor) -> torch.Tensor:
        """
        Sample multi-hop neighbors for seed nodes.
        
        Args:
            seed_nodes: Initial node indices.
        
        Returns:
            All sampled node indices (seed + neighbors).
        """
        sampled = seed_nodes.unique()
        
        for num_neighbors in self.num_neighbors:
            if len(sampled) == 0:
                break
            
            mask = torch.isin(self.edge_index[0], sampled)
            neighbor_edges = self.edge_index[:, mask]
            
            if neighbor_edges.shape[1] == 0:
                break
            
            neighbors = neighbor_edges[1]
            
            if len(neighbors) > num_neighbors:
                neighbors = neighbors[torch.randperm(len(neighbors))[:num_neighbors]]
            
            sampled = torch.cat([sampled, neighbors]).unique()
        
        return sampled
    
    def _create_subgraph(self, node_indices: torch.Tensor) -> Data:
        edge_index, _, edge_mask = subgraph(
            node_indices,
            self.edge_index,
            relabel_nodes=True,
            return_edge_mask=True,
        )
        
        subgraph_data = Data(
            x=self.data.x[node_indices],
            edge_index=edge_index,
            n_id=node_indices,
        )
        
        if hasattr(self.data, 'y') and self.data.y is not None:
            subgraph_data.y = self.data.y[node_indices]
        
        return subgraph_data
    
    def __iter__(self):
        for _ in range(self.num_steps):
            seed_nodes = torch.randint(0, self.num_nodes, (self.batch_size,))
            sampled_nodes = self._sample_neighbors(seed_nodes)
            batch = self._create_subgraph(sampled_nodes)
            yield batch
    
    def __len__(self):
        return self.num_steps


class GraphSAINTConfig:
    """Configuration for GraphSAINT-style sampling."""
    
    def __init__(
        self,
        batch_size: int = 256,
        num_neighbors: list = None,
        num_steps: int = 10,
    ):
        self.batch_size = batch_size
        self.num_neighbors = num_neighbors if num_neighbors is not None else [256, 128, 64]
        self.num_steps = num_steps
    
    def __repr__(self):
        return (
            f"GraphSAINTConfig(\n"
            f"  batch_size={self.batch_size},\n"
            f"  num_neighbors={self.num_neighbors},\n"
            f"  num_steps={self.num_steps}\n"
            f")"
        )


def create_graph_saint_loader(
    data: Data,
    batch_size: int = 256,
    num_neighbors: list = None,
    num_steps: int = 10,
):
    """
    Create a GraphSAINT-style sampler for large graph training.
    
    Args:
        data: Full graph data with edge_index and features.
        batch_size: Seed nodes per batch (default: 256).
        num_neighbors: Layer-wise neighbor counts (default: [256, 128, 64]).
        num_steps: Batches per epoch (default: 10).
    
    Returns:
        GraphSAINTSampler: Mini-batch iterator.
    """
    return GraphSAINTSampler(
        data=data,
        batch_size=batch_size,
        num_neighbors=num_neighbors,
        num_steps=num_steps,
    )


def main():
    """Example usage with dummy data."""
    print("=" * 60)
    print("GraphSAINT-style Sampler Demo")
    print("=" * 60)
    
    print("\nCreating dummy graph (10k nodes)...")
    N = 10000
    num_features = 48
    
    edge_index = torch.randint(0, N, (2, N * 3))
    x = torch.randn(N, num_features)
    y = torch.randint(0, 3, (N,))
    
    data = Data(x=x, y=y, edge_index=edge_index)
    print(f"  Nodes: {data.num_nodes:,}")
    print(f"  Edges: {data.num_edges:,}")
    print(f"  Features: {data.num_node_features}")
    
    print("\nCreating sampler...")
    config = GraphSAINTConfig(
        batch_size=256,
        num_neighbors=[256, 128, 64],
        num_steps=5,
    )
    print(config)
    
    loader = create_graph_saint_loader(
        data=data,
        batch_size=config.batch_size,
        num_neighbors=config.num_neighbors,
        num_steps=config.num_steps,
    )
    
    print(f"\nIterating through {len(loader)} batches...")
    for i, batch in enumerate(loader):
        print(f"  Batch {i+1}/{len(loader)}:")
        print(f"    Nodes: {batch.num_nodes:,}")
        print(f"    Edges: {batch.num_edges:,}")
        print(f"    Features: {batch.x.shape}")
        print(f"    Node IDs: {batch.n_id.shape}")
        
        if i >= 2:
            break
    
    print("\n" + "=" * 60)
    print("Sampler ready for training!")
    print("=" * 60)
    
    print("\nFor your 140k node graph:")
    print("  1. Load graph: data = load_your_graph()")
    print("  2. Create loader: loader = create_graph_saint_loader(data, batch_size=256)")
    print("  3. Train: for batch in loader: out = model(batch.x, batch.edge_index)")
    print("\nRecommended for 4GB GPU:")
    print("  - batch_size: 256")
    print("  - num_neighbors: [256, 128, 64]")


if __name__ == '__main__':
    main()