Update demo.py

demo.py

@@ -1,18 +1,21 @@
 #!/usr/bin/env python3
 """
-Quick demo script to test Mamba Graph implementation
-Device-safe version
+Complete test script for Mamba Graph implementation
+Tests training, evaluation, and visualization
 """
 
 import torch
 import os
+import time
 from core.graph_mamba import GraphMamba
+from core.trainer import GraphMambaTrainer
 from data.loader import GraphDataLoader
 from utils.metrics import GraphMetrics
+from utils.visualization import GraphVisualizer
 
 def main():
-    print("🧠 Testing Mamba Graph Neural Network")
-    print("=" * 50)
+    print("🧠 Mamba Graph Neural Network - Complete Test")
+    print("=" * 60)
     
     # Configuration
     config = {
@@ -28,6 +31,14 @@ def main():
             'batch_size': 16,
             'test_split': 0.2
         },
+        'training': {
+            'learning_rate': 0.01,
+            'weight_decay': 0.0005,
+            'epochs': 50,  # Quick test
+            'patience': 10,
+            'warmup_epochs': 5,
+            'min_lr': 1e-6
+        },
         'ordering': {
             'strategy': 'bfs',
             'preserve_locality': True
@@ -39,7 +50,7 @@ def main():
         device = torch.device('cpu')
     else:
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    print(f"Device: {device}")
+    print(f"💾 Device: {device}")
     
     # Load dataset
     print("\n📊 Loading Cora dataset...")
@@ -48,13 +59,15 @@ def main():
         dataset = data_loader.load_node_classification_data('Cora')
         data = dataset[0].to(device)
         
-        # Dataset info
         info = data_loader.get_dataset_info(dataset)
-        print(f"✅ Success!")
-        print(f"Nodes: {data.num_nodes}")
-        print(f"Edges: {data.num_edges}")
-        print(f"Features: {info['num_features']}")
-        print(f"Classes: {info['num_classes']}")
+        print(f"✅ Dataset loaded successfully!")
+        print(f"   Nodes: {data.num_nodes:,}")
+        print(f"   Edges: {data.num_edges:,}")
+        print(f"   Features: {info['num_features']}")
+        print(f"   Classes: {info['num_classes']}")
+        print(f"   Train nodes: {data.train_mask.sum()}")
+        print(f"   Val nodes: {data.val_mask.sum()}")
+        print(f"   Test nodes: {data.test_mask.sum()}")
         
     except Exception as e:
         print(f"❌ Error loading dataset: {e}")
@@ -66,70 +79,143 @@ def main():
         model = GraphMamba(config).to(device)
         total_params = sum(p.numel() for p in model.parameters())
         print(f"✅ Model initialized!")
-        print(f"Parameters: {total_params:,}")
+        print(f"   Parameters: {total_params:,}")
+        print(f"   Memory usage: ~{total_params * 4 / 1024**2:.1f} MB")
         
     except Exception as e:
         print(f"❌ Error initializing model: {e}")
         return
     
-    # Forward pass test
+    # Test forward pass
     print("\n🚀 Testing forward pass...")
     try:
         model.eval()
         with torch.no_grad():
             h = model(data.x, data.edge_index)
             print(f"✅ Forward pass successful!")
-            print(f"Input shape: {data.x.shape}")
-            print(f"Output shape: {h.shape}")
-            print(f"Output range: [{h.min():.3f}, {h.max():.3f}]")
+            print(f"   Input shape: {data.x.shape}")
+            print(f"   Output shape: {h.shape}")
+            print(f"   Output range: [{h.min():.3f}, {h.max():.3f}]")
             
     except Exception as e:
         print(f"❌ Forward pass failed: {e}")
         return
     
-    # Test different ordering strategies
+    # Test ordering strategies
     print("\n🔄 Testing ordering strategies...")
-    
     strategies = ['bfs', 'spectral', 'degree', 'community']
     
     for strategy in strategies:
         try:
             config['ordering']['strategy'] = strategy
-            model_test = GraphMamba(config).to(device)
-            model_test.eval()
+            test_model = GraphMamba(config).to(device)
+            test_model.eval()
             
+            start_time = time.time()
             with torch.no_grad():
-                h = model_test(data.x, data.edge_index)
-            print(f"✅ {strategy}: Success - Shape {h.shape}")
+                h = test_model(data.x, data.edge_index)
+            end_time = time.time()
+            
+            print(f"✅ {strategy:12} | Shape: {h.shape} | Time: {(end_time-start_time)*1000:.2f}ms")
             
         except Exception as e:
-            print(f"❌ {strategy}: Failed - {str(e)}")
+            print(f"❌ {strategy:12} | Failed: {str(e)}")
     
-    # Test evaluation
-    print("\n📈 Testing evaluation...")
+    # Initialize trainer
+    print("\n🏋️ Testing training system...")
     try:
-        # Initialize classifier
-        num_classes = info['num_classes']
-        model._init_classifier(num_classes, device)
+        trainer = GraphMambaTrainer(model, config, device)
+        print(f"✅ Trainer initialized!")
+        print(f"   Optimizer: {type(trainer.optimizer).__name__}")
+        print(f"   Learning rate: {trainer.lr}")
+        print(f"   Epochs: {trainer.epochs}")
         
-        # Create test mask if not available
-        if hasattr(data, 'test_mask'):
-            mask = data.test_mask
-        else:
-            mask = torch.zeros(data.num_nodes, dtype=torch.bool, device=device)
-            mask[data.num_nodes//2:] = True
+    except Exception as e:
+        print(f"❌ Trainer initialization failed: {e}")
+        return
+    
+    # Run training
+    print("\n🎯 Running training...")
+    try:
+        start_time = time.time()
+        history = trainer.train_node_classification(data, verbose=True)
+        training_time = time.time() - start_time
+        
+        print(f"✅ Training completed!")
+        print(f"   Training time: {training_time:.2f}s")
+        print(f"   Epochs trained: {len(history['train_loss'])}")
+        print(f"   Best val accuracy: {trainer.best_val_acc:.4f}")
+        
+    except Exception as e:
+        print(f"❌ Training failed: {e}")
+        return
+    
+    # Test evaluation
+    print("\n📊 Testing evaluation...")
+    try:
+        test_results = trainer.test(data)
+        print(f"✅ Evaluation completed!")
+        print(f"   Test accuracy: {test_results['test_acc']:.4f}")
+        print(f"   Test loss: {test_results['test_loss']:.4f}")
         
-        metrics = GraphMetrics.evaluate_node_classification(model, data, mask, device)
-        print("✅ Evaluation successful!")
-        for metric, value in metrics.items():
-            if isinstance(value, float):
-                print(f"  {metric}: {value:.4f}")
+        # Per-class results
+        class_accs = test_results['class_acc']
+        print(f"   Per-class accuracy:")
+        for i, acc in enumerate(class_accs):
+            print(f"     Class {i}: {acc:.4f}")
             
     except Exception as e:
         print(f"❌ Evaluation failed: {e}")
+        return
+    
+    # Test visualization
+    print("\n🎨 Testing visualization...")
+    try:
+        # Create visualizations
+        graph_fig = GraphVisualizer.create_graph_plot(data, max_nodes=200)
+        metrics_fig = GraphVisualizer.create_metrics_plot(test_results)
+        training_fig = GraphVisualizer.create_training_history_plot(history)
+        
+        print(f"✅ Visualizations created!")
+        print(f"   Graph plot: {type(graph_fig).__name__}")
+        print(f"   Metrics plot: {type(metrics_fig).__name__}")
+        print(f"   Training plot: {type(training_fig).__name__}")
+        
+        # Save plots
+        graph_fig.write_html("graph_visualization.html")
+        metrics_fig.write_html("metrics_plot.html")
+        training_fig.write_html("training_history.html")
+        print(f"   Plots saved as HTML files")
+        
+    except Exception as e:
+        print(f"❌ Visualization failed: {e}")
+    
+    # Performance summary
+    print("\n🏆 Performance Summary")
+    print("=" * 40)
+    print(f"📊 Dataset: Cora ({data.num_nodes:,} nodes)")
+    print(f"🧠 Model: {total_params:,} parameters")
+    print(f"⏱️ Training: {training_time:.2f}s ({len(history['train_loss'])} epochs)")
+    print(f"🎯 Test Accuracy: {test_results['test_acc']:.4f} ({test_results['test_acc']*100:.2f}%)")
+    print(f"🏅 Best Val Accuracy: {trainer.best_val_acc:.4f} ({trainer.best_val_acc*100:.2f}%)")
+    
+    # Compare with baselines
+    cora_baselines = {
+        'GCN': 0.815,
+        'GAT': 0.830,
+        'GraphSAGE': 0.824,
+        'GIN': 0.800
+    }
+    
+    print(f"\n📈 Comparison with Baselines:")
+    test_acc = test_results['test_acc']
+    for model_name, baseline in cora_baselines.items():
+        diff = test_acc - baseline
+        status = "🟢" if diff > 0 else "🟡" if diff > -0.05 else "🔴"
+        print(f"   {status} {model_name:12}: {baseline:.3f} (diff: {diff:+.3f})")
     
-    print("\n✨ Demo completed!")
-    print("🚀 Ready for production deployment!")
+    print(f"\n✨ Test completed successfully!")
+    print(f"🚀 Ready for production deployment!")
 
 if __name__ == "__main__":
     main()