test_phase2.py

# %% [markdown]
# # Phase 2 Testing Script
# 
# This script tests all Phase 2 components:
# - Vision Transformer model
# - Training pipeline
# - Band importance analysis
# - Data loading integration

# %%
import os
import sys
import torch
import numpy as np
from pathlib import Path

# Add current directory to path for imports
sys.path.append('.')

# Test imports
try:
    from models.vision_transformer import create_eurosat_model, MultispectralVisionTransformer
    print("[OK] Vision Transformer model imported successfully")
except ImportError as e:
    print(f"[ERROR] Failed to import Vision Transformer: {e}")

try:
    from train_eurosat import EuroSATTrainer, create_training_config
    print("[OK] Training pipeline imported successfully")
except ImportError as e:
    print(f"[ERROR] Failed to import training pipeline: {e}")

try:
    from band_importance import BandImportanceAnalyzer
    print("[OK] Band importance analyzer imported successfully")
except ImportError as e:
    print(f"[ERROR] Failed to import band importance analyzer: {e}")

try:
    from data_loader_eurosat import create_eurosat_dataloaders
    print("[OK] EuroSAT data loader imported successfully")
except ImportError as e:
    print(f"[ERROR] Failed to import EuroSAT data loader: {e}")

# %%
def test_vision_transformer():
    """Test Vision Transformer model creation and forward pass."""
    print("\n[TEST] Testing Vision Transformer...")
    
    # Create model
    model = create_eurosat_model(
        num_classes=10,
        model_size='tiny',  # Use tiny for faster testing
        use_pretrained=False
    )
    
    # Test forward pass
    batch_size = 4
    x = torch.randn(batch_size, 13, 64, 64)
    
    with torch.no_grad():
        output = model(x)
        print(f"[OK] Model forward pass successful")
        print(f"   Input shape: {x.shape}")
        print(f"   Output shape: {output.shape}")
        print(f"   Model parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # Test attention weights (if supported)
    try:
        with torch.no_grad():
            attention_weights = model.get_attention_weights(x)
            print(f"[OK] Attention weights extraction successful")
            print(f"   Attention weights shape: {attention_weights.shape}")
    except Exception as e:
        print(f"[WARNING] Attention weights extraction failed: {e}")
    
    return model

# %%
def test_data_loading():
    """Test data loading functionality."""
    print("\n[TEST] Testing Data Loading...")
    
    try:
        # Create data loaders with small batch size for testing
        train_loader, val_loader, test_loader, metadata = create_eurosat_dataloaders(
            data_dir='./eurosat_npz',
            batch_size=8,
            train_split=0.7,
            val_split=0.15,
            test_split=0.15,
            use_augmentation=False,  # Disable augmentation for testing
            num_workers=0  # Use 0 workers for testing
        )
        
        print(f"[OK] Data loaders created successfully")
        print(f"   Train batches: {len(train_loader)}")
        print(f"   Val batches: {len(val_loader)}")
        print(f"   Test batches: {len(test_loader)}")
        print(f"   Classes: {metadata['class_names']}")
        print(f"   Bands: {metadata['band_names']}")
        
        # Test a single batch
        for images, labels in train_loader:
            print(f"[OK] Sample batch loaded successfully")
            print(f"   Images shape: {images.shape}")
            print(f"   Labels shape: {labels.shape}")
            print(f"   Image range: [{images.min():.2f}, {images.max():.2f}]")
            break
        
        return train_loader, val_loader, test_loader, metadata
        
    except Exception as e:
        print(f"[ERROR] Data loading failed: {e}")
        return None, None, None, None

# %%
def test_training_pipeline():
    """Test training pipeline components."""
    print("\n[TEST] Testing Training Pipeline...")
    
    # Create model and data loaders
    model = create_eurosat_model(num_classes=10, model_size='tiny')
    train_loader, val_loader, test_loader, metadata = test_data_loading()
    
    if train_loader is None:
        print("[ERROR] Skipping training pipeline test due to data loading failure")
        return
    
    # Create training configuration
    config = create_training_config()
    config['epochs'] = 1  # Just test for 1 epoch
    config['batch_size'] = 8
    
    # Create trainer
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    trainer = EuroSATTrainer(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        test_loader=test_loader,
        device=device,
        config=config,
        class_names=metadata['class_names']
    )
    
    print(f"[OK] Trainer created successfully")
    print(f"   Device: {device}")
    print(f"   Optimizer: {type(trainer.optimizer).__name__}")
    print(f"   Scheduler: {type(trainer.scheduler).__name__ if trainer.scheduler else 'None'}")
    
    # Test one training epoch
    try:
        train_loss, train_acc = trainer.train_epoch()
        print(f"[OK] Training epoch completed")
        print(f"   Train Loss: {train_loss:.4f}")
        print(f"   Train Acc: {train_acc:.4f}")
    except Exception as e:
        print(f"[ERROR] Training epoch failed: {e}")
    
    # Test validation
    try:
        val_loss, val_acc, val_metrics = trainer.validate_epoch()
        print(f"[OK] Validation completed")
        print(f"   Val Loss: {val_loss:.4f}")
        print(f"   Val Acc: {val_acc:.4f}")
        print(f"   Val F1: {val_metrics['f1_macro']:.4f}")
    except Exception as e:
        print(f"[ERROR] Validation failed: {e}")

# %%
def test_band_importance():
    """Test band importance analysis."""
    print("\n[TEST] Testing Band Importance Analysis...")
    
    # Create model and data loaders
    model = create_eurosat_model(num_classes=10, model_size='tiny')
    _, _, test_loader, metadata = test_data_loading()
    
    if test_loader is None:
        print("[ERROR] Skipping band importance test due to data loading failure")
        return
    
    # Create analyzer
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    analyzer = BandImportanceAnalyzer(
        model=model,
        test_loader=test_loader,
        device=device,
        band_names=metadata['band_names'],
        class_names=metadata['class_names']
    )
    
    print(f"[OK] Band importance analyzer created successfully")
    
    # Test baseline performance
    try:
        baseline_results = analyzer.compute_baseline_performance()
        print(f"[OK] Baseline performance computed")
        print(f"   Accuracy: {baseline_results['accuracy']:.4f}")
        print(f"   F1 Macro: {baseline_results['f1_macro']:.4f}")
    except Exception as e:
        print(f"[ERROR] Baseline performance computation failed: {e}")
        return
    
    # Test ablation study (faster than permutation)
    try:
        ablation_results = analyzer.ablation_study()
        print(f"[OK] Ablation study completed")
        print(f"   Top band: {list(ablation_results['ablation_results'].keys())[0]}")
    except Exception as e:
        print(f"[ERROR] Ablation study failed: {e}")

# %%
def main():
    """Run all tests."""
    print("[START] Starting Phase 2 Component Tests")
    print("=" * 50)
    
    # Test Vision Transformer
    model = test_vision_transformer()
    
    # Test data loading
    test_data_loading()
    
    # Test training pipeline
    test_training_pipeline()
    
    # Test band importance
    test_band_importance()
    
    print("\n" + "=" * 50)
    print("[OK] Phase 2 testing completed!")
    print("\nNext steps:")
    print("1. Run full training: python train_eurosat.py")
    print("2. Analyze band importance: python band_importance.py")
    print("3. View results in logs/ and results/ directories")

# %%
if __name__ == "__main__":
    main()