bl_utils.py

import numpy as np
import pandas as pd
import pickle
import json
import random
from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score, roc_curve, roc_auc_score, precision_recall_curve, average_precision_score, confusion_matrix, balanced_accuracy_score
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV, StratifiedKFold
from torch.utils.data import DataLoader, Dataset
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
import matplotlib.pyplot as plt
from collections import defaultdict
import shap

class RecurrenceDataset(Dataset):
    def __init__(self, features, labels, transform = None, target_transform = None):
        self.labels = labels
        self.features = features

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        label = self.labels[idx]
        features = self.features[idx, :]
        return features, label
    
    def print(self):
        print("labels", self.labels)
        print("features", self.features)

class RecurrenceClassifier(torch.nn.Module):
    def __init__(self, activation, input_dim, hidden_dim1, hidden_dim2, output_dim, dropout_rate):
        super().__init__()

        self.activation_func = activation

        self.layer1 =  torch.nn.Sequential(
            torch.nn.Linear(input_dim, hidden_dim1),
            # torch.nn.BatchNorm1d(hidden_dim1),
            torch.nn.LayerNorm(hidden_dim1),
            self.activation_func(),
            torch.nn.Dropout(dropout_rate)
            )   

        self.layer2 =  torch.nn.Sequential(
            torch.nn.Linear(hidden_dim1, hidden_dim2),
            # torch.nn.BatchNorm1d(hidden_dim2),
            torch.nn.LayerNorm(hidden_dim2),
            self.activation_func(),
            torch.nn.Dropout(dropout_rate)
            )

        self.output_layer = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim2, output_dim),
            torch.nn.Sigmoid()
        )        

        self.output_layer_logits = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim2, output_dim)
        )        

    def forward(self, x):
        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        output = self.output_layer(out2)
        return output
    
    # def forward_logits(self, x):
    #     out1 = self.layer1(x)
    #     out2 = self.layer2(out1)
    #     output = self.output_layer_logits(out2)
    #     return output

# class RecurrenceDatasetBranched(Dataset):
#     def __init__(self, clinical_features, expression_features, mutation_features, labels, transform = None, target_transform = None):
#         self.labels = labels
#         self.clin_features = clinical_features
#         self.mut_features = mutation_features
#         self.exp_features = expression_features

#     def __len__(self):
#         return len(self.labels)

#     def __getitem__(self, idx):
#         clin = self.clin_features[idx, :]
#         mut = self.mut_features[idx, :]
#         exp = self.exp_features[idx, :]
#         label = self.labels[idx]
#         return (clin, exp, mut), label
    
#     def print(self):
#         print("labels", self.labels)
#         print("clinical features", self.clin_features)
#         print("mutation features", self.mut_features)
#         print("expression features", self.exp_features)

# class RecurrenceClassifierBranched(torch.nn.Module):
#     def __init__(self, activation, clin_input_dim, exp_input_dim, mut_input_dim, clin_dim1, exp_dim1, mut_dim1, hidden_dim1, hidden_dim2, output_dim, dropout_rate):
#         super().__init__()

#         self.activation_func = activation
        
#         self.clin_model = torch.nn.Sequential(
#             torch.nn.Linear(clin_input_dim, clin_dim1),
#             torch.nn.LayerNorm(clin_dim1),
#             torch.nn.LeakyReLU(),
#             torch.nn.Dropout(dropout_rate),
#             )
        
#         self.mut_model =  torch.nn.Sequential(
#             torch.nn.Linear(mut_input_dim, mut_dim1),
#             torch.nn.LayerNorm(mut_dim1),
#             torch.nn.LeakyReLU(),
#             torch.nn.Dropout(dropout_rate),
#             )

#         self.exp_model = torch.nn.Sequential(
#             torch.nn.Linear(exp_input_dim, exp_dim1),
#             torch.nn.LayerNorm(exp_dim1),
#             torch.nn.LeakyReLU(),
#             torch.nn.Dropout(dropout_rate),
#             )
        
#         combined_dim = clin_dim1 + exp_dim1
#         self.combined_model =  torch.nn.Sequential(
#             torch.nn.Linear(combined_dim, hidden_dim1),
#             torch.nn.LayerNorm(hidden_dim1),
#             torch.nn.LeakyReLU(),
#             torch.nn.Dropout(dropout_rate),
#             torch.nn.Linear(hidden_dim1, hidden_dim2),  
#             torch.nn.LeakyReLU(),
#             torch.nn.Linear(hidden_dim2, output_dim),
#             torch.nn.Sigmoid()
#             )
        
#     #take in a list of tensors
#     def forward(self, x):
#         clin, exp = x
#         clin_output = self.clin_model(clin)
#         mut_output = self.mut_model(mut)
#         exp_output = self.exp_model(exp)
#         combined = torch.cat((clin_output, mut_output, exp_output), dim = 1)
#         # combined = torch.cat((clin_output, exp_output), dim = 1)
#         output = self.combined_model(combined)

#         return output

class EarlyStopping:
    def __init__(self, patience=5, delta=0):
        self.patience = patience
        self.delta = delta
        self.best_score = None
        self.early_stop = False
        self.counter = 0
        self.best_model_state = None

    def __call__(self, loss, model):
        score = -loss
        if self.best_score is None:
            self.best_score = score
            self.best_model_state = model.state_dict()
        elif score < self.best_score + self.delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.best_model_state = model.state_dict()
            self.counter = 0

    def load_best_model(self, model):
        model.load_state_dict(self.best_model_state)

def train_validate_loop(epochs, model, device, train_dataloader, val_dataloader, optimizer, loss_fn = torch.nn.BCELoss(), scheduler=None, early_stopping = None, print_bool = False, print_early_stop = True):
    overall_train_losses = []
    overall_val_losses = []

    for epoch in range(epochs):
        train_loss = 0.0
        #training loop
        model.train()
        for i, data in enumerate(train_dataloader):
            #extract inputs and labels
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            #zero your gradients for every batch
            optimizer.zero_grad()
            # Forward Pass
            outputs = model(inputs)
            #compute the loss and its gradients
            loss = loss_fn(outputs, labels)
            loss.backward()
            #adjust learning weights
            optimizer.step()
            #record loss
            train_loss += loss.item() * inputs.size(0)

            if type(scheduler) == torch.optim.lr_scheduler.CyclicLR: 
                scheduler.step()
        
        current_train_loss = train_loss / len(train_dataloader.dataset)
        overall_train_losses.append(current_train_loss)
        train_metrics = compute_metrics_dl(model, device, train_dataloader)

        val_loss = 0.0
        #validation loop
        model.eval()
        with torch.no_grad():
            for i, data in enumerate(val_dataloader):
                inputs, labels = data
                inputs, labels = inputs.to(device), labels.to(device)
                val_outputs = model(inputs)
                loss = loss_fn(val_outputs, labels)
                val_loss += loss.item() * inputs.size(0)
            
        current_val_loss = val_loss / len(val_dataloader.dataset)
        overall_val_losses.append(current_val_loss)
        val_metrics = compute_metrics_dl(model, device, val_dataloader)

        if type(scheduler) == torch.optim.lr_scheduler.ReduceLROnPlateau: 
            scheduler.step(val_metrics['AUC']) #changed from curent_val_loss

        if early_stopping:
            early_stopping(val_metrics['AUC'], model) #changed from curent_val_loss
            if early_stopping.early_stop:
                if print_early_stop:
                    print(f"Early stopping triggered at epoch {epoch}")
                break

        if print_bool:
            print(f"Epoch {epoch+1}/{epochs} | Train Loss: {current_train_loss:.4f} | Val Loss: {current_val_loss:.4f}")
            train_df = pd.DataFrame(train_metrics, index=['Training'])
            val_df = pd.DataFrame(val_metrics, index=['Validation'])
            result = pd.concat([train_df, val_df], axis=0)
            display(result)

    return train_metrics, val_metrics, overall_train_losses, overall_val_losses

def test_model(model, dataloader, device, loss_fn = torch.nn.BCELoss(), print_bool = True):
    model.eval()
    total_loss = 0.0

    with torch.no_grad():
        for i, data in enumerate(dataloader):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)

            # Forward Pass: Make predictions for this batch
            outputs = model(inputs)

            #Update running loss
            loss = loss_fn(outputs, labels)
            total_loss += loss.item() * inputs.size(0)

    overall_loss = total_loss / len(dataloader.dataset)
    metrics = compute_metrics_dl(model, device, dataloader)
    if print_bool:
        display(pd.DataFrame(metrics, index=['Testing']))
        print(f'Testing Loss: {overall_loss}')

    return metrics, overall_loss

def plot_losses(train_losses, val_losses, model_title):
    plt.figure(figsize=(8, 5))
    plt.plot(range(1, len(train_losses) + 1), train_losses, label='Training Loss', marker='o')
    plt.plot(range(1, len(val_losses) + 1), val_losses, label='Validation Loss', marker='o')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'Training and Validation Loss for {model_title}')
    plt.legend()
    plt.grid(True)
    plt.show()

def compute_metrics_dl(model, device, dataloader, all = True):
    """
    Returns the accuracy, recall, precision, f1 score, and AUC
    """
    model.eval()
    all_labels = []
    all_predictions = []
    all_probs = []

    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            
            probs = outputs.cpu()
            # Apply threshold to get binary predictions
            predictions = (probs >= 0.5).int().cpu().numpy()
            
            # Collect all predictions, probabilities, and true labels
            all_probs.extend(probs)
            all_predictions.extend(predictions)
            all_labels.extend(labels.cpu().numpy())
            
    acc = accuracy_score(all_labels, all_predictions)
    rec = recall_score(all_labels, all_predictions)
    prec = precision_score(all_labels, all_predictions, zero_division=0)
    f1 = f1_score(all_labels, all_predictions)
    auc = roc_auc_score(all_labels, all_probs)
    
    metrics_dict = {"Accuracy": acc, "Recall" : rec, "Precision" : prec, "F1 Score" : f1, "AUC" : auc}

    return metrics_dict

def compute_metrics_ml(model, X, y, print = False):
    """
    Returns the accuracy, recall, precision, f1 score, and AUC
    """
    if hasattr(model, 'predict_proba'):
        probs = model.predict_proba(X)[:, 1] 
        predictions = model.predict(X)
    else:
        probs = model.predict(X)
        predictions = (probs >= 0.5)

    acc = accuracy_score(y, predictions)
    rec = recall_score(y, predictions)
    prec = precision_score(y, predictions, zero_division=0)
    f1 = f1_score(y, predictions)
    auc = roc_auc_score(y, probs)
    bal_acc = balanced_accuracy_score(y, predictions)

    tn, fp, fn, tp = confusion_matrix(y, predictions).ravel()
    spec = tn / (tn + fp)

    metrics_dict = {"Accuracy": acc, "Recall" : rec, "Precision" : prec, "F1 Score" : f1, "AUC" : auc, "Balanced Accuracy" : bal_acc, "Specificity" : spec}
    if print == True:
        display(pd.DataFrame(metrics_dict, index=[0]))

    return metrics_dict

def plot_roc_auc(model, device, dataloader, selector):
    model.eval()
    all_labels = []
    all_predictions = []
    all_probs = []

    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            
            probs = outputs.cpu()
            predictions = (probs >= 0.5).int().cpu().numpy()
            
            all_probs.extend(probs)
            all_predictions.extend(predictions)
            all_labels.extend(labels.cpu().numpy())

    auc = roc_auc_score(all_labels, all_probs)
    
    fpr, tpr, thresholds = roc_curve(all_labels, all_probs)
    plt.plot(fpr, tpr, label='ROC curve (area = %0.3f)' % auc)
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'ROC Curve using {selector} Features')
    plt.legend()
    plt.show()

def plot_auprc(model, device, dataloader, selector):
    model.eval()
    all_labels = []
    all_predictions = []
    all_probs = []

    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            
            probs = outputs.cpu()
            predictions = (probs >= 0.5).int().cpu().numpy()
            
            all_probs.extend(probs)
            all_predictions.extend(predictions)
            all_labels.extend(labels.cpu().numpy())

    auprc = average_precision_score(all_labels, all_probs)
    precision, recall, _ = precision_recall_curve(all_labels, all_probs)
    plt.plot(precision, recall, label=f'Precision-Recall curve (AUPRC = {auprc:.2f})')
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title(f'Precision-Recall Curve using {selector} Features')
    plt.legend(loc='lower left')
    plt.legend()
    plt.show()

def plot_confusion_matrix(model, device, dataloader, threshold):
    model.eval()
    all_labels = []
    all_predictions = []
    all_probs = []

    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            
            probs = outputs.cpu()
            predictions = (probs >= threshold).int().cpu().numpy()
            
            all_probs.extend(probs)
            all_predictions.extend(predictions)
            all_labels.extend(labels.cpu().numpy())

    cm = confusion_matrix(all_labels, all_predictions)
    cm_df = pd.DataFrame(cm,
                     index=["Actual Negative", "Actual Positive"],
                     columns=["Predicted Negative", "Predicted Positive"])

    print("\nConfusion Matrix:")
    display(cm_df)

def create_shap(selector):
    train_dataloader = selector_to_dataloader[selector][0]
    test_dataloader = selector_to_dataloader[selector][1]
    train_dataloader = torch.utils.data.DataLoader(train_dataloader.dataset, batch_size=64, shuffle=True)

    batch = next(iter(train_dataloader))
    inputs, _ = batch

    background = inputs
    background = background.to(device)

    batch = next(iter(train_dataloader))
    inputs, _ = batch

    samples = inputs
    samples = samples.to(device)

    recurrence_model = model_dict_new[selector].to(device)

    recurrence_model.eval()
    e = shap.GradientExplainer(recurrence_model, background)

    shap_values = e.shap_values(samples)
    shap_values = np.squeeze(shap_values)
    print(shap_values.shape)
    print(samples.shape)

    X_train_df = pd.concat([X_train_clin, X_train_exp_dict[selector]], axis=1)
    selected_feature_names = X_train_df.columns.tolist()

    shap.summary_plot(shap_values, features=samples, feature_names = np.array(selected_feature_names))

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)