Update data/loader.py

data/loader.py

@@ -1,14 +1,13 @@
 import torch
-from torch_geometric.datasets import Planetoid, TUDataset
+from torch_geometric.datasets import Planetoid, TUDataset, Amazon, Coauthor
 from torch_geometric.loader import DataLoader
-from torch_geometric.transforms import NormalizeFeatures
+from torch_geometric.transforms import NormalizeFeatures, Compose
 import yaml
 import os
 
 class GraphDataLoader:
     """
-    Production data loading with real datasets only
-    Device-safe implementation
+    Production data loading with comprehensive dataset support
     """
     
     def __init__(self, config_path='config.yaml'):
@@ -16,7 +15,7 @@ class GraphDataLoader:
             with open(config_path, 'r') as f:
                 self.config = yaml.safe_load(f)
         else:
-            # Default config if file doesn't exist
+            # Default config
             self.config = {
                 'data': {
                     'batch_size': 32,
@@ -27,70 +26,139 @@ class GraphDataLoader:
         self.batch_size = self.config['data']['batch_size']
         self.test_split = self.config['data']['test_split']
         
+        # Standard transform
+        self.transform = Compose([
+            NormalizeFeatures()
+        ])
+        
     def load_node_classification_data(self, dataset_name='Cora'):
-        """Load real node classification datasets"""
+        """Load node classification datasets with proper splits"""
         
         try:
             if dataset_name in ['Cora', 'CiteSeer', 'PubMed']:
                 dataset = Planetoid(
                     root=f'./data/{dataset_name}',
                     name=dataset_name,
-                    transform=NormalizeFeatures()
+                    transform=self.transform
+                )
+                
+            elif dataset_name in ['Computers', 'Photo']:
+                dataset = Amazon(
+                    root=f'./data/Amazon{dataset_name}',
+                    name=dataset_name,
+                    transform=self.transform
+                )
+                
+            elif dataset_name in ['CS', 'Physics']:
+                dataset = Coauthor(
+                    root=f'./data/Coauthor{dataset_name}',
+                    name=dataset_name,
+                    transform=self.transform
                 )
+                
             else:
-                # Fallback to Cora
+                print(f"Unknown dataset {dataset_name}, falling back to Cora")
                 dataset = Planetoid(
                     root='./data/Cora',
                     name='Cora',
-                    transform=NormalizeFeatures()
+                    transform=self.transform
                 )
+                
         except Exception as e:
             print(f"Error loading {dataset_name}: {e}")
             # Fallback to Cora
             dataset = Planetoid(
                 root='./data/Cora',
                 name='Cora',
-                transform=NormalizeFeatures()
+                transform=self.transform
             )
-            
+        
+        # Ensure proper masks exist
+        data = dataset[0]
+        self._ensure_masks(data)
+        
         return dataset
     
+    def _ensure_masks(self, data):
+        """Ensure train/val/test masks exist"""
+        num_nodes = data.num_nodes
+        
+        if not hasattr(data, 'train_mask') or data.train_mask is None:
+            # Create random splits
+            indices = torch.randperm(num_nodes)
+            
+            train_size = int(0.6 * num_nodes)
+            val_size = int(0.2 * num_nodes)
+            
+            train_mask = torch.zeros(num_nodes, dtype=torch.bool)
+            val_mask = torch.zeros(num_nodes, dtype=torch.bool)
+            test_mask = torch.zeros(num_nodes, dtype=torch.bool)
+            
+            train_mask[indices[:train_size]] = True
+            val_mask[indices[train_size:train_size + val_size]] = True
+            test_mask[indices[train_size + val_size:]] = True
+            
+            data.train_mask = train_mask
+            data.val_mask = val_mask
+            data.test_mask = test_mask
+    
     def load_graph_classification_data(self, dataset_name='MUTAG'):
-        """Load real graph classification datasets"""
+        """Load graph classification datasets"""
         
-        valid_datasets = ['MUTAG', 'ENZYMES', 'PROTEINS', 'COLLAB', 'IMDB-BINARY']
+        valid_datasets = ['MUTAG', 'ENZYMES', 'PROTEINS', 'COLLAB', 'IMDB-BINARY', 'DD']
         
         try:
             if dataset_name not in valid_datasets:
-                dataset_name = 'MUTAG'  # Default fallback
+                dataset_name = 'MUTAG'
                 
             dataset = TUDataset(
                 root=f'./data/{dataset_name}',
                 name=dataset_name,
-                transform=NormalizeFeatures()
+                transform=self.transform
             )
+            
+            # Handle missing features
+            if dataset[0].x is None:
+                # Use degree as features
+                max_degree = 0
+                for data in dataset:
+                    if data.edge_index.shape[1] > 0:
+                        degree = torch.zeros(data.num_nodes)
+                        degree.index_add_(0, data.edge_index[0], torch.ones(data.edge_index.shape[1]))
+                        max_degree = max(max_degree, degree.max().item())
+                
+                for data in dataset:
+                    if data.edge_index.shape[1] > 0:
+                        degree = torch.zeros(data.num_nodes)
+                        degree.index_add_(0, data.edge_index[0], torch.ones(data.edge_index.shape[1]))
+                        data.x = degree.unsqueeze(1) / max(max_degree, 1)
+                    else:
+                        data.x = torch.zeros(data.num_nodes, 1)
+                        
         except Exception as e:
             print(f"Error loading {dataset_name}: {e}")
-            # Create a minimal synthetic dataset as fallback
+            # Create minimal synthetic dataset
             from torch_geometric.data import Data
-            dataset = [Data(
-                x=torch.randn(10, 5),
-                edge_index=torch.randint(0, 10, (2, 20)),
-                y=torch.randint(0, 2, (1,))
-            )]
+            dataset = [
+                Data(
+                    x=torch.randn(10, 5),
+                    edge_index=torch.randint(0, 10, (2, 20)),
+                    y=torch.randint(0, 2, (1,))
+                ) for _ in range(100)
+            ]
             
         return dataset
     
     def create_dataloaders(self, dataset, task_type='node_classification'):
-        """Create train/val/test splits"""
+        """Create train/val/test splits with dataloaders"""
         
         if task_type == 'node_classification':
-            # Use predefined splits for node classification
+            # Single graph with masks
             data = dataset[0]
-            return data, None, None  # Single graph with masks
+            return data, None, None
             
         elif task_type == 'graph_classification':
-            # Random split for graph classification
+            # Split dataset
             num_graphs = len(dataset)
             indices = torch.randperm(num_graphs)
             
@@ -108,42 +176,70 @@ class GraphDataLoader:
             return train_loader, val_loader, test_loader
             
     def get_dataset_info(self, dataset):
-        """Get dynamic dataset information"""
+        """Get comprehensive dataset information"""
         try:
             if hasattr(dataset, 'num_features'):
                 num_features = dataset.num_features
             else:
-                num_features = dataset[0].x.size(1)
+                num_features = dataset[0].x.size(1) if dataset[0].x is not None else 1
                 
             if hasattr(dataset, 'num_classes'):
                 num_classes = dataset.num_classes
             else:
-                if hasattr(dataset[0], 'y'):
+                if hasattr(dataset[0], 'y') and dataset[0].y is not None:
                     if len(dataset) > 1:
-                        all_labels = torch.cat([data.y.flatten() for data in dataset])
-                        num_classes = len(torch.unique(all_labels))
+                        all_labels = []
+                        for data in dataset:
+                            if data.y is not None:
+                                all_labels.extend(data.y.flatten().tolist())
+                        num_classes = len(set(all_labels)) if all_labels else 2
                     else:
                         num_classes = len(torch.unique(dataset[0].y))
                 else:
-                    num_classes = 2  # Default binary
+                    num_classes = 2
                     
             num_graphs = len(dataset)
-            avg_nodes = sum([data.num_nodes for data in dataset]) / len(dataset)
-            avg_edges = sum([data.num_edges for data in dataset]) / len(dataset)
+            
+            # Calculate statistics
+            total_nodes = sum([data.num_nodes for data in dataset])
+            total_edges = sum([data.num_edges for data in dataset])
+            
+            avg_nodes = total_nodes / num_graphs
+            avg_edges = total_edges / num_graphs
+            
+            # Additional statistics
+            node_counts = [data.num_nodes for data in dataset]
+            edge_counts = [data.num_edges for data in dataset]
+            
+            stats = {
+                'num_features': num_features,
+                'num_classes': num_classes,
+                'num_graphs': num_graphs,
+                'avg_nodes': avg_nodes,
+                'avg_edges': avg_edges,
+                'min_nodes': min(node_counts),
+                'max_nodes': max(node_counts),
+                'min_edges': min(edge_counts),
+                'max_edges': max(edge_counts),
+                'total_nodes': total_nodes,
+                'total_edges': total_edges
+            }
             
         except Exception as e:
             print(f"Error getting dataset info: {e}")
-            # Return defaults
-            num_features = 1433  # Cora default
-            num_classes = 7      # Cora default
-            num_graphs = 1
-            avg_nodes = 2708
-            avg_edges = 10556
-            
-        return {
-            'num_features': num_features,
-            'num_classes': num_classes,
-            'num_graphs': num_graphs,
-            'avg_nodes': avg_nodes,
-            'avg_edges': avg_edges
-        }
+            # Return safe defaults
+            stats = {
+                'num_features': 1433,
+                'num_classes': 7,
+                'num_graphs': 1,
+                'avg_nodes': 2708.0,
+                'avg_edges': 10556.0,
+                'min_nodes': 2708,
+                'max_nodes': 2708,
+                'min_edges': 10556,
+                'max_edges': 10556,
+                'total_nodes': 2708,
+                'total_edges': 10556
+            }
+            
+        return stats