Added majority voting evaluation

protac_degradation_predictor/optuna_utils.py

@@ -2,7 +2,7 @@ import os
 from typing import Literal, List, Tuple, Optional, Dict
 import logging
 
-from .pytorch_models import train_model
+from .pytorch_models import train_model, PROTAC_Model
 from .sklearn_models import (
     train_sklearn_model,
     suggest_random_forest,
@@ -11,6 +11,7 @@ from .sklearn_models import (
     suggest_gradient_boosting,
 )
 
+import torch
 import optuna
 from optuna.samplers import TPESampler
 import joblib
@@ -27,6 +28,56 @@ from sklearn.model_selection import (
 )
 import numpy as np
 import pytorch_lightning as pl
+from torchmetrics import (
+    Accuracy,
+    AUROC,
+    Precision,
+    Recall,
+    F1Score,
+)
+
+
+def get_dataframe_stats(
+        train_df = None,
+        val_df = None,
+        test_df = None,
+        active_label = 'Active',
+    ) -> Dict:
+    """ Get some statistics from the dataframes.
+    
+    Args:
+        train_df (pd.DataFrame): The training set.
+        val_df (pd.DataFrame): The validation set.
+        test_df (pd.DataFrame): The test set.
+    """
+    stats = {}
+    if train_df is not None:
+        stats['train_len'] = len(train_df)
+        stats['train_active_perc'] = train_df[active_label].sum() / len(train_df)
+        stats['train_inactive_perc'] = (len(train_df) - train_df[active_label].sum()) / len(train_df)
+    if val_df is not None:
+        stats['val_len'] = len(val_df)
+        stats['val_active_perc'] = val_df[active_label].sum() / len(val_df)
+        stats['val_inactive_perc'] = (len(val_df) - val_df[active_label].sum()) / len(val_df)
+    if test_df is not None:
+        stats['test_len'] = len(test_df)
+        stats['test_active_perc'] = test_df[active_label].sum() / len(test_df)
+        stats['test_inactive_perc'] = (len(test_df) - test_df[active_label].sum()) / len(test_df)
+    if train_df is not None and val_df is not None:
+        leaking_uniprot = list(set(train_df['Uniprot']).intersection(set(val_df['Uniprot'])))
+        leaking_smiles = list(set(train_df['Smiles']).intersection(set(val_df['Smiles'])))
+        stats['num_leaking_uniprot_train_val'] = len(leaking_uniprot)
+        stats['num_leaking_smiles_train_val'] = len(leaking_smiles)
+        stats['perc_leaking_uniprot_train_val'] = len(train_df[train_df['Uniprot'].isin(leaking_uniprot)]) / len(train_df)
+        stats['perc_leaking_smiles_train_val'] = len(train_df[train_df['Smiles'].isin(leaking_smiles)]) / len(train_df)
+    if train_df is not None and test_df is not None:
+        leaking_uniprot = list(set(train_df['Uniprot']).intersection(set(test_df['Uniprot'])))
+        leaking_smiles = list(set(train_df['Smiles']).intersection(set(test_df['Smiles'])))
+        stats['num_leaking_uniprot_train_test'] = len(leaking_uniprot)
+        stats['num_leaking_smiles_train_test'] = len(leaking_smiles)
+        stats['perc_leaking_uniprot_train_test'] = len(train_df[train_df['Uniprot'].isin(leaking_uniprot)]) / len(train_df)
+        stats['perc_leaking_smiles_train_test'] = len(train_df[train_df['Smiles'].isin(leaking_smiles)]) / len(train_df)
+    return stats
 
 
 def pytorch_model_objective(
@@ -77,15 +128,15 @@ def pytorch_model_objective(
     X = train_val_df.copy().drop(columns=active_label)
     y = train_val_df[active_label].tolist()
     report = []
+    val_preds = []
+    test_preds = []
     for k, (train_index, val_index) in enumerate(kf.split(X, y, groups)):
         logging.info(f'Fold {k + 1}/{kf.get_n_splits()}')
         # Get the train and val sets
         train_df = train_val_df.iloc[train_index]
         val_df = train_val_df.iloc[val_index]
 
-        # Check for data leakage and get some statistics
-        leaking_uniprot = list(set(train_df['Uniprot']).intersection(set(val_df['Uniprot'])))
-        leaking_smiles = list(set(train_df['Smiles']).intersection(set(val_df['Smiles'])))
+        # Get some statistics from the dataframes
         stats = {
             'model_type': 'Pytorch',
             'fold': k,
@@ -93,22 +144,15 @@ def pytorch_model_objective(
             'val_len': len(val_df),
             'train_perc': len(train_df) / len(train_val_df),
             'val_perc': len(val_df) / len(train_val_df),
-            'train_active_perc': train_df[active_label].sum() / len(train_df),
-            'train_inactive_perc': (len(train_df) - train_df[active_label].sum()) / len(train_df),
-            'val_active_perc': val_df[active_label].sum() / len(val_df),
-            'val_inactive_perc': (len(val_df) - val_df[active_label].sum()) / len(val_df),
-            'num_leaking_uniprot': len(leaking_uniprot),
-            'num_leaking_smiles': len(leaking_smiles),
-            'train_leaking_uniprot_perc': len(train_df[train_df['Uniprot'].isin(leaking_uniprot)]) / len(train_df),
-            'train_leaking_smiles_perc': len(train_df[train_df['Smiles'].isin(leaking_smiles)]) / len(train_df),
         }
+        stats.update(get_dataframe_stats(train_df, val_df, test_df, active_label))
         if groups is not None:
             stats['train_unique_groups'] = len(np.unique(groups[train_index]))
             stats['val_unique_groups'] = len(np.unique(groups[val_index]))
 
         # At each fold, train and evaluate the Pytorch model
         # Train the model with the current set of hyperparameters
-        _, trainer, metrics = train_model(
+        ret = train_model(
             protein2embedding=protein2embedding,
             cell2embedding=cell2embedding,
             smiles2fp=smiles2fp,
@@ -127,22 +171,47 @@ def pytorch_model_objective(
             use_logger=False,
             fast_dev_run=fast_dev_run,
             active_label=active_label,
+            return_predictions=True,
             disabled_embeddings=disabled_embeddings,
         )
+        if test_df is not None:
+            _, trainer, metrics, val_pred, test_pred = ret
+            test_preds.append(test_pred)
+            logging.info(f'Test predictions: {test_pred}')
+        else:
+            _, trainer, metrics, val_pred = ret
         train_metrics = {m: v.item() for m, v in trainer.callback_metrics.items() if 'train' in m}
         stats.update(metrics)
         stats.update(train_metrics)
         report.append(stats.copy())
+        val_preds.append(val_pred)
+
+    # Save the report in the trial
+    trial.set_user_attr('report', report)
+
+    # Get the majority vote for the test predictions
+    if test_df is not None:
+        # Get the majority vote for the test predictions
+        test_preds = torch.stack(test_preds)
+        test_preds, _ = torch.mode(test_preds, dim=0)
+        y = torch.tensor(test_df[active_label].tolist())
+        # Measure the test accuracy and ROC AUC
+        majority_vote_metrics = {
+            'test_acc': Accuracy(task='binary')(test_preds, y).item(),
+            'test_roc_auc': AUROC(task='binary')(test_preds, y).item(),
+            'test_precision': Precision(task='binary')(test_preds, y).item(),
+            'test_recall': Recall(task='binary')(test_preds, y).item(),
+            'test_f1': F1Score(task='binary')(test_preds, y).item(),
+        }
+        majority_vote_metrics.update(get_dataframe_stats(train_df, val_df, test_df, active_label))
+        trial.set_user_attr('majority_vote_metrics', majority_vote_metrics)
+        logging.info(f'Majority vote metrics: {majority_vote_metrics}')
 
     # Get the average validation accuracy and ROC AUC accross the folds
-    val_acc = np.mean([r['val_acc'] for r in report])
     val_roc_auc = np.mean([r['val_roc_auc'] for r in report])
 
-    # Save the report in the trial
-    trial.set_user_attr('report', report)
-    
     # Optuna aims to minimize the pytorch_model_objective
-    return - val_acc - val_roc_auc
+    return - val_roc_auc
 
 
 def hyperparameter_tuning_and_training(
@@ -162,6 +231,7 @@ def hyperparameter_tuning_and_training(
         active_label: str = 'Active',
         max_epochs: int = 100,
         study_filename: Optional[str] = None,
+        force_study: bool = False,
 ) -> tuple:
     """ Hyperparameter tuning and training of a PROTAC model.
     
@@ -181,10 +251,11 @@ def hyperparameter_tuning_and_training(
     pl.seed_everything(42)
 
     # Define the search space
-    hidden_dim_options = [256, 512, 768]
-    batch_size_options = [8, 16, 32]
+    hidden_dim_options = [32, 64, 128, 256, 512, 768]
+    batch_size_options = [4, 8, 16, 32, 64, 128]
     learning_rate_options = (1e-5, 1e-3) # min and max values for loguniform distribution
     smote_k_neighbors_options = list(range(3, 16))
+    dropout_options = (0.1, 0.9)
 
     # Set the verbosity of Optuna
     optuna.logging.set_verbosity(optuna.logging.WARNING)
@@ -193,13 +264,13 @@ def hyperparameter_tuning_and_training(
     study = optuna.create_study(direction='minimize', sampler=sampler)
 
     study_loaded = False
-    if study_filename:
+    if study_filename and not force_study:
         if os.path.exists(study_filename):
             study = joblib.load(study_filename)
             study_loaded = True
             logging.info(f'Loaded study from {study_filename}')
 
-    if not study_loaded:
+    if not study_loaded or force_study:
         study.optimize(
             lambda trial: pytorch_model_objective(
                 trial=trial,
@@ -214,6 +285,7 @@ def hyperparameter_tuning_and_training(
                 batch_size_options=batch_size_options,
                 learning_rate_options=learning_rate_options,
                 smote_k_neighbors_options=smote_k_neighbors_options,
+                dropout_options=dropout_options,
                 fast_dev_run=fast_dev_run,
                 active_label=active_label,
                 max_epochs=max_epochs,
@@ -228,9 +300,11 @@ def hyperparameter_tuning_and_training(
 
     # Retrain N models with the best hyperparameters (measure model uncertainty)
     test_report = []
+    test_preds = []
+    dfs_stats = get_dataframe_stats(train_val_df, test_df=test_df, active_label=active_label)
     for i in range(n_models_for_test):
         pl.seed_everything(42 + i + 1)
-        _, trainer, metrics = train_model(
+        _, trainer, metrics, test_pred = train_model(
             protein2embedding=protein2embedding,
             cell2embedding=cell2embedding,
             smiles2fp=smiles2fp,
@@ -245,29 +319,52 @@ def hyperparameter_tuning_and_training(
             logger_name=f'{logger_name}_best_model_n{i}',
             enable_checkpointing=True,
             checkpoint_model_name=f'best_model_n{i}_{split_type}',
+            return_predictions=True,
             **study.best_params,
         )
         # Rename the keys in the metrics dictionary
         metrics = {k.replace('val_', 'test_'): v for k, v in metrics.items()}
         metrics['model_type'] = 'Pytorch'
         metrics['test_model_id'] = i
-        metrics['test_len'] = len(test_df)
-        metrics['test_active_perc'] = test_df[active_label].sum() / len(test_df)
-        metrics['test_inactive_perc'] = (len(test_df) - test_df[active_label].sum()) / len(test_df)
+        metrics.update(dfs_stats)
 
         # Add the training metrics        
-        train_metrics = {m.replace('train_', 'train_val_'): v.item() for m, v in trainer.callback_metrics.items() if 'train' in m}
+        train_metrics = {m: v.item() for m, v in trainer.callback_metrics.items() if 'train' in m}
         logging.info(f'Training metrics: {train_metrics}')
         logging.info(f'Training trainer.logged_metrics: {trainer.logged_metrics}')
         logging.info(f'Training trainer.callback_metrics: {trainer.callback_metrics}')
         
         metrics.update(train_metrics)
-
         test_report.append(metrics.copy())
+        test_preds.append(test_pred)
     test_report = pd.DataFrame(test_report)
 
+    # Get the majority vote for the test predictions
+    test_preds = torch.stack(test_preds)
+    test_preds, _ = torch.mode(test_preds, dim=0)
+    y = torch.tensor(test_df[active_label].tolist())
+    # Measure the test accuracy and ROC AUC
+    majority_vote_metrics = {
+        'cv_models': False,
+        'test_acc': Accuracy(task='binary')(test_preds, y).item(),
+        'test_roc_auc': AUROC(task='binary')(test_preds, y).item(),
+        'test_precision': Precision(task='binary')(test_preds, y).item(),
+        'test_recall': Recall(task='binary')(test_preds, y).item(),
+        'test_f1': F1Score(task='binary')(test_preds, y).item(),
+    }
+    majority_vote_metrics.update(get_dataframe_stats(train_val_df, test_df=test_df, active_label=active_label))
+    majority_vote_metrics_cv = study.best_trial.user_attrs['majority_vote_metrics']
+    majority_vote_metrics_cv['cv_models'] = True
+    majority_vote_report = pd.DataFrame([
+        majority_vote_metrics,
+        majority_vote_metrics_cv,
+    ])
+    majority_vote_report['model_type'] = 'Pytorch'
+    majority_vote_report['split_type'] = split_type
+
     # Ablation study: disable embeddings at a time
     ablation_report = []
+    dfs_stats = get_dataframe_stats(train_val_df, test_df=test_df, active_label=active_label)
     for disabled_embeddings in [['e3'], ['poi'], ['cell'], ['smiles'], ['e3', 'cell'], ['poi', 'e3', 'cell']]:
         logging.info('-' * 100)
         logging.info(f'Ablation study with disabled embeddings: {disabled_embeddings}')
@@ -291,9 +388,10 @@ def hyperparameter_tuning_and_training(
         metrics = {k.replace('val_', 'test_'): v for k, v in metrics.items()}
         metrics['disabled_embeddings'] = 'disabled ' + ' '.join(disabled_embeddings)
         metrics['model_type'] = 'Pytorch'
+        metrics.update(dfs_stats)
 
         # Add the training metrics        
-        train_metrics = {m.replace('train_', 'train_val_'): v.item() for m, v in trainer.callback_metrics.items() if 'train' in m}
+        train_metrics = {m: v.item() for m, v in trainer.callback_metrics.items() if 'train' in m}
         metrics.update(train_metrics)
 
         ablation_report.append(metrics.copy())
@@ -304,7 +402,14 @@ def hyperparameter_tuning_and_training(
         report['split_type'] = split_type
 
     # Return the reports
-    return cv_report, hparam_report, test_report, ablation_report
+    ret = {
+        'cv_report': cv_report,
+        'hparam_report': hparam_report,
+        'test_report': test_report,
+        'ablation_report': ablation_report,
+        'majority_vote_report': majority_vote_report,
+    }
+    return ret
 
 
 def sklearn_model_objective(

protac_degradation_predictor/pytorch_models.py

@@ -315,26 +315,6 @@ class PROTAC_Model(pl.LightningModule):
         e3_emb = batch['e3_emb']
         cell_emb = batch['cell_emb']
         smiles_emb = batch['smiles_emb']
-
-        if self.apply_scaling:
-            if self.join_embeddings == 'beginning':
-                embeddings = np.hstack([
-                    np.array(smiles_emb.tolist()),
-                    np.array(poi_emb.tolist()),
-                    np.array(e3_emb.tolist()),
-                    np.array(cell_emb.tolist()),
-                ])
-                embeddings = self.scalers.transform(embeddings)
-                smiles_emb = embeddings[:, :self.smiles_emb_dim]
-                poi_emb = embeddings[:, self.smiles_emb_dim:self.smiles_emb_dim+self.poi_emb_dim]
-                e3_emb = embeddings[:, self.smiles_emb_dim+self.poi_emb_dim:self.smiles_emb_dim+2*self.poi_emb_dim]
-                cell_emb = embeddings[:, -self.cell_emb_dim:]
-            else:
-                poi_emb = self.scalers['Uniprot'].transform(poi_emb)
-                e3_emb = self.scalers['E3 Ligase Uniprot'].transform(e3_emb)
-                cell_emb = self.scalers['Cell Line Identifier'].transform(cell_emb)
-                smiles_emb = self.scalers['Smiles'].transform(smiles_emb)
-
         y_hat = self.forward(poi_emb, e3_emb, cell_emb, smiles_emb)
         return torch.sigmoid(y_hat)
 
@@ -416,6 +396,7 @@ def train_model(
         enable_checkpointing: bool = False,
         checkpoint_model_name: str = 'protac',
         disabled_embeddings: List[str] = [],
+        return_predictions: bool = False,
 ) -> tuple:
     """ Train a PROTAC model using the given datasets and hyperparameters.
     
@@ -540,12 +521,19 @@ def train_model(
         warnings.simplefilter("ignore")
         trainer.fit(model)
     metrics = trainer.validate(model, verbose=False)[0]
-
-    # Add train metrics to metrics
-
+    # Add test metrics to metrics
     if test_df is not None:
         test_metrics = trainer.test(model, verbose=False)[0]
         metrics.update(test_metrics)
+    if return_predictions:
+        val_dl = DataLoader(val_ds, batch_size=batch_size, shuffle=False)
+        val_pred = trainer.predict(model, val_dl)
+        val_pred = torch.concat(trainer.predict(model, val_dl)).squeeze()
+        if test_df is not None:
+            test_dl = DataLoader(test_ds, batch_size=batch_size, shuffle=False)
+            test_pred = torch.concat(trainer.predict(model, test_dl)).squeeze()
+            return model, trainer, metrics, val_pred, test_pred
+        return model, trainer, metrics, val_pred
     return model, trainer, metrics
 
 

src/run_experiments.py

@@ -3,6 +3,7 @@ import sys
 from collections import defaultdict
 import warnings
 import logging
+from typing import Literal
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
@@ -214,6 +215,8 @@ def main(
     cv_n_splits: int = 5,
     max_epochs: int = 100,
     run_sklearn: bool = False,
+    force_study: bool = False,
+    experiments: str | Literal['all', 'random', 'e3_ligase', 'tanimoto', 'uniprot'] = 'all',
 ):
     """ Train a PROTAC model using the given datasets and hyperparameters.
     
@@ -244,10 +247,15 @@ def main(
     ## Get the test sets
     test_indeces = {}
     active_df = protac_df[protac_df[active_col].notna()].copy()
-    test_indeces['random'] = get_random_split_indices(active_df, test_split)
-    test_indeces['e3_ligase'] = get_e3_ligase_split_indices(active_df)
-    test_indeces['tanimoto'] = get_tanimoto_split_indices(active_df, active_col, test_split)
-    test_indeces['uniprot'] = get_target_split_indices(active_df, active_col, test_split)
+    
+    if experiments == 'random' or experiments == 'all':
+        test_indeces['random'] = get_random_split_indices(active_df, test_split)
+    if experiments == 'uniprot' or experiments == 'all':
+        test_indeces['uniprot'] = get_target_split_indices(active_df, active_col, test_split)
+    if experiments == 'e3_ligase' or experiments == 'all':
+        test_indeces['e3_ligase'] = get_e3_ligase_split_indices(active_df)
+    if experiments == 'tanimoto' or experiments == 'all':
+        test_indeces['tanimoto'] = get_tanimoto_split_indices(active_df, active_col, test_split)
     
     # Make directory ../reports if it does not exist
     if not os.path.exists('../reports'):
@@ -296,22 +304,18 @@ def main(
             logger_name=f'logs_{experiment_name}',
             active_label=active_col,
             study_filename=f'../reports/study_{experiment_name}.pkl',
+            force_study=force_study,
         )
-        cv_report, hparam_report, test_report, ablation_report = optuna_reports
 
         # Save the reports to file
-        for report, filename in zip([cv_report, hparam_report, test_report, ablation_report], ['cv_train', 'hparams', 'test', 'ablation']):
-            report.to_csv(f'../reports/report_{filename}_{experiment_name}.csv', index=False)
-
-        reports['cv'].append(cv_report.copy())
-        reports['hparam'].append(hparam_report.copy())
-        reports['test'].append(test_report.copy())
-        reports['ablation'].append(ablation_report.copy())
+        for report_name, report in optuna_reports.items():
+            report.to_csv(f'../reports/report_{report_name}_{experiment_name}.csv', index=False)
+            reports[report_name].append(report.copy())
     
     # Save the reports to file after concatenating them
-    for key, report in reports.items():
+    for report_name, report in reports.items():
         report = pd.concat(report)
-        report.to_csv(f'../reports/report_{key}_{active_name}_test_split_{test_split}.csv', index=False)
+        report.to_csv(f'../reports/report_{report_name}_{active_name}_test_split_{test_split}.csv', index=False)