EpilepsyNet_model.py

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim



def extract_upper_triangle(corr_matrices):
    """
    Extract upper triangles from correlation matrices
    
    Args:
        corr_matrices: numpy array of shape (n_segments, n_channels, n_channels)
        
    Returns:
        numpy array of shape (n_segments, n_features) where n_features = n_channels*(n_channels-1)/2
    """
    n_segments, n_channels, _ = corr_matrices.shape
    n_features = n_channels * (n_channels - 1) // 2
    
    flattened = np.zeros((n_segments, n_features))
    
    for i in range(n_segments):
        # Get upper triangle indices (excluding diagonal)
        upper_indices = np.triu_indices(n_channels, k=1)
        # Extract values
        flattened[i] = corr_matrices[i][upper_indices]
    
    return flattened

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.3):
        super(MultiHeadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
        
        # Linear projections for Q, K, V
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        
        # Final projection after concatenating heads
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        # Dropout
        self.dropout = nn.Dropout(dropout)
        
        # Softmax for attention weights
        self.softmax = nn.Softmax(dim=-1)
        
    def forward(self, x, mask=None):
        batch_size = x.size(0)
        
        # Project Q, K, V
        Q = self.q_proj(x)  # (batch_size, seq_len, embed_dim)
        K = self.k_proj(x)    # (batch_size, seq_len, embed_dim)
        V = self.v_proj(x)  # (batch_size, seq_len, embed_dim)
        
        # Split into multiple heads
        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
        K = K.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
        V = V.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
        
        # Calculate attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)  # (batch_size, num_heads, seq_len, seq_len)
        
        # Apply mask (if provided)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
            
        # Apply softmax to get attention weights
        attn_weights = self.softmax(scores)  # (batch_size, num_heads, seq_len, seq_len)
        attn_weights = self.dropout(attn_weights)
        
        # Calculate weighted output
        attn_output = torch.matmul(attn_weights, V)  # (batch_size, num_heads, seq_len, head_dim)
        
        # Recompose heads
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)  # (batch_size, seq_len, embed_dim)
        
        # Pass through final projection
        output = self.out_proj(attn_output)  # (batch_size, seq_len, embed_dim)
        
        return output, attn_weights

class PositionalEncoding(nn.Module):
    def __init__(self, embed_dim, max_seq_length=100):
        super(PositionalEncoding, self).__init__()
        
        # Create positional encoding matrix
        pe = torch.zeros(max_seq_length, embed_dim)
        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2).float() * (-np.log(10000.0) / embed_dim))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # Register as buffer (not a parameter)
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        # Add positional encoding to input
        # x: [batch_size, seq_len, embed_dim]
        return x + self.pe[:, :x.size(1)]

class TimeSeriesAttentionClassifier(nn.Module):
    def __init__(self, input_dim, embed_dim, num_heads, num_classes=2, dropout=0.2):
        super(TimeSeriesAttentionClassifier, self).__init__()
        
        # Project flattened correlation features to embedding space
        self.embedding = nn.Linear(input_dim, embed_dim)
        
        # Positional encoding
        self.pos_encoding = PositionalEncoding(embed_dim)
        
        # Multi-head attention
        self.attention = MultiHeadAttention(embed_dim, num_heads, dropout)
        
        # Layer normalization
        self.layer_norm1 = nn.LayerNorm(embed_dim)
        self.layer_norm2 = nn.LayerNorm(embed_dim)
        
        # Feed-forward network
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim * 4, embed_dim)
        )
        
        # Output layer
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim // 2, 1),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        # batch_size, seq_len, input_dim = x.shape
        
        # Project to embedding space
        x = self.embedding(x)
        
        # Add positional encoding
        x = self.pos_encoding(x)
        
        # Self-attention (use x for query, key, and value)
        residual = x
        x, attention_weights = self.attention(x)
        x = self.layer_norm1(x + residual)
        
        # Feed-forward network with residual connection
        residual = x
        x = self.ffn(x)
        x = self.layer_norm2(x + residual)
        
        # Global average pooling over sequence dimension
        x = torch.mean(x, dim=1)
        
        # Classification
        logits = self.classifier(x)
        
        return logits, attention_weights

def train_model(model, train_loader, val_loader, num_epochs=50, learning_rate=1e-4, weight_decay=1e-5,  patience=10, scheduler_factor=0.5, min_lr=1e-6):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    
    # Changed from CrossEntropyLoss to BCELoss for binary classification with sigmoid
    criterion = nn.BCELoss()
    # optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # Add L2 regularization through weight_decay parameter in Adam
    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    
    # Learning rate scheduler - reduce LR when validation loss plateaus
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, 
        mode='min',
        factor=scheduler_factor,
        patience=patience,
        verbose=True,
        min_lr=min_lr
    )

    train_losses = []
    val_losses = []
    val_accuracies = []
    
    
    # Track best model and early stopping
    best_val_loss = float('inf')
    best_model_state = None
    early_stop_counter = 0
    early_stop_patience = patience * 2  # Stop after 2x the scheduler patience
    
    for epoch in range(num_epochs):
        # Training
        model.train()
        train_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            # Convert labels to float and reshape for BCE loss
            labels = labels.float().view(-1, 1)
            
            optimizer.zero_grad()
            outputs, _ = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        
        train_loss /= len(train_loader)
        train_losses.append(train_loss)
        
        # Validation
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                # Convert labels to float and reshape for BCE loss
                labels = labels.float().view(-1, 1)
                
                outputs, _ = model(inputs)
                loss = criterion(outputs, labels)
                val_loss += loss.item()
                
                # For binary classification with sigmoid, prediction is 1 if output > 0.5
                predicted = (outputs > 0.5).float()
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        val_loss /= len(val_loader)
        val_losses.append(val_loss)
        
        accuracy = 100 * correct / total
        val_accuracies.append(accuracy)
        
        # Learning rate scheduler step based on validation loss
        scheduler.step(val_loss)
        
        # Print current learning rate
        current_lr = optimizer.param_groups[0]['lr']


        # Print epoch results
        print(f'Epoch {epoch+1}/{num_epochs}, LR: {current_lr:.6f}, Train Loss: {train_loss:.4f}, '
              f'Val Loss: {val_loss:.4f}, Val Accuracy: {accuracy:.2f}%')
        
        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_model_state = model.state_dict().copy()
            early_stop_counter = 0
        else:
            early_stop_counter += 1
        
        # Early stopping
        if early_stop_counter >= early_stop_patience:
            print(f"Early stopping triggered after {epoch+1} epochs")
            break

    # Load best model weights
    if best_model_state is not None:
        model.load_state_dict(best_model_state)
        print(f"Loaded best model with validation loss: {best_val_loss:.4f}")
    
    return train_losses, val_losses, val_accuracies