Tiny Lab BioML

From Fingerprints to Message Passing: Benchmarking Three Approaches to Molecular Toxicity Prediction on Tox21

Jay Chung — Fri, 24 Jul 2026 07:00:00 GMT

Summary

In this post, I compare three approaches to predicting chemical toxicity directly from molecular structure, using the Tox21 dataset (~7,800 compounds tested across 12 nuclear receptor and stress-response assays). The first is a Random Forest trained on molecular fingerprints, a standard cheminformatics baseline. The second is a Graph Neural Network (GINE) built from scratch that learns directly from atoms and bonds. The third is Chemprop, a widely-used graph neural network package, combined with RDKit 2D molecular descriptors. All three models are evaluated using a scaffold split, which keeps structurally similar molecules out of both the training and test sets for a more realistic test of generalization to new chemical structures. All three models achieve similar performance that aligns with previous literature benchmarks. I also walk through a real numerical pitfall with RDKit descriptors that can silently corrupt training, and use t-SNE to check whether the learned graph representations capture meaningful chemical structure.

Introduction

Tox21 is a multi-task binary classification dataset for predicting the toxicity of chemical compounds. The dataset contains 12 different toxicity tasks, each corresponding to a specific biological target or pathway. The goal is to predict whether a given compound is toxic or non-toxic for each of these tasks. Each compound is represented by its SMILES string (Chemprop additionally uses RDKit 2D descriptors as extra molecule-level features, described in Model 3). The dataset is highly imbalanced, with some tasks having very few positive examples (3-15% positive rates). There are also many missing labels. This provides a challenge for training machine learning models. We will tackle these problems by using specific metrics and loss function as described later.

Three models to train:

Random Forest (RF) with ECFP2048 (Morgan fingerprints) as input features: this is considered as a baseline model with classical machine learning approach trained on molecular fingerprints. The RF model will be well-tuned with hyperparameter search and cross-validation.
GINE (Graph Isomorphism Network with Edge features), trained directly on atom- and bond-level graph features: GIN or GINE is a type of Graph Neural Network (GNN) with message-passing similar to a Graph Convolution Network (GCN). Unlike GCN, which combines messages by averaging neighbors’ embeddings, GINE sums the neighboring messages instead, which better preserves information about the number of neighbors a node has. This allows for better differentiation of certain non-isomorphic graphs than GCN. For an introduction into GIN vs GCN, see the TeachOpenCADD GNN tutorial. The GINE model will be built from scratch using PyG’s GINEConv - a variant of GIN that includes both node and edge features during the message-passing steps (a typical GIN only calculates node features). This allows for learning of molecular properties from both atom and bond features from the compounds. I will be using the PyTorch Geometric (PyG) library to build the GINE model from scratch.
Chemprop (D-MPNN), trained on the molecular graph plus RDKit 2D (whole-molecule) descriptors as extra features: Chemprop is a high-level library for building and training D-MPNN (Directed Message Passing Neural Network) models to predict molecular properties. D-MPNN passes messages along directed bonds and explicitly excludes the reverse-direction message when aggregating each edge’s update, avoiding the “tottering” problem (a signal bouncing immediately back to the atom it came from). D-MPNN has been shown to perform well on larger industrial proprietary datasets. Chemprop combines D-MPNN with a feed-forward neural network (FFNN) that takes in additional molecular features, such as RDKit 2D features, to improve the model’s performance. Chemprop abstracts away the details of building and training D-MPNN models, making it easier to use. I will be using Chemprop’s Python-first API that is built on PyTorch Lightning to train the D-MPNN model.

To prevent data leakage, a scaffold split will be performed on the Tox21 dataset, which ensures that structurally related molecules do not end up in different partitions (train/validation/test). This is important because if structurally similar molecules are present in both training and test sets, the model may perform well on the test set simply because it has seen similar molecules during training, rather than learning to generalize to unseen molecules.

The final performance of each model will be summarized using AUROC, AUPRC and binary cross-entropy (BCE) loss. It is important to note that the RF model will be well-tuned with hyperparameter search and cross-validation, while the GINE and Chemprop models will be trained once with pre-selected hyperparameters. The learned graph-level embeddings from the GINE and Chemprop models will be visualized using t-SNE to see if the embeddings capture any meaningful chemical and predictive properties in the data.

Data processing

Loading required libraries and setting up the environment:

Code

Querying ChEMBL for Target-Specific Bioactivity and Compound Data in Python

Jay Chung — Sun, 05 Jul 2026 07:00:00 GMT

chEMBL target query

Small, reusable Python utilities for pulling target-specific bioactivity and compound data from ChEMBL given a UniProt ID, and evaluating those compounds against Lipinski’s rule of five — adapted from TeachOpenCADD T001: Compound data acquisition (ChEMBL) and T002: Molecular filtering (Ro5).

Synthetic data representing pIC50 vs. molecular weight for EGFR inhibitors

Given a UniProt accession (e.g. P00533 for EGFR), it:

Looks up the matching ChEMBL target(s)
Fetches IC50 bioactivity data (human, exact measurements, binding assays)
Filters to nM units and deduplicates by compound
Fetches canonical SMILES for the resulting compounds
Merges bioactivity + compound data and computes pIC50
Optionally adds Lipinski’s rule of five (Ro5) property columns

… and returns a tidy pandas.DataFrame with columns:

| molecule_chembl_id | IC50 | units | smiles | pIC50 |

Adding Ro5 properties appends: molecular_weight, n_hba, n_hbd, logp, ro5_fulfilled.

Installation

git clone https://github.com/jaychung10010/chembl-target-query.git
cd chembl-target-query
pip install -r requirements.txt

Usage

A UniProt accession can map to multiple ChEMBL target entries (single protein, protein family, chimeric construct, protein-protein interaction, etc.), so the workflow is split into two steps — inspect, then extract. This avoids blocking input() prompts, so it works in any environment (plain scripts, Jupyter, agent-driven IDEs like Antigravity/Cursor where stdin isn’t interactive).

Code

Diagnosing Overfitting in a PPI Transformer: Hyperparameter Search and What Actually Helped

Jay Chung — Mon, 22 Jun 2026 07:00:00 GMT

Summary

A hyperparameter search was conducted to diagnose and mitigate overfitting in a transformer-based model for predicting Protein-Protein Interactions (PPIs). The search explored variations in dropout rates, latent dimensions, and weight decay. The results showed that increasing dropout and reducing model size mildly improved generalization, while weight decay had a less pronounced effect. The best configuration achieved an AUROC improvement of ~0.02 on the leakage-reduced test set compared to the baseline. These results suggest that while hyperparameter tuning can provide modest gains, the model may have reached a possible performance ceiling given the current architecture and mean-pooled ESM2 embeddings. Further improvements are likely to require architectural changes or data augmentation strategies.

The problem

In my previous post, I built a transformer-based model to predict PPIs trained from the HuRI dataset. The model achieved an AUROC of 0.75 and accuracy of 0.67 on leakage-reduced test data, which is comparable to the performance of existing models in the literature (Reim et al. 2025). However, looking at the training and validation AUROC curves, it was clear that the model was overfitting to the training data quite early on. The training AUROC quickly reached 0.9 within the first few epochs, while the validation AUROC plateaued around 0.74 and even started to decline after epoch 10. This indicated that the model was learning patterns specific to the training data that did not generalize well to unseen data.

Training and validation AUROC curves showing overfitting.

Hypothesis

I hypothesized that the overfitting was due to the model being too complex for the amount of training data available, and that tuning the hyperparameters could help mitigate this issue. Specifically, I wanted to investigate how different hyperparameters such as latent dimension (model size), dropout rate, and weight decay affected the model’s performance and generalization ability.

Experimental design

Let’s look at the model architecture again:

Model architecture diagram.

I defined 8 sets of hyperparameter configurations to test, including variations in dropout rates for the projection, attention, and MLP layers, as well as different latent dimensions and weight decay values. The configurations were designed to systematically explore the effects of increasing dropout, reducing model size, and applying stronger regularization through weight decay.

Code

Predicting Protein-Protein Interactions with ESM2 Embeddings and Transformer-based Interaction Model

Jay Chung — Mon, 01 Jun 2026 07:00:00 GMT

Summary

Using the PyTorch framework, I built a transformer-based model to predict Protein-Protein Interactions (PPIs) trained from the HuRI dataset. The workflow takes a pair of protein sequences as input, extracts ESM2 embeddings, and predicts the probability of interaction. The model achieved an AUPRC of 0.75 and accuracy of 0.67 on leakage-reduced test data, which is comparable to the performance of existing models in the literature. The transformer model significantly outperformed a baseline random forest model, which achieved an AUPRC of 0.62 and accuracy of 0.59 on the same test set.

Research Impact

Why is it important to build a PPI predictor?

Drug discovery and disease mechanism: PPI is fundamental to how proteins carry out their functions in the cell. Identifying interaction partners of a drug target can help us understand the target’s biological role and its involvement in disease pathways. Many diseases involve dysregulated PPIs, often by disrupting normal interactions or creating new ones through functional mutations. Predicting how these mutations affect PPIs can help us better understand disease mechanisms and identify potential therapeutic targets. In addition, a PPI predictor model can be used for off-target screening of protein biologics, which is crucial for drug safety and efficacy.
Proteome-wide interaction networks: Current large-scale PPI screen like HuRI contains roughly ~9,000 proteins, while the human proteome contains ~20,000 proteins, many of which are understudied. A trained model can be used to explore these unknown PPIs, potentially expanding our knowledge of interactome without the costs and time of experimental screens. A generalized PPI prediction model can be applied to other species as well, or to cross-species interaction like virus-host protein interaction, helping us understand the biology of non-model organisms and infectious diseases.
Contributing to the PPI deep learning field: By utilizing state-of-the-art models like ESM2 and transformer architectures, I aim to see if building a model with comparable performance to existing models is possible. This contributes to the broader question of how much interaction-relevant information is already encoded in ESM2 embeddings, and whether further optimization of the model architecture can lead to significant performance gains. I will also address the issue of data leakage in PPI prediction, which is a common pitfall that can lead to overestimated performance. By carefully designing the training and evaluation strategy to avoid data leakage, I hope to provide a more realistic assessment of the model’s predictive power.

Background

Last year, a paper by Reim et al. (2025) compared various deep learning architectures for PPI prediction, including models with attention mechanisms and incorporation of ESM2 protein language model. They found that regardless of the model architecture or the hyperparameter space, all models seemed to plateau at an accuracy of 0.65. They concluded that any performance gains observed are attributable to the use of ESM2 embeddings, rather than the model architecture itself. The paper also provided some interesting observations about the ESM2 model usage; for example, models profited from smaller embeddings (t33), and that per-token embeddings (sequence information preserved) did not yield better performance than per-protein embeddings (mean embeddings across sequence).

Inspired by this paper, I want to see if I can build a model with comparable performance, albeit with a different model design. I decided to use the HuRI PPI dataset for training and evaluation, rather than the HIPPIE dataset used in the paper. The HuRI dataset came from a Yeast-2-Hybrid (Y2H) screen that contains ~9,000 proteins and ~64,000 experimental validated interactions. The HIPPIE dataset is a literature-curated dataset that is much larger but may also bias toward more well-studied proteins. Due to the difference in dataset size and quality, my result may not be directly comparable to the Reim et al. paper, but I will compare the transformer performance with a baseline random forest model to see if the transformer architecture provides any performance gain over a simpler model.

Many papers in the PPI prediction field have pointed out the issue of data leakage, which is when the same or similar proteins appear in both training and test sets, leading to overestimated performance. To address this, I will implement a leakage-reduced splitting strategy that ensures that proteins in the test set are not similar or present in the training set. This will provide a more realistic assessment of the model’s generalization ability.

Model architecture

The model contains four main components:

ESM2 mean embedding extractor: a pre-trained ESM2 model that takes amino acid sequences of two proteins and generates their per-protein mean embeddings (EmbA and EmbB). I will use the t33_650M_UR50D variant of ESM2, which has been shown to perform well for PPI prediction in the Reim et al. paper.
Protein projector: a feedforward neural network that takes the embeddings and projects them to a lower-dimensional space. Both proteins pass through identical weight, which forces a common geometry before interaction modeling.
Interaction transformer: a transformer encoder that takes the projected embeddings of the two proteins and models their interactions. Because the protein pair is only 2 tokens, self-attention can be considered as cross-attention, and each protein attends to the other and itself. The transformer will have multiple layers of self-attention and feedforward networks to capture complex interaction patterns between the two proteins.
PPI classifier: the MLP concatenates the output of the interaction transformer with two symmetric features of |EmbA - EmbB| and EmbA * EmbB (absolute difference and products), which gave the model more flexibility to learn the interaction-relevant information. The MLP will output a probability score indicating the likelihood of interaction between the two proteins.

One caveat for the HuRI data is that the Y2H experiment is prone to false positives, which might exacerbate the issue of overfitting, meaning that the model will predict well on training data but could not generalized to unseen test data. To deal with this, these features were added to the model design:

Model training features that deals with overfitting: dropouts, L2 regularization (AdamW + weight decay), adaptive learning rate, and early stopping.
Label smoothing function to make sure that the target label is not overly confident on positive or negative outcomes.
Symmetric augmentation function that switch protein pairs randomly before training.

Another consideration is the size of the model vs. training sample size. With around 50k of HuRI training samples, this is on the lower side of training a 9.3 million parameters transformer model. 9.3M params / 50k samples = 180 params per sample. Normally it would be great to get this ratio to < 100 to prevent overfitting, but given that the model is not learning everything from scratch - ESM2 did the initial heavy lifting - this setting might be acceptable.

1. Data processing

First, I will download the positive interactions HuRI data from TDC, and then process and clean up the data. A protein similarity-aware split will be performed to obtain the train, validation and test data. HuRI data only contains positive interactions, so I will also perform negative sampling to generate negative interaction samples for each of the split. This is done after the splitting to avoid further data leakage, as the negative samples are generated based on the proteins present in each split.

Data download

Code

Evaluating Data Leakage in Protein Binding Affinity Prediction

Jay Chung — Wed, 25 Feb 2026 08:00:00 GMT

Research Impact

Data leakage is a critical issue in machine learning, especially in the context of protein-protein interaction (PPI) prediction. Data leakage occurs when information from the training data is inadvertently used in the testing phase, allowing the model to “remember” specific examples rather than learning generalizable patterns. To address this issue, this post compares two data splitting strategies, regular C3 split and strict C3 split, to evaluate the extent of data leakage when using pre-trained protein large language models (pLLMs) to predict peptide-MHC2 binding. The findings showed that while the neural network model generalized well on predicting binding peptides of seen MHC2 alleles, it struggled to generalize to unseen MHC2 alleles. This highlights the importance of using strict data splitting strategies to ensure that models are evaluated on truly unseen data, which is crucial for developing robust and generalizable models in protein binding affinity prediction.

Introduction

I’ve recently come across this paper about data leakage in PPI prediction models using pLLMs, and it got me thinking about how data leakage might affect other types of protein interaction predictions, such as peptide-MHC2 binding. In this project, I decided to explore this issue by comparing two different data splitting strategies: a regular C3 split and a strict C3 split. The regular C3 split allows for some overlap between the training and testing sets, while the strict C3 split ensures that there is no overlap at all. By evaluating the performance of a neural network model trained on embeddings extracted from a pre-trained pLLM (ESM2) under these two splitting strategies, I aimed to understand the extent of data leakage and its impact on model generalization in the context of peptide-MHC2 binding prediction.

Key Steps

Splitting Strategies: Peptide or MHC sequences were first clustered independently using two different approaches. In both cases, the goal was to ensure that similar protein sequences were not present in both the training and testing sets. Peptides were clustered based on a 3-gram sequence similarity approach, while MHC pseudo-sequences were clustered based on their BLOSUM62 amino acid similarity. After that, the cluster classes were split into training and testing sets using two different strategies largely based on the C3 approach:
- Regular C3 split: This split ensures that the test set contains pairs and . This is a less strict C3 split, as it allows for some overlap in the individual components (peptides and MHCs) between the training and testing sets, but not in the interactions.
- Strict C3 split: This split ensures that for every pair in the test set, both and classes are entirely absent from the training set. This is the “Double-Cold” strategy, which is more stringent and ensures that the model is evaluated on completely unseen peptides and MHCs, thus providing a more accurate assessment of the model’s generalization capabilities. A down side of this split is that it results in a smaller training and testing set. To ensure a fair comparison between the two splits, the sample number from the “regular C3 split” was downsampled to match the sample number from the “strict C3 split”.
Embeddings Extraction and Model Training: Similar to my previous post, we will extract ESM2 embeddings for both peptides and MHC2 pseudo-sequences. The embeddings will be extracted in a way that retains the 2D structure, which is crucial for applying convolutional layers in the model. After extracting the embeddings, we will concatenate the peptide and MHC2 embeddings along the sequence length dimension, retaining the sequence information and the 2D structure. We will then define a model architecture that includes two 1D convolutional layers followed by a self-attention layer and five dense layers. The model will be trained separately on the datasets generated from the regular C3 split and the strict C3 split, allowing us to compare the performance of the model under both splitting strategies.
Comparing Regular vs. Strict C3 split: After training the model on both datasets, we will evaluate its performance using appropriate metrics such as score, root mean squared error (RMSE), and loss.

1. Data Splitting Strategies

Loading required libraries:

Code

Peptide-MHCII Binding Affinity Prediction with Convolutional Neural Networks and Protein Large Language Model

Jay Chung — Tue, 03 Feb 2026 08:00:00 GMT

Research Impact

To accurately predict peptide-MHCII binding affinity, I developed a 1D Convolutional Neural Network (CNN) combined with Multi-Layer Perceptron (MLP) model using protein embeddings from the ESM2 large language model.

Comparing with a simpler MLP-only model, the CNN + MLP model demonstrated over 12% improvement in peptide-MHCII binding prediction accuracy, achieving an score of 0.62 for all MHC alleles, and close to 0.8 for specific alleles. Saliency Map and In Silico Mutagenesis analyses revealed key amino acid positions that significantly influence binding affinity predictions, providing insights into the underlying biological interactions.

This exercise showcases the potential of combining Large Language Models with machine learning to tackle complex biological challenges. These could range from predicting protein-protein interactions and drug-target binding to forecasting drug responses from gene expression data.

Schematic overview of the CNN + MLP pipeline for predicting peptide-MHC binding affinity (figure generated with the assistance of Gemini).

Introduction

In my last post, I explored the use of protein large language model (LLM) embeddings from ESM2 for predicting peptide-MHCII binding affinity using a Multi-Layer Perceptron (MLP). While the MLP performed reasonably well, I believe that incorporating Convolutional Neural Networks (CNNs) could further enhance the model’s ability to capture local sequence patterns in the protein embeddings.

The protein sequence embeddings generated by ESM2 are typically 2-dimensional arrays, where one dimension represents the amino acid sequence and the other represents the embedding features. In my previous approach, I flattened these embeddings into 1-dimensional vectors by taking the mean before feeding them into the MLP. However, this flattening process may lead to the loss of important spatial relationships between amino acids in the sequence.

In this post, I will implement a 1D CNN to process the full embeddings from ESM2 without pre-flattening. I will apply 2 layers of 1D convolution followed by average pooling to capture local patterns in the sequence embeddings. After the CNN layers, I will add fully connected MLP layers to perform the final prediction task. I will compare the performance of this CNN + MLP model with a similar MLP-only model to evaluate the impact of incorporating convolutional layers.

Finally, to understand which amino acid positions contribute most to the binding affinity predictions, I will perform Saliency Map analysis on both the peptide and MHCII sequences, using gradients calculated from the model outputs with respect to the input embeddings. I will also perform an In Silico Mutagenesis (ISM) analysis to see how disruptive single amino acid changes affect the predicted binding affinity.

Key Steps

I will not go through all the steps in detail, as they are similar to my previous post. Instead, I will highlight the key steps involved in this approach:

Embeddings Extraction: We will have to modify the embeddings extraction code to retain the 2D structure of the embeddings. We need to input the max sequence length to ensure consistent input sizes for the model. After that, we will concatenate the peptide and MHC2 embeddings along the sequence length dimension, retaining the sequence information and the 2D structure. This allows for the saliency map analysis later on.
Model Architecture: We will define a new model architecture that includes 1D convolutional layers followed by dense layers. Using the functional API in TensorFlow/Keras, we will create a model that takes the 2D embeddings and peptide/MHC sequence lengths as inputs, applies convolutional layers, and then concatenates them to pass it through dense layers to make the final prediction.
Training and Evaluation: We will train the new model and evaluate its performance against a similar MLP-only model, to see if the CNN layers improve binding affinity prediction.
Saliency Map and In Silico Mutagenesis Analysis: We will compute saliency maps to identify important amino acid positions in the peptide and MHC2 sequences that influence binding affinity predictions. We will also perform ISM to assess the impact of single amino acid mutations on the predicted affinity.

1. Embeddings Extraction

First, let’s see what the input data looks like:

Code

Predicting Peptide-MHC Class II Binding with ESM2 and Neural Networks

Jay Chung — Sun, 18 Jan 2026 08:00:00 GMT

In this post, I will build a neural network to predict the binding affinity between peptides and MHC Class II alleles, utilizing embeddings extracted from the Evolutionary Scale Modeling 2 (ESM2) protein large language model (LLM).

Background

Why is this important?

Understanding Immune Response: MHC Class II molecules present peptide fragments to CD4+ T helper cells, which are crucial for initiating and regulating adaptive immune responses. Understanding which peptides bind strongly to MHC II allows researchers to identify potential T cell epitopes.
Vaccine Design: Accurate prediction of MHC II-binding peptides can guide the design of subunit vaccines. By selecting peptides likely to bind to a broad range of MHC II alleles (covering genetic diversity in a population), vaccines can be designed to elicit robust T cell responses.
Autoimmune Diseases: Inappropriate presentation of self-peptides by MHC II molecules can contribute to autoimmune diseases. Predicting these interactions helps in understanding disease mechanisms and developing therapeutic interventions.
Allergy Research: Similarly, understanding how allergens bind to MHC II can shed light on allergic reactions and aid in treatment development.
Personalized Medicine: With advancements in sequencing, predicting an individual’s MHC II genotype and the peptides they present could lead to personalized immunotherapies.

Analysis Plan

Download and split the TDC MHC2_IEDB-Jensen dataset.
Perform exploratory data analysis (EDA).
Extract and merge LLM embeddings for peptides and MHC sequences.
Setup TensorFlow datasets, build a Multilayer Perceptron (MLP) model, with defined loss function and optimizer.
Train the model and evaluate metrics.
Assess prediction accuracy on test data.

Summary Figure

Schematic overview of the Deep Neural Network pipeline for predicting peptide-MHC binding affinity (Figure generated with the assistance of Gemini).

1. Data Selection and Splitting

We will use datasets from Therapeutics Data Commons (TDC), a large-scale data repository for machine learning projects.

Code

Machine learning for drug sensitivity prediction (Part 3): Training an Elastic Net gene expression model to predict Erlotinib sensitivity

Jay Chung — Fri, 26 Dec 2025 08:00:00 GMT

In my previous posts (part 1 and part 2), I explored whether using autoencoders (AE) to compress high-dimensional multi-omics data could improve the performance of drug sensitivity prediction models. In my specific context, I showed that using AE did not necessarily improve model performance compared to using a common feature selection method.

In this post, I will use the best-performing model from my previous analyses (elastic net using features selected by Pearson correlation), and apply it to all available DepMap data to generate predictions for cell lines without measured drug sensitivity. I will then see if these predictions provide additional useful insights beyond the measured data.

I will demonstrate how to use R’s Caret package to train an elastic net model with hyperparameter tuning via cross-validation.

Aim 1: Apply the elastic net model to all available DepMap data to generate Erlotinib sensitivity predictions

Data location (pre-processed for this post): https://zenodo.org/records/17970100

Load data and libraries

library(tidyverse)
library(data.table)
library(ggpubr)
library(caret)
library(doParallel)

# Load data from the files directory
load("files/CCLE_24Q2_GE_match_sample_info.RData")
load("files/PRISM_24Q2_compound_screen_match_sample_info.RData")
cmpd_dat <- read_csv("files/Repurposing_Public_24Q2_Extended_Primary_Compound_List.csv")
load("files/sample_info_match_biomarkers.RData")

Preprocessing

First, we remove low variance genes and any cell lines containing NAs.

# Remove low variance genes
ge_var <- apply(ccle_ge_match_sam, 2, var, na.rm = TRUE) # calculate variance for each gene
ge_dat_hv <- ccle_ge_match_sam[, which(ge_var > quantile(ge_var, 0.2))] # select high variance genes

# Remove any cell lines (rows) that contains NAs
ge_dat_hv <- ge_dat_hv |>
    na.omit() |>
    as.data.frame()

# Match data
match_idx <- match(rownames(ge_dat_hv), sample_info$StrippedCellLineName)
sample_info_filt <- sample_info[match_idx, ]
prism_dat_match_sam_filt <- prism_dat_match_sam[match_idx, ]

# Verify alignment
all(rownames(ge_dat_hv) == sample_info_filt$StrippedCellLineName)

[1] TRUE

Now we prepare the drug response data for Erlotinib.

cmpd_id <- cmpd_dat |>
    filter(Drug.Name == "ERLOTINIB") |>
    pull(IDs)
drug_LFC <- prism_dat_match_sam_filt |> pull(cmpd_id) # Erlotinib LFC

# Remove NAs in drug response
valid_idx <- which(!is.na(drug_LFC))
ge_dat_final <- ge_dat_hv[valid_idx, ]
drug_LFC_final <- drug_LFC[valid_idx]
sample_info_final <- sample_info_filt[valid_idx, ]

# Verify alignment
all(rownames(ge_dat_final) == sample_info_final$StrippedCellLineName)

[1] TRUE

# Check drug response distribution
ggplot(data.frame(Erlotinib_LFC = drug_LFC_final), aes(x = Erlotinib_LFC)) +
    geom_density(fill = "purple", alpha = 0.7) +
    theme_classic(base_size = 20) +
    labs(title = "Density of Erlotinib LFC", x = "Erlotinib LFC", y = "Density")

Feature Selection

We select the top predictors correlating with the drug response.

# Select top predictors correlating with drug response, Pearson correlation > 0.1
cor_values <- apply(ge_dat_final, 2, function(x) cor(x, drug_LFC_final, method = "pearson"))
cor_values <- abs(cor_values)
ggplot(data.frame(Correlation = cor_values), aes(x = Correlation)) +
    geom_density(fill = "darkgreen", alpha = 0.7) +
    theme_classic(base_size = 20) +
    labs(title = "Density of Gene-Drug Correlation Values", x = "Absolute Pearson Correlation", y = "Density")

top_predictors <- names(cor_values)[which(cor_values > 0.1)]
ge_dat_final <- ge_dat_final |> select(all_of(top_predictors))

Elastic Net Hyperparameter Tuning

We perform hyperparameter tuning using caret.

Note: The training step is computationally intensive and is not evaluated here. We load the pre-trained model for subsequent analysis.

# Elastic net hyperparameter tuning with caret
detectCores()
cl <- makePSOCKcluster(22)
registerDoParallel(cl)

glmnetGrid <- expand.grid(
    .lambda = round(seq(0.005, 5, length = 30), 3),
    .alpha = round(seq(0, 0.1, length = 10), 3)
)

set.seed(524)
enet_full_model <- train(ge_dat_final, drug_LFC_final,
    method = "glmnet",
    tuneGrid = glmnetGrid,
    preProcess = c("center", "scale"),
    metric = "Rsquared",
    trControl = trainControl(method = "repeatedcv", number = 10, repeats = 10)
)
stopCluster(cl)

# Save the final model (Already saved in files/)
# save(enet_full_model, file = "files/enet_full_model_erlotinib.RData")

Model Analysis

Let’s examine the tuning results and the best hyperparameters.

# Load the pre-trained model
load("files/enet_full_model_erlotinib.RData")

# View hyperparameter tuning results
enet_full_model$results |> arrange(desc(Rsquared)) |> head(20) |> gt::gt()

alpha	lambda	RMSE	Rsquared	MAE	RMSESD	RsquaredSD	MAESD
0.011	1.038	0.7246115	0.3234841	0.5678419	0.06244775	0.09552666	0.05504541
0.011	1.211	0.7243195	0.3231371	0.5668798	0.06237722	0.09620741	0.05451755
0.011	0.866	0.7259216	0.3227335	0.5697183	0.06251476	0.09492452	0.05557950
0.011	1.383	0.7246391	0.3221409	0.5665405	0.06232075	0.09689575	0.05382773
0.011	1.555	0.7254462	0.3205662	0.5668465	0.06223333	0.09753524	0.05317660
0.011	0.694	0.7287441	0.3204541	0.5731421	0.06255456	0.09446370	0.05612583
0.011	1.727	0.7265929	0.3185779	0.5676531	0.06210627	0.09802932	0.05253257
0.022	0.694	0.7276708	0.3185370	0.5695071	0.06120405	0.09481046	0.05397744
0.022	0.522	0.7298551	0.3180041	0.5726883	0.06144171	0.09315651	0.05466701
0.022	0.866	0.7277250	0.3168255	0.5686165	0.06109334	0.09620112	0.05314454
0.011	1.900	0.7279475	0.3163917	0.5686306	0.06199278	0.09836348	0.05195792
0.011	0.522	0.7339054	0.3158314	0.5786809	0.06262747	0.09381218	0.05677975
0.033	0.522	0.7303634	0.3142904	0.5715719	0.06047743	0.09405803	0.05348058
0.011	2.072	0.7294330	0.3141275	0.5697778	0.06194162	0.09867437	0.05146113
0.022	1.038	0.7289549	0.3140533	0.5692525	0.06118321	0.09735362	0.05229037
0.022	0.349	0.7362917	0.3135250	0.5805138	0.06171108	0.09216744	0.05541844
0.033	0.694	0.7296913	0.3129244	0.5700855	0.06080031	0.09625073	0.05273691
0.011	2.244	0.7309768	0.3118847	0.5711072	0.06195602	0.09901043	0.05102660
0.044	0.522	0.7306761	0.3116728	0.5711300	0.06048624	0.09528180	0.05289960
0.033	0.349	0.7355536	0.3116308	0.5781206	0.06063449	0.09156073	0.05410146

The best hyperparameters for the model were alpha = 0.011 and lambda = 1.038. This indicates that the model prefers a nearly ridge regression approach (alpha close to 0) with moderate regularization (lambda value). This makes sense given that the predictors were pre-selected based on correlation with the response, so we want to retain most predictors while controlling for multicollinearity.

# Plot tuning results
ggplot(enet_full_model) +
    theme_classic(base_size = 20) +
    labs(title = "Elastic Net tuning")

A major advantage of a linear model like elastic net is its interpretability. The coefficients correspond to how important each gene is for predicting Erlotinib sensitivity.

final_coef <- as.matrix(coef(enet_full_model$finalModel, s = enet_full_model$bestTune$lambda))[-1, ]
final_coef_df <- data.frame(
    Gene = names(final_coef),
    Coefficient = final_coef
)

# Plot a bar plot for the top 20 sensitivity and top 20 resistance predictor genes
final_coef_sorted <- final_coef[order(final_coef, decreasing = TRUE)]
top20_pos_coef <- head(final_coef_sorted[final_coef_sorted > 0], 20)
top20_neg_coef <- tail(final_coef_sorted[final_coef_sorted < 0], 20)

top_coef_df <- data.frame(
    Gene = c(names(top20_pos_coef), names(top20_neg_coef)),
    Coefficient = c(top20_pos_coef, top20_neg_coef),
    Direction = c(rep("Resistance predictors", length(top20_pos_coef)), rep("Sensitivity predictors", length(top20_neg_coef)))
)

ggplot(top_coef_df, aes(x = reorder(Gene, Coefficient), y = Coefficient, fill = Direction)) +
    geom_bar(stat = "identity", alpha = 0.7) +
    theme_classic(base_size = 18) +
    labs(
        title = "Top 40 predictor genes",
        x = "Gene",
        y = "Coefficient"
    ) +
    scale_fill_manual(values = c("Sensitivity predictors" = "darkgreen", "Resistance predictors" = "darkred")) +
    theme(
        legend.position = "top",
        axis.text.x = element_text(angle = 45, hjust = 1),
        axis.text = element_text(size = 10)
    )

Prediction on All Cell Lines

We now use the trained model to predict Erlotinib sensitivity for all CCLE cell lines with available RNA expression data.

# Predict on all CCLE cell lines with available RNA expression
all_ccle_ge <- ccle_ge_match_sam |> as.data.frame()
all_ccle_ge <- all_ccle_ge |>
    select(all_of(top_predictors)) |>
    na.omit() # 1517 cell lines with complete GE data

enet_pred_all_ccle <- predict(enet_full_model, newdata = all_ccle_ge)

# Save predicted Erlotinib LFC for all CCLE cell lines
predicted_erlotinib_df <- data.frame(
    Cell_Line = rownames(all_ccle_ge),
    Predicted_Erlotinib_LFC = enet_pred_all_ccle
)
# fwrite(predicted_erlotinib_df, "files/predicted_erlotinib_LFC_all_CCLE_cell_lines.csv", sep = ",", row.names = FALSE, quote = FALSE)

Aim 2: Evaluate whether the predicted Erlotinib sensitivity provides additional insights beyond the measured data

Comparison of Actual vs. Predicted

# Load actual and predicted Erlotinib LFC data
actual_erlotinib_df <- data.frame(
    Cell_Line = rownames(ge_dat_final),
    Actual_Erlotinib_LFC = drug_LFC_final
)

# If we didn't run the prediction above, we could load it:
# predicted_erlotinib_df <- fread("files/predicted_erlotinib_LFC_all_CCLE_cell_lines.csv", data.table = FALSE)

# Merge actual and predicted data
erlotinib_compare_df <- merge(actual_erlotinib_df, predicted_erlotinib_df, by = "Cell_Line", all = TRUE)

# Add cancer lineage information
erlotinib_compare_df <- erlotinib_compare_df |> left_join(sample_info |> select(StrippedCellLineName, OncotreeLineage),
    by = c("Cell_Line" = "StrippedCellLineName")
)

# Add EGFR mutation status
# Note: Using tryCatch or existence check in case file is missing in user environment
if (file.exists("files/CCLE_24Q2_HOTMUT_match_sample_info.RData")) {
    load("files/CCLE_24Q2_HOTMUT_match_sample_info.RData")
    egfr_mutation_status <- hotmut_dat_match_sam |>
        rownames_to_column(var = "Cell_Line") |>
        select(Cell_Line, EGFR) |>
        mutate(EGFR_Mutation_Status = recode(as.factor(EGFR), `0` = "Wildtype", `1` = "Heterozygous Mutant", `2` = "Homozygous Mutant")) |>
        select(Cell_Line, EGFR_Mutation_Status)
    erlotinib_compare_df <- erlotinib_compare_df |> left_join(egfr_mutation_status, by = "Cell_Line")
} else {
    warning("files/CCLE_24Q2_HOTMUT_match_sample_info.RData not found. Skipping EGFR analysis.")
}

# Plot actual vs predicted Erlotinib LFC for cell lines with actual data
plot_df <- erlotinib_compare_df |> filter(!is.na(Actual_Erlotinib_LFC))
ggplot(plot_df, aes(x = Actual_Erlotinib_LFC, y = Predicted_Erlotinib_LFC)) +
    geom_point(color = "steelblue", size = 3, alpha = 0.7) +
    theme_classic(base_size = 20) +
    theme(plot.title = element_text(size = 20)) +
    labs(
        title = "Elastic Net prediction of Erlotinib LFC",
        x = "Actual Erlotinib LFC",
        y = "Predicted Erlotinib LFC"
    ) +
    stat_cor(method = "pearson", label.x.npc = "left", label.y.npc = "top", size = 6, color = "darkred") +
    geom_smooth(method = "lm", color = "darkred", se = TRUE)

Note: This prediction performance result is evaluated on the training data, so it is expected to be better than on independent test data (overfitting). It is important to always evaluate model performance on independent test data. From my previous post, the prediction performance on independent test data looks like this (not bad but not as good as training data):

Cancer Lineage Distribution

We compare the cancer lineages covered by the actual measured data versus the predicted data.

actual_lineages <- erlotinib_compare_df |>
    filter(!is.na(Actual_Erlotinib_LFC)) |>
    select(Cell_Line, OncotreeLineage) |>
    distinct()
predicted_lineages <- erlotinib_compare_df |>
    select(Cell_Line, OncotreeLineage) |>
    distinct()

actual_lineage_counts <- actual_lineages |>
    group_by(OncotreeLineage) |>
    summarise(Count = n()) |>
    mutate(Type = "Actual data")
predicted_lineage_counts <- predicted_lineages |>
    group_by(OncotreeLineage) |>
    summarise(Count = n()) |>
    mutate(Type = "Predicted data")
lineage_counts_df <- rbind(actual_lineage_counts, predicted_lineage_counts) |> filter(OncotreeLineage != "")

# Make missing lineage in actual data have count 0
all_lineages <- unique(lineage_counts_df$OncotreeLineage)
for (lineage in all_lineages) {
    if (!(lineage %in% actual_lineage_counts$OncotreeLineage)) {
        lineage_counts_df <- rbind(lineage_counts_df, data.frame(OncotreeLineage = lineage, Count = 0, Type = "Actual data"))
    }
}

ggplot(lineage_counts_df, aes(x = reorder(OncotreeLineage, -Count), y = Count, fill = Type)) +
    geom_bar(stat = "identity", position = position_dodge()) +
    theme_classic(base_size = 18) +
    labs(
        title = "Cancer Lineage Distribution",
        x = "Cancer Lineage",
        y = "Number of Cell Lines"
    ) +
    geom_text(aes(label = Count, color = Type), position = position_dodge(width = 0.9), vjust = -0.5, size = 3) +
    scale_fill_manual(values = c("Actual data" = "steelblue", "Predicted data" = "darkorange")) +
    scale_color_manual(values = c("Actual data" = "steelblue", "Predicted data" = "darkorange")) +
    theme(
        legend.position = "top",
        axis.text.x = element_text(angle = 45, hjust = 1),
        axis.text = element_text(size = 10)
    )

You can see here that the predicted data covers many more cancer lineages than the actual data, which may improve the power of downstream analyses.

Predicted Sensitivity by Lineage

Next, let’s see if the predicted Erlotinib LFC can better capture differences across cancer lineages.

# Plot actual vs predicted Erlotinib LFC boxplot faceted by LFC type and stratified by OncotreeLineage
plot_df <- erlotinib_compare_df |>
    select(Cell_Line, Actual_Erlotinib_LFC, Predicted_Erlotinib_LFC, OncotreeLineage) |>
    pivot_longer(
        cols = c(Actual_Erlotinib_LFC, Predicted_Erlotinib_LFC),
        names_to = "LFC_Type",
        values_to = "Erlotinib_LFC"
    ) |>
    na.omit()

lineage_counts <- plot_df |>
    group_by(OncotreeLineage) |>
    summarise(Count = n()) |>
    filter(Count >= 10)
plot_df <- plot_df |> filter(OncotreeLineage %in% lineage_counts$OncotreeLineage)

lineage_sort_by_predicted_LFC <- plot_df |>
    filter(LFC_Type == "Predicted_Erlotinib_LFC") |>
    group_by(OncotreeLineage) |>
    summarise(Median_Predicted_LFC = median(Erlotinib_LFC)) |>
    arrange(Median_Predicted_LFC) |>
    pull(OncotreeLineage)

plot_df$OncotreeLineage <- factor(plot_df$OncotreeLineage, levels = lineage_sort_by_predicted_LFC)

ggplot(plot_df, aes(x = OncotreeLineage, y = Erlotinib_LFC, color = LFC_Type)) +
    geom_jitter(width = 0.2, size = 2, alpha = 0.5) +
    geom_boxplot(outliers = FALSE, fill = "grey99", alpha = 0.3) +
    facet_wrap(~LFC_Type,
        nrow = 2, scales = "free_y",
        labeller = as_labeller(c(Actual_Erlotinib_LFC = "Actual Erlotinib LFC", Predicted_Erlotinib_LFC = "Predicted Erlotinib LFC"))
    ) +
    theme_classic(base_size = 18) +
    scale_color_manual(
        values = c(Actual_Erlotinib_LFC = "steelblue", Predicted_Erlotinib_LFC = "darkorange"),
        labels = c("Actual Erlotinib LFC", "Predicted Erlotinib LFC")
    ) +
    labs(
        title = "Erlotinib LFC by Cancer Lineage",
        x = "Cancer Lineage",
        y = "Erlotinib LFC"
    ) +
    geom_hline(yintercept = 0, linetype = "dashed", color = "grey30") +
    theme(
        axis.text.x = element_text(angle = 60, hjust = 1),
        legend.position = "none"
    )

Result: Predicted data discovered lineage sensitivity for prostate, cervix, biliary tract; and showed that blood lineages are likely not sensitive.

Predicted Sensitivity by EGFR Mutation Status

Finally, let’s see whether the predicted Erlotinib LFC can capture known biological associations, such as sensitivity in EGFR mutant cell lines.

if (exists("egfr_mutation_status")) {
    # Compare predicted vs actual Erlotinib LFC between EGFR mutant and wildtype cell lines
    plot_df <- erlotinib_compare_df |>
        select(Cell_Line, Actual_Erlotinib_LFC, Predicted_Erlotinib_LFC, EGFR_Mutation_Status) |>
        pivot_longer(
            cols = c(Actual_Erlotinib_LFC, Predicted_Erlotinib_LFC),
            names_to = "LFC_Type",
            values_to = "Erlotinib_LFC"
        ) |>
        na.omit()

    # Plot boxplot: EGFR mutant vs wildtype facet by actual vs predicted with test p-value
    ggplot(plot_df, aes(x = EGFR_Mutation_Status, y = Erlotinib_LFC, color = EGFR_Mutation_Status)) +
        geom_boxplot(alpha = 1, outliers = FALSE) +
        geom_jitter(width = 0.2, size = 2, alpha = 0.5) +
        facet_wrap(~LFC_Type, nrow = 1, labeller = as_labeller(c(Actual_Erlotinib_LFC = "Actual Erlotinib LFC", Predicted_Erlotinib_LFC = "Predicted Erlotinib LFC"))) +
        theme_classic(base_size = 18) +
        labs(
            title = "Erlotinib LFC by EGFR Mutation Status",
            x = "EGFR Hotspot Mutation Status",
            y = "Erlotinib LFC"
        ) +
        scale_color_manual(values = c("Wildtype" = "darkgreen", "Heterozygous Mutant" = "darkorange", "Homozygous Mutant" = "darkred")) +
        theme(
            legend.position = "none",
            axis.text.x = element_text(angle = 45, hjust = 1)
        ) +
        stat_compare_means(aes(group = EGFR_Mutation_Status), label.y.npc = 0.9, size = 5)
} else {
    print("EGFR mutation data not loaded.")
}

It seems that even in the predicted data, there is only a marginal difference in Erlotinib LFC between EGFR mutant and wildtype cell lines. This could be due to the fact that EGFR mutation instances are relatively rare in all of DepMap cell lines.

Summary

Here I demonstrated how to train an elastic net model using R’s Caret package with hyperparameter tuning via cross-validation from end to end. I showed the top 20 sensitivity and resistance predictor genes identified by the final model. I then applied the final model to all available DepMap cell lines to generate Erlotinib sensitivity predictions.

Finally, I showed that the predicted data can provide additional insights beyond the measured data, such as covering more cancer lineages. However, some known biological associations (e.g. EGFR mutation status) may still be difficult to capture due to data limitations.

Overall, this framework can be useful for generating drug sensitivity predictions for cell lines without measured data, which can aid in drug repurposing and precision oncology efforts.

References

DepMap Portal: https://depmap.org/portal/
Caret Package Documentation: https://topepo.github.io/caret/
Applied Predictive Modeling by Kuhn and Johnson: https://www.springer.com/gp/book/9781461468486

Training a Random Forest model with Scikit-Learn on DepMap data

Jay Chung — Mon, 15 Dec 2025 08:00:00 GMT

Introduction

In this specific blog post, we will explore the Cancer Dependency Map (DepMap) dataset to understand the relationship between genetic dependencies (CRISPR knockout effects) and gene expression levels. We will use Python’s pandas ecosystem for data loading, exploratory data analysis, and visualization. We will also use scikit-learn’s functions to build a Random Forest predictor pipeline that performs preprocessing, imputes missing values, and tunes hyperparameters. Finally, we will evaluate the model’s performance and extract genes that are most important for predicting gene dependency.

Data Loading

First, we import the necessary libraries and load the datasets. We are using three main datasets:

Sample Info: Metadata about the cell lines (lineage, disease subtype, etc.).
Chronos Data: CRISPR knockout scores representing gene dependency (lower score = higher dependency).
Gene Expression (GE) Data: RNA-seq expression levels for various genes.

Data can be downloaded from here (modified from DepMap 24Q2): https://zenodo.org/records/17970100

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import os

# Define file paths (relative to the ML_practice directory)
# Here I have preprocessed the data for this demo so the cell lines (rows) are aligned
sample_info_path = "ML_practice/sample_info.csv"
chronos_path = "ML_practice/chronos_dat.csv"
ge_path = "ML_practice/ccle_ge.csv"

# Load the datasets
# Sample metadata containing cell line information
sample_data = pd.read_csv(sample_info_path)

# Chronos scores: Gene effect scores from CRISPR knockout screens
chronos_dat = pd.read_csv(chronos_path)

# Gene Expression data: mRNA expression levels log2(TPM + 1)
ge_dat = pd.read_csv(ge_path)

Exploratory Data Analysis

Let’s inspect the structure and basic statistics of our datasets.

Sample Information

Checking the metadata to understand the cell lines we are working with.

# Display dataset information: columns, non-null counts, and data types
print("Sample Data Info:")
sample_data.info()

Sample Data Info:

RangeIndex: 1959 entries, 0 to 1958
Data columns (total 9 columns):
 #   Column                   Non-Null Count  Dtype 
---  ------                   --------------  ----- 
 0   ModelID                  1959 non-null   object
 1   StrippedCellLineName     1959 non-null   object
 2   CCLEName                 1902 non-null   object
 3   OncotreeLineage          1954 non-null   object
 4   OncotreeSubtype          1959 non-null   object
 5   OncotreePrimaryDisease   1959 non-null   object
 6   LegacySubSubtype         831 non-null    object
 7   LegacyMolecularSubtype   151 non-null    object
 8   PatientMolecularSubtype  138 non-null    object
dtypes: object(9)
memory usage: 137.9+ KB

Chronos (Dependency) Data

This dataset contains dependency scores for key genes like SMARCA2, SOX10, and KRAS.

# Display dataset information for Chronos data
print("\nChronos Data Info:")
chronos_dat.info()


Chronos Data Info:

RangeIndex: 1959 entries, 0 to 1958
Data columns (total 4 columns):
 #   Column     Non-Null Count  Dtype  
---  ------     --------------  -----  
 0   cell_line  1959 non-null   object 
 1   SMARCA2    1150 non-null   float64
 2   SOX10      1150 non-null   float64
 3   KRAS       1150 non-null   float64
dtypes: float64(3), object(1)
memory usage: 61.3+ KB

Gene Expression Data

This dataset provides expression levels for a wide range of genes across the cell lines.

# Display dataset information for Gene Expression data
print("\nGene Expression Data Info:")
ge_dat.info()
print("\nFirst 5 rows of Gene Expression Data:")
display(ge_dat.head())


Gene Expression Data Info:

RangeIndex: 1959 entries, 0 to 1958
Columns: 19153 entries, TSPAN6 to CDR1
dtypes: float64(19153)
memory usage: 286.3 MB

First 5 rows of Gene Expression Data:

	TSPAN6	TNMD	DPM1	SCYL3	FIRRM	FGR	CFH	FUCA2	GCLC	NFYA	...	SPDYE11	H3C2	H3C3	DUS4L-BCAP29	C8orf44-SGK3	ELOA3BP	NPBWR1	ELOA3P	CDR1
0	5.183487	0.000000	7.497612	2.107688	4.217231	0.042644	0.903038	5.722193	4.676944	3.720278	...	0.056584	1.137504	0.000000	1.794936	0.201634	0.000000	0.028569	0.214125	0.014355
1	0.176323	0.000000	5.702103	1.238787	3.119356	4.141596	0.163499	4.134221	4.111866	2.347666	...	0.000000	1.952334	0.238787	1.516015	0.000000	0.000000	0.028569	0.000000	0.000000
2	5.309976	0.084064	7.846117	1.875780	3.894333	0.000000	0.056584	6.666615	4.738768	3.984589	...	0.014355	0.918386	0.000000	1.655352	0.000000	0.028569	0.000000	0.028569	0.000000
3	2.176323	0.000000	5.454505	2.480265	3.921246	0.887525	4.958843	3.949535	4.877253	4.829850	...	0.014355	0.536053	0.505891	2.643856	0.097611	0.000000	0.000000	0.000000	0.000000
4	2.451541	0.000000	5.884842	2.927896	5.299391	0.201634	5.759156	4.150560	5.531069	5.029895	...	0.000000	0.895303	1.350497	2.709291	0.678072	0.000000	0.000000	0.000000	0.000000

5 rows × 19153 columns

Data Distribution and Correlations

We can look at the overall distribution of the dependency data using a scatter matrix. This helps us spot potential correlations or clusters between different gene dependencies.

from pandas.plotting import scatter_matrix

# Create a scatter matrix to visualize pair-wise relationships in the Chronos dataset
# This includes histograms on the diagonal and scatter plots on off-diagonals
scatter_matrix(chronos_dat)
plt.show()

Scatter matrix of Chronos dependency scores for selected genes.

Let’s verify the correlations numerically. We drop the ‘cell_line’ column as it is categorical.

# Calculate the correlation matrix for the numeric columns in Chronos data
chronos_corr = chronos_dat.drop("cell_line", axis=1).corr()
print("\nCorrelation Matrix (Chronos Data):")
print(chronos_corr)


Correlation Matrix (Chronos Data):
          SMARCA2     SOX10      KRAS
SMARCA2  1.000000 -0.032416  0.005902
SOX10   -0.032416  1.000000 -0.052455
KRAS     0.005902 -0.052455  1.000000

Visualization: Dependency vs. Expression

A key question in cancer biology is whether gene expression predicts dependency. For example, if a cell line highly expresses SOX10, is it more dependent on SOX10 for survival (lower Chronos score)?

Let’s visualize the relationship between SOX10 dependency (Chronos score) and SOX10 gene expression.

# Plot SOX10 Chronos score (x-axis) vs SOX10 Gene Expression (y-axis)
# Note: Lower Chronos score means higher dependency.
plt.figure(figsize=(8, 6))
plt.scatter(chronos_dat["SOX10"], ge_dat["SOX10"], alpha=0.6)
plt.xlabel("SOX10 Chronos Score (Dependency)")
plt.ylabel("SOX10 Gene Expression")
plt.title("SOX10: Dependency vs Expression")
plt.grid(True, linestyle='--', alpha=0.5)
plt.show()

Scatter plot showing the relationship between SOX10 Gene Expression and SOX10 Dependency (Chronos Score).

Finally, let’s calculate the Pearson correlation coefficient between these two variables to quantify the relationship.

# Calculate Pearson correlation between SOX10 dependency and expression
sox10_corr = chronos_dat["SOX10"].corr(ge_dat["SOX10"])
print(f"Pearson Correlation between SOX10 Chronos and Gene Expression: {sox10_corr:.4f}")

Pearson Correlation between SOX10 Chronos and Gene Expression: -0.8491

This correlation value suggests a strong relationship between SOX10 gene expression and dependency.

Machine Learning

Now we will build a machine learning model to predict the SOX10 dependency score based on the gene expression profile of the cell lines.

Data Preprocessing

Before training, we need to prepare our data.

Alignment: Ensure the cell lines in the dependency dataset match those in the gene expression dataset.
Merging: Combined the target variable (SOX10 Chronos score) with the feature set (Gene Expression).
Imputation: Dealing with missing values using K-Nearest Neighbors (KNN) imputation.

# Align and merge data
# The rows are already pre-aligned by cell line index
# Adding target variable to the dataframe for alignment
ge_dat = pd.read_csv(ge_path)
ge_dat["SOX10_chronos"] = chronos_dat["SOX10"]

# Remove rows where the target (SOX10_chronos) is NaN, as we can't train/test on them
ge_dat_clean = ge_dat.dropna(subset=["SOX10_chronos"])

# Separate features (X) and target (y)
X = ge_dat_clean.drop("SOX10_chronos", axis=1)
y = ge_dat_clean["SOX10_chronos"]

print(f"Data shape after removing missing targets: {X.shape}")

Data shape after removing missing targets: (1150, 19153)

We will handle missing values in the features (Gene Expression) using KNN Imputation within our modeling pipeline.

Model Training

We will use a Random Forest Regressor to predict the dependency score.

Data Splitting

First, we split the data into training and testing sets (80% train, 20% test).

from sklearn.model_selection import train_test_split

# Split data into training and testing set
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

print(f"Training samples: {X_train.shape[0]}")
print(f"Testing samples: {X_test.shape[0]}")

Training samples: 920
Testing samples: 230

Pipeline and Hyperparameter Tuning

We create a pipeline that: 1. Imputes missing values using KNNImputer. 2. Scales features using StandardScaler (optional for RF but good practice). 3. Trains a RandomForestRegressor.

We will use RandomizedSearchCV to find the best hyperparameters (e.g., max_features for the Random Forest).

from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import RandomizedSearchCV
from scipy.stats import randint
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.impute import KNNImputer

# Create a pipeline
pipeline = Pipeline([
    ("imputer", KNNImputer(n_neighbors=10)),
    ("scaler", StandardScaler()),
    ("rf", RandomForestRegressor(random_state=42))
])

# Define hyperparameter search space
param_dist = {
    'rf__max_features': randint(low=100, high=5000), # Number of features to consider at each split
    'rf__n_estimators': randint(low=50, high=200) # Number of trees in the forest
}

# Initialize RandomizedSearchCV
# n_iter=5 to keep run time reasonable for this demo
rnd_search = RandomizedSearchCV(
    pipeline, 
    param_distributions=param_dist, 
    n_iter=5, # Number of parameter settings that are sampled
    cv=3, # Number of folds in cross-validation
    random_state=42, # Random state for reproducibility
    scoring="neg_root_mean_squared_error", # Scoring metric
    n_jobs=-1 # Use all available CPU cores
)

# Fit the model
print("Training Random Forest model...")
rnd_search.fit(X_train, y_train)

print(f"Best RMSE: {-rnd_search.best_score_:.4f}")
print(f"Best Parameters: {rnd_search.best_params_}")

# Get the best model
final_model = rnd_search.best_estimator_

Training Random Forest model...

Best RMSE: 0.2097
Best Parameters: {'rf__max_features': 3019, 'rf__n_estimators': 180}

Model Evaluation

Now we evaluate the model’s performance on the unseen test set and investigate which genes are most important for predicting SOX10 dependency.

Feature Importance

Which genes’ expression levels are most predictive of SOX10 dependency?

# Extract feature importances
rf_model = final_model.named_steps['rf']
importances = rf_model.feature_importances_
feature_names = X.columns

# Create a dataframe for visualization
feat_importances = pd.Series(importances, index=feature_names)
top_features = feat_importances.nlargest(20)

# Plot top 20 features
plt.figure(figsize=(10, 6))
top_features.plot(kind='bar')
plt.title("Top 20 Feature Importances")
plt.ylabel("Importance")
plt.xlabel("Gene")
plt.xticks(rotation=45, ha='right')
plt.show()

Top 20 most important features (genes) for predicting SOX10 dependency.

As expected, SOX10 gene expression is the most important feature for predicting SOX10 dependency.

Prediction Performance

We calculate the Root Mean Squared Error (RMSE) on the test set and visualize the predictions vs actual values.

from sklearn.metrics import mean_squared_error

# Predict on test set
test_predictions = final_model.predict(X_test)

# Calculate RMSE
test_mse = mean_squared_error(y_test, test_predictions)
test_rmse = np.sqrt(test_mse)

print(f"Test Set RMSE: {test_rmse:.4f}")

# Plot Predicted vs Expected
plt.figure(figsize=(8, 6))
plt.scatter(y_test, test_predictions, alpha=0.6)
plt.plot([y.min(), y.max()], [y.min(), y.max()], 'k--', lw=2) # Identity line
plt.xlabel("Expected Score (True)")
plt.ylabel("Predicted Score")
plt.title("Test Set: Predicted vs Expected")
plt.grid(True, linestyle='--', alpha=0.5)
plt.show()

Test Set RMSE: 0.2139

Predicted vs Expected SOX10 Chronos Scores on Test Set.

We can see that the model predicts a more bimodal distribution of scores than the actual data, although it does somewhat differentiate between high and low dependency populations. To improve the model, we could try to tune the hyperparameters further, or try different feature selection methods.

References

Machine learning for drug sensitivity prediction (Part 2): will deep features improve accuracy?

Jay Chung — Mon, 08 Dec 2025 08:00:00 GMT

In Part 1, I explored using autoencoders (AE) to compress high-dimensional DepMap omics data into deep features. I visualized these features using UMAP and observed that they preserved biological signals related to cell lineage.

In this post (Part 2), I will evaluate whether these deep features actually improve drug sensitivity prediction accuracy compared to using either the full RNA expression data or a feature-selected subset of RNA expression data. The motivation is that deep features may capture complex, non-linear relationships in the data while reducing dimensionality, potentially enhancing model performance. We will focus on two drugs with known biomarkers but with different mechanisms of action: Erlotinib (an EGFR inhibitor) and JQ1 (a BET inhibitor).

Setup and Libraries

library(tidyverse)
library(data.table)
library(ggpubr)
library(caret)
library(doParallel)

Methods: Machine Learning Models

We will use two machine learning algorithms to predict drug sensitivity (log2 fold change) from the DepMap PRISM drug screen data:

Random Forest (RF): A tree-based ensemble method that can capture non-linear relationships.
Elastic Net: A linear regression method with and regularization to prevent overfitting.

We are using the R Caret package for model training and hyperparameter tuning. Caret is a versatile package that provides a unified interface for training various machine learning models with built-in cross-validation and hyperparameter tuning. Documentation can be found here: https://topepo.github.io/caret/.

We will evaluate performance using R-squared () on six independent test-train sets across multiple random splits.

Prediction Functions

We define helper functions to perform repeated train-test splits and model training.

# Random Forest Prediction Function
rf_drug_prediction <- function(feature_dat, drug_response, output_file_prefix, n_repeats = 6) {
    r2_values <- c()
    for (i in 1:n_repeats) {
        set.seed(524 + i)
        # 80-20 Train-Test Split
        train_idx <- sample(1:nrow(feature_dat), size = round(0.8 * nrow(feature_dat)), replace = FALSE)
        test_idx <- setdiff(1:nrow(feature_dat), train_idx)

        train_x <- feature_dat[train_idx, ]
        train_y <- drug_response[train_idx]
        test_x <- feature_dat[test_idx, ]
        test_y <- drug_response[test_idx]

        # Train Random Forest with caret
        # caret will perform tuning for all available hyperparameters, and automatically pick the best model (highest R2 on cross-validation)
        set.seed(524 + i)
        rfTune <- train(train_x, train_y,
            method = "rf",
            importance = TRUE,
            ntree = 1000,
            preProcess = c("center", "scale"),
            metric = "Rsquared",
            tuneLength = 10, # tune over 10 different mtry values
            trControl = trainControl(method = "repeatedcv", number = 10, repeats = 3) # 10-fold cross-validation repeated 3 times
        )

        # Predict on Test Set
        rf_pred <- predict(rfTune, newdata = test_x)
        r2 <- cor(rf_pred, test_y, method = "pearson")^2
        r2_values <- c(r2_values, r2)
    }

    # Return R2 values
    return(data.frame(Repeat = 1:n_repeats, R2 = r2_values))
}

# Elastic Net Prediction Function
enet_drug_prediction <- function(feature_dat, drug_response, output_file_prefix, n_repeats = 6) {
    r2_values <- c()
    for (i in 1:n_repeats) {
        set.seed(524 + i)
        # 80-20 Train-Test Split
        train_idx <- sample(1:nrow(feature_dat), size = round(0.8 * nrow(feature_dat)), replace = FALSE)
        test_idx <- setdiff(1:nrow(feature_dat), train_idx)

        train_x <- feature_dat[train_idx, ]
        train_y <- drug_response[train_idx]
        test_x <- feature_dat[test_idx, ]
        test_y <- drug_response[test_idx]

        # Train Elastic Net
        set.seed(524 + i)
        enetTune <- train(train_x, train_y,
            method = "glmnet",
            preProcess = c("center", "scale"),
            metric = "Rsquared",
            tuneLength = 10, # tune over 10 X 10 different alpha/lambda combinations
            # alpha: mixing parameter (0 = ridge, 1 = lasso)
            # lambda: regularization strength
            trControl = trainControl(method = "repeatedcv", number = 10, repeats = 3) # 10-fold cross-validation repeated 3 times
        )

        # Predict on Test Set
        enet_pred <- predict(enetTune, newdata = test_x)
        r2 <- cor(enet_pred, test_y, method = "pearson")^2
        r2_values <- c(r2_values, r2)
    }

    # Return R2 values
    return(data.frame(Repeat = 1:n_repeats, R2 = r2_values))
}

Analysis 1: Erlotinib Prediction

Erlotinib is an EGFR inhibitor used in cancer treatment, which is a type of targeted cancer therapy used primarily to treat certain types of non-small cell lung cancer and pancreatic cancer.

Data Preparation

We load the data and prepare four feature sets:

Full RNA: All high-variance genes.
RNA AE: Deep features from RNA autoencoder.
Multi-omics AE: Deep features from multi-omics autoencoder (including RNA, mutation, and CRISPR).
Selected Top RNA: Top ~4000 gene expressions most correlated with Erlotinib response (a common feature selection strategy).

# Load Data
load("CCLE_24Q2_GE_match_sample_info.RData") # RNA expression in log2(TPM+1)
load("PRISM_24Q2_compound_screen_match_sample_info.RData") # PRISM drug response data in log2 fold change (LFC)
cmpd_dat <- read_csv("Repurposing_Public_24Q2_Extended_Primary_Compound_List.csv") # Compound metadata
load("sample_info_match_biomarkers.RData") # Sample info with DepMap IDs

# ... (Data filtering and matching code similar to Part 1) ...

# Get Erlotinib Response
cmpd_id <- cmpd_dat |>
    filter(Drug.Name == "ERLOTINIB") |>
    pull(IDs)
drug_LFC <- prism_dat_match_sam |> pull(cmpd_id)

Let’s take a look at the PRISM drug screen data. Rows are cell lines and columns are compounds. Values are log2 fold change (LFC) in viability after drug treatment. You can see that there are a lot of NAs here, as many cell lines were not screened with all compounds.

load("PRISM_24Q2_compound_screen_match_sample_info.RData")
prism_dat_match_sam[1:5, 1:5] |> gt::gt()

BRD:BRD-A00047421-001-01-7	BRD:BRD-A00055058-001-01-0	BRD:BRD-A00077618-236-07-6	BRD:BRD-A00092689-236-04-9	BRD:BRD-A00100033-001-08-9
-1.207281	0.5157434	-0.01557664	-0.39512253	-0.4493321
-4.231563	NA	NA	-0.53837559	NA
NA	NA	NA	NA	NA
-3.860672	NA	NA	0.30697134	NA
-2.271411	NA	NA	0.03509603	NA

For brevity, we skip the data cleaning steps here. After filtering for missing data and matching cell lines across datasets, we obtain the final datasets for modeling. We also skip the codes for selecting the top RNA genes correlated with Erlotinib response - it is a simple apply function calculating Pearson correlation for each gene across the transcriptome, and selecting the top genes.

Model Training (Erlotinib)

We train models on all four feature sets.

# Parallel Processing
cl <- makePSOCKcluster(22) # Adjust number of cores as needed
registerDoParallel(cl)

# 1. Full RNA
rf_drug_prediction(ge_dat_final, drug_LFC_final, "full_RNA")
enet_drug_prediction(ge_dat_final, drug_LFC_final, "full_RNA")

# 2. RNA AE Features
rf_drug_prediction(RNA_AE_features_final, drug_LFC_final, "RNA_AE")
enet_drug_prediction(RNA_AE_features_final, drug_LFC_final, "RNA_AE")

# 3. Multi-omics AE Features
rf_drug_prediction(MultiOmics_AE_features_final, drug_LFC_final, "MultiOmics_AE")
enet_drug_prediction(MultiOmics_AE_features_final, drug_LFC_final, "MultiOmics_AE")

# 4. Selected Top RNA
rf_drug_prediction(ge_dat_selected, drug_LFC_final, "selected_top_RNA")
enet_drug_prediction(ge_dat_selected, drug_LFC_final, "selected_top_RNA")

stopCluster(cl)

Results: Erlotinib

Let’s look at the performance () of the different models.

# Load results (code omitted for brevity, loading CSVs from results folder)

# Plot Random Forest Results
ggplot(rf_r2_df, aes(x = Model, y = R2, fill = Model, color = Model)) +
    geom_boxplot(alpha = 0.5, outliers = FALSE) +
    geom_jitter(width = 0.2, size = 4) +
    theme_classic(base_size = 20) +
    labs(title = "Random Forest Erlotinib LFC Prediction R2", x = "Model", y = "R2") +
    theme(legend.position = "none")

# Plot Elastic Net Results
ggplot(enet_r2_df, aes(x = Model, y = R2, fill = Model, color = Model)) +
    geom_boxplot(alpha = 0.5, outliers = FALSE) +
    geom_jitter(width = 0.2, size = 4) +
    theme_classic(base_size = 20) +
    labs(title = "Elastic Net Erlotinib LFC Prediction R2", x = "Model", y = "R2") +
    theme(legend.position = "none")

This result is a bit surprising! In both models, the selected top RNA genes outperform all other feature sets. The full RNA set ranks second, while the AE features perform worse. The multi-omics AE features do not seem to add value in this case, and it performs even worse than the RNA AE features alone. I’ll speculate why this might be the case in the summary section. Is this specific to Erlotinib, or a general trend? Let’s check with another drug.

Analysis 2: JQ1 Prediction

Now let’s try to predict sensitivity to JQ1, a BET inhibitor that has a very different mechanism of action compared to Erlotinib.

Model Training (JQ1)

We repeat the same process for JQ1.

Results: JQ1

# ... (Similar plotting code as Erlotinib) ...

So the conclusion is pretty much the same for JQ1 as well. The selected top RNA genes outperform all other feature sets, followed by the full RNA set. The AE features again perform worse, and adding multi-omics data does not help.

Summary and Conclusions

Contrary to my initial hypothesis, the deep features extracted from autoencoders did not improve drug sensitivity prediction accuracy for either Erlotinib or JQ1. Instead, using a feature-selected subset of RNA expression data yielded the best performance, followed by using the full RNA expression data. The AE features, both RNA-only and multi-omics, underperformed in comparison.

So why is this the case? Here are some possible explanations:

Information Loss: The AE compression may have discarded important predictive information present in the original RNA expression data. Rather than preserving features relevant for drug response, the AE might have prioritized reconstructing general patterns, which in this case is likely the lineage programs between different cell lines. A possible remedy could be to use supervised or semi-supervised autoencoders that incorporate drug response information during training. I also have not experimented with different AE architectures, latent dimensions, or training strategies that might better capture drug response signals.
Multi-omics Integration Challenges: Combining multiple omics data types (RNA, mutation, CRISPR) into a single AE may introduce noise or conflicting signals that obscure relevant features for drug response. Each omics type has different sparsity and scales, making it difficult for a single AE to effectively learn a unified representation. More sophisticated integration methods or separate AEs for each omics type followed by feature fusion might yield better results.
Drug-Specific Biology: Assuming that the AE mainly captures broad biological variation (e.g., lineage), it may not align well with the specific molecular mechanisms driving sensitivity to Erlotinib and JQ1. For example, one of the key biomarkers for Erlotinib sensitivity - EGFR hotspot mutations - are actually quite rare across the cell lines (I’ve counted around 10), and may not be well represented in the AE features. I also did not explore including other omics data types (e.g., proteomics, methylation) that might be relevant for these drugs. Finally, only two drugs were tested here; results may vary for other drugs with different mechanisms of action.
Limitation of ML Models: The Random Forest and Elastic Net models used here may not fully leverage the complex representations learned by the AEs. More advanced models (e.g., deep neural networks) that can capture non-linear relationships in the AE feature space might yield different results. However, one significant drawback of deep learning is that it generally requires much larger training datasets to avoid overfitting, which may not be feasible with the limited number of cell lines available in a regular drug screen.

Interestingly, between the 2 models, the Elastic Net seemed to slightly outperform Random Forest in most cases. Given that the training time for Elastic Net is significantly shorter than Random Forest, this suggests that for these specific drug response predictions, simpler linear models with regularization may be more effective than more complex non-linear models.

So there you have it! In the next and final part of this series, I will revisit the concept I introduced in my first post: enhancing statistical power of drug biomarker detection by ML. I will use the best performed model (selected top RNA with Elastic Net) to expand drug prediction to all DepMap available cell lines (e.g., there are only ~500 cell lines screened for Erlotinib, but DepMap has close to 2000 cell lines), and see if we can better recover known sensitivity biomarkers.

References

DepMap Portal: https://depmap.org/portal/
Caret Package Documentation: https://topepo.github.io/caret/
Applied Predictive Modeling by Kuhn and Johnson: https://www.springer.com/gp/book/9781461468486
An Introduction to Statistical Learning by James et al.: https://www.statlearning.com/

Machine learning for drug sensitivity prediction (Part 1): using autoencoders to extract deep features from DepMap OMICS

Jay Chung — Mon, 01 Dec 2025 08:00:00 GMT

During my time in the industry, I have analyzed quite a few cell panel screens aiming to identify biomarkers of drug sensitivity. These screens are often carried out by CROs, and due to the costs associated with these experiments, the number of cell lines screened is often limited. Very often the average number of cell lines in each lineage is less than 10, which limits the power of biomarker detection. One way to mitigate this issue is to train a machine learning model based on the available cell lines and their corresponding DepMap omics data, and then use the trained model to predict sensitivity across thousands of DepMap cell lines. This expanded set of predicted sensitivities can then be used for biomarker detection with improved power.

In this post, I want to explore whether using autoencoders (AE) to compress high-dimensional omics data into deep features can improve drug sensitivity prediction accuracy. The concept can be summarized in the figure below:

I demonstrate how to use AE to compress DepMap RNA expression data alone and multi-omics data (RNA expression + mutation + CRISPR dependency scores) into deep features. I then visualize and compare the embeddings from RNA alone vs. multi-omics data. In future work (part 2), I will compare the drug sensitivity prediction accuracy using deep features vs. full omics data.

These concepts can be summarized in Component 1 (AE embedding) and 2 (ML training) in the following figures:

Component 1: Autoencoder embedding

Component 2: ML training and prediction

Introduction

Background: Cell panel drug screens are very useful during drug development in identifying sensitive biomarkers and understanding the mechanism-of-action; however, it can be costly and often result in only up to hundreds of cell lines available, limiting the power of biomarker detection. By leveraging machine learning (ML) approaches, one can expand sensitivity prediction to thousands of DepMap cell lines and thus improve the power of biomarker detection.

Problem: DepMap OMICS data are high-dimensional, and such high-dimensional data may hinder ML prediction due to overfitting problems and noisy data.

Proposal: By leveraging AutoEncoder (AE), a deep learning approach, high-dimensional OMICS data can be compressed into Deep Features (DF), thus overcoming the problem of overfitting and resulting in better drug sensitivity prediction accuracy. Unlike traditional dimensionality reduction methods (e.g., PCA), AE can capture non-linear relationships in the data, leading to more informative embeddings.

Setup and Libraries

library(tidyverse)
library(data.table)
library(ggpubr)
library(h2o)
library(umap)
library(ggrepel)
library(microViz)

Aim 1: RNA Expression Autoencoder

First, we use an autoencoder to embed cell lines based on RNA expression data alone.

Data Loading and Preprocessing

We filter for high-variance genes to reduce noise and dimensionality before feeding the data into the autoencoder.

# load data: RNA expression
# Note: these RData files contain preprocessed data matrices with matched sample information
# DepMap data can be downloaded from https://depmap.org/portal/download/
# Here I am using the 24Q2 release
load("CCLE_24Q2_GE_match_sample_info.RData") # rows are cell lines, columns are genes, values are RNA-seq log2(TPM+1)
load("sample_info_match_biomarkers.RData") # sample info with lineage annotations

# remove low variance genes
ge_var <- apply(ccle_ge_match_sam, 2, var, na.rm = TRUE)
ge_dat_hv <- ccle_ge_match_sam[, which(ge_var > quantile(ge_var, 0.2))] # select high variance genes (top 80th percentile)

# remove any cell lines (rows) that contains NAs
ge_dat_hv <- ge_dat_hv |> na.omit()

# standardize the data
ge_dat_hv <- scale(ge_dat_hv) |> as.data.frame()

Let’s see what the RNA expression data looks like:

ge_dat_hv[1:5, 1:5] |>
    rownames_to_column("CellLine") |>
    gt::gt(rowname_col = "CellLine")

	TSPAN6	DPM1	FIRRM	FGR	CFH
NIHOVCAR3	1.0926276	1.5011722	0.6858109	-0.3170060	-0.5737765
HL60	-1.9693314	-1.2610758	-0.6848542	3.0178580	-0.8997009
CACO2	1.1699782	2.0373191	0.2826820	-0.3517010	-0.9468198
HEL	-0.7463001	-1.6419864	0.3162823	0.3703802	1.2136680
HEL9217	-0.5779999	-0.9799468	2.0368574	-0.1876538	1.5663762

Let’s see what the sample info looks like:

sample_info[1:5, ] |> gt::gt()

ModelID	StrippedCellLineName	CCLEName	OncotreeLineage	OncotreeSubtype	OncotreePrimaryDisease	LegacySubSubtype	PatientMolecularSubtype
ACH-000001	NIHOVCAR3	NIHOVCAR3_OVARY	Ovary/Fallopian Tube	High-Grade Serous Ovarian Cancer	Ovarian Epithelial Tumor	high_grade_serous
ACH-000002	HL60	HL60_HAEMATOPOIETIC_AND_LYMPHOID_TISSUE	Myeloid	Acute Myeloid Leukemia	Acute Myeloid Leukemia	M3	TP53(del), CDKN2A and NRAS mutations [PubMed=288488], No PML-RARA fusion
ACH-000003	CACO2	CACO2_LARGE_INTESTINE	Bowel	Colon Adenocarcinoma	Colorectal Adenocarcinoma
ACH-000004	HEL	HEL_HAEMATOPOIETIC_AND_LYMPHOID_TISSUE	Myeloid	Acute Myeloid Leukemia	Acute Myeloid Leukemia	M6	JAK2 and TP53 mutations,
ACH-000005	HEL9217	HEL9217_HAEMATOPOIETIC_AND_LYMPHOID_TISSUE	Myeloid	Acute Myeloid Leukemia	Acute Myeloid Leukemia	M6	JAK2 and TP53 mutations

Model Training

We perform a hyperparameter grid search to find the optimal autoencoder architecture.

# Hyperparameter search grid to find best AE model
h2o.init(nthreads = 20, max_mem_size = "60G") # initialize H2O with 20 threads and 60GB memory
rna_h2o <- as.h2o(ge_dat_hv)

hyper_grid <- list(hidden = list(
    c(1000, 500, 1000),
    c(1000, 500, 200, 500, 1000),
    c(1000, 100, 50, 100, 1000),
    c(1000, 500, 100, 500, 1000),
    c(2000, 1000, 200, 1000, 2000),
    c(1000, 500, 300, 500, 1000)
))

ae_grid <- h2o.grid(
    algorithm = "deeplearning",
    x = colnames(rna_h2o),
    training_frame = rna_h2o,
    grid_id = "RNA_ae1",
    autoencoder = TRUE,
    activation = "TanhWithDropout", # I've also tried "RectifierWithDropout" but seemed to encounter exploding gradients
    hyper_params = hyper_grid,
    nesterov_accelerated_gradient = TRUE,
    epochs = 30,
    stopping_rounds = 5,
    seed = 524
)

# get grid results, sorted by reconstruction error (MSE)
h2o.getGrid("RNA_ae1", sort_by = "mse", decreasing = FALSE)

# AE learning with best parameters (selected from grid search)
ae_model <- h2o.deeplearning(
    x = colnames(rna_h2o),
    training_frame = rna_h2o,
    autoencoder = TRUE,
    hidden = c(1000, 500, 200, 500, 1000),
    hidden_dropout_ratios = c(0.5, 0.5, 0.5, 0.5, 0.5),
    activation = "TanhWithDropout",
    epochs = 30,
    stopping_rounds = 5,
    nesterov_accelerated_gradient = TRUE,
    seed = 524
)

# extract the compressed features
compressed_features <- h2o.deepfeatures(ae_model, rna_h2o, layer = 3) |> as.data.frame() # get deep features from the middle layer
rownames(compressed_features) <- rownames(ge_dat_hv)

# Save results
fwrite(compressed_features, file = "./results/output_files/RNA_AE_compressed_features.csv", sep = ",", row.names = TRUE)
h2o.shutdown()

I want to point out that the choice of architecture (number of layers and nodes) and hyperparameters (dropout rates, activation functions) can significantly impact the quality of the learned embeddings. For example, using batch normalization layers or experimenting with different activation functions (e.g., ReLU, Leaky ReLU) might yield better results depending on the data characteristics. Readers interested in this topic should explore more sophisticated MLP training platform like Keras and TensorFlow.

Aim 2: Multi-omics Autoencoder

Next, we integrate RNA expression, mutation data (damaging, hotspot), and CRISPR dependency scores.

Data Integration

We process each data type separately (filtering, standardizing) and then combine them into a single multi-omics dataset.

# load data (all rows are cell lines and columns are genes)
load("CCLE_24Q2_GE_match_sample_info.RData") # RNA expression
load("DepMap_24Q2_Chronos_match_sample_info.RData") # CRISPR dependency scores (Chronos)
load("CCLE_24Q2_DAMMUT_match_sample_info.RData") # Damaging mutations (0: non-mutated, 1: mutated heterozygous, 2: mutated homozygous)
load("CCLE_24Q2_HOTMUT_match_sample_info.RData") # Hotspot mutations (0: non-mutated, 1: mutated heterozygous, 2: mutated homozygous)
load("sample_info_match_biomarkers.RData") # sample info

# RNA processing
ge_var <- apply(ccle_ge_match_sam, 2, var, na.rm = TRUE)
ge_dat_hv <- ccle_ge_match_sam[, which(ge_var > quantile(ge_var, 0.2))]
ge_dat_hv <- scale(ge_dat_hv) |> as.data.frame()

# Mutation processing
dam_mut_dat_filt <- dammut_dat_match_sam[, which(apply(na.omit(dammut_dat_match_sam), 2, function(x) sum(x != 0)) > 20)] # filter genes with >20 mutated events across cell lines
hot_mut_dat_filt <- hotmut_dat_match_sam[, which(apply(na.omit(hotmut_dat_match_sam), 2, function(x) sum(x != 0)) > 20)]

# CRISPR processing
crispr_var <- apply(chronos_dat_match_sam, 2, var, na.rm = TRUE)
crispr_dat_hv <- chronos_dat_match_sam[, which(crispr_var > quantile(crispr_var, 0.2))]
crispr_dat_hv <- scale(crispr_dat_hv) |> as.data.frame()

# Combine data
colnames(ge_dat_hv) <- paste0(colnames(ge_dat_hv), "_RNA")
colnames(dam_mut_dat_filt) <- paste0(colnames(dam_mut_dat_filt), "_DAMMUT")
colnames(hot_mut_dat_filt) <- paste0(colnames(hot_mut_dat_filt), "_HOTMUT")
colnames(crispr_dat_hv) <- paste0(colnames(crispr_dat_hv), "_CRISPR")

multi_omics_dat <- cbind(ge_dat_hv, dam_mut_dat_filt, hot_mut_dat_filt, crispr_dat_hv)

# Remove rows with too many NAs
multi_omics_dat <- multi_omics_dat[rowSums(is.na(multi_omics_dat)) <= ncol(multi_omics_dat) * 0.8, ] # keep rows with <=80% NAs

Multi-omics Model Training

We train another autoencoder on this combined dataset.

h2o.init(nthreads = 20, max_mem_size = "60G")
rna_h2o <- as.h2o(multi_omics_dat)

# (Grid search code omitted for brevity, similar to Aim 1)

# AE learning with best parameters
ae_model <- h2o.deeplearning(
    x = colnames(rna_h2o),
    training_frame = rna_h2o,
    autoencoder = TRUE,
    hidden = c(1000, 500, 300, 500, 1000),
    hidden_dropout_ratios = c(0.5, 0.5, 0.5, 0.5, 0.5),
    activation = "TanhWithDropout",
    epochs = 30,
    stopping_rounds = 5,
    nesterov_accelerated_gradient = TRUE,
    seed = 524
)

compressed_features <- h2o.deepfeatures(ae_model, rna_h2o, layer = 3) |> as.data.frame() # here we take 300-dim deep features
rownames(compressed_features) <- rownames(multi_omics_dat)

fwrite(compressed_features, file = "./results/output_files/MultiOmics_AE_compressed_features.csv", sep = ",", row.names = TRUE)
h2o.shutdown()

Aim 3: Visualization and Comparison

Finally, we visualize the embeddings using UMAP and compare the clustering with known lineages.
Here we are testing if the deep features preserve and/or enrich biological signals present in the original data.

Full RNA UMAP (without Autoencoder)

# load data
load("CCLE_24Q2_GE_match_sample_info.RData")
load("sample_info_match_biomarkers.RData")

# remove low variance genes
ge_var <- apply(ccle_ge_match_sam, 2, var, na.rm = TRUE) # calculate variance for each gene
sum(ge_var > quantile(ge_var, 0.2)) # how many genes passed variance cutoff: 15322
ge_dat_hv <- ccle_ge_match_sam[, which(ge_var > quantile(ge_var, 0.2))] # select high variance genes

# standardize the data and match with sample info
ge_dat_hv <- scale(ge_dat_hv) |>
    as.data.frame() |>
    na.omit()
sample_info <- sample_info[match(rownames(ge_dat_hv), sample_info$StrippedCellLineName), ]
all(rownames(ge_dat_hv) == sample_info$StrippedCellLineName)

# Run UMAP
library(umap)
umap_dat <- umap(
    ge_dat_hv,
    n_components = 2,
    n_neighbors = 50,
    random_state = 524,
    verbose = TRUE,
    n_threads = 15
)
umap_dat <- as.data.frame(umap_dat$layout)
umap_dat$cell_line <- sample_info$StrippedCellLineName
umap_dat$lineage <- sample_info$OncotreeLineage
umap_dat$subtype <- sample_info$OncotreeSubtype

# Plot UMAP embedding colored by lineage
ggplot(umap_dat, aes(V1, -V2, color = lineage, fill = lineage)) +
    geom_point(size = 3, alpha = 0.5) +
    scale_fill_manual(values = distinct_palette(n = length(unique(umap_dat$lineage)), pal = "brewerPlus", add = "lightgrey")) +
    scale_color_manual(values = distinct_palette(n = length(unique(umap_dat$lineage)), pal = "brewerPlus", add = "lightgrey")) +
    theme_classic(base_size = 20) +
    # add lineage labels at the center of each lineage cluster
    geom_label_repel(
        data = umap_dat %>%
            group_by(lineage) %>%
            summarize(V1 = median(V1), V2 = median(V2)),
        aes(label = lineage, color = lineage),
        size = 4,
        # color = "black",
        fill = "white",
        alpha = 0.7,
        show.legend = FALSE,
        max.overlaps = 100
    ) +
    theme(
        legend.position = "none",
        legend.title = element_text(size = 10),
        legend.text = element_text(size = 10)
    ) +
    labs(
        x = "UMAP1",
        y = "UMAP2",
        title = "UMAP of full CCLE RNA expression (no AE)"
    )

As you can see, the full RNA data UMAP already shows some clustering by lineage, indicating that the transcriptomic profiles capture biological differences among cell lines.

RNA deep features UMAP

# Load compressed features
RNA_AE_compressed_features <- fread("./results/output_files/RNA_AE_compressed_features.csv", data.table = FALSE) |>
    column_to_rownames(var = "V1")
RNA_AE_compressed_features <- RNA_AE_compressed_features[rowSums(RNA_AE_compressed_features) != 0, ]

# Match with sample info
sample_info <- sample_info[match(rownames(RNA_AE_compressed_features), sample_info$StrippedCellLineName), ]

# Run UMAP
umap_dat <- umap(RNA_AE_compressed_features, n_components = 2, n_neighbors = 50, random_state = 524)
umap_dat <- as.data.frame(umap_dat$layout)
umap_dat$lineage <- sample_info$OncotreeLineage

# Plot
ggplot(umap_dat, aes(V1, V2, color = lineage, fill = lineage)) +
    geom_point(size = 3, alpha = 0.5) +
    scale_fill_manual(values = distinct_palette(n = length(unique(umap_dat$lineage)), pal = "brewerPlus", add = "lightgrey")) +
    scale_color_manual(values = distinct_palette(n = length(unique(umap_dat$lineage)), pal = "brewerPlus", add = "lightgrey")) +
    theme_classic(base_size = 20) +
    geom_label_repel(
        data = umap_dat %>% group_by(lineage) %>% summarize(V1 = median(V1), V2 = median(V2)),
        aes(label = lineage, color = lineage),
        size = 4, fill = "white", alpha = 0.7, show.legend = FALSE, max.overlaps = 100
    ) +
    theme(legend.position = "none") +
    labs(title = "UMAP of AE compressed CCLE RNA expression")

The RNA deep features UMAP also shows clustering by lineage, similar to the full RNA data, indicating that the autoencoder effectively preserves the biological signal in a lower-dimensional space.

Multi-omics deep features UMAP

# Load compressed features
MultiOmics_AE_compressed_features <- fread("./results/output_files/MultiOmics_AE_compressed_features.csv", data.table = FALSE) |>
    column_to_rownames(var = "V1")
MultiOmics_AE_compressed_features <- MultiOmics_AE_compressed_features[rowSums(MultiOmics_AE_compressed_features) != 0, ]

# Match with sample info
sample_info <- sample_info[match(rownames(MultiOmics_AE_compressed_features), sample_info$StrippedCellLineName), ]

# Run UMAP
umap_dat <- umap(MultiOmics_AE_compressed_features, n_components = 2, n_neighbors = 50, random_state = 524)
umap_dat <- as.data.frame(umap_dat$layout)
umap_dat$lineage <- sample_info$OncotreeLineage

# Plot
ggplot(umap_dat, aes(-V1, -V2, color = lineage, fill = lineage)) +
    geom_point(size = 3, alpha = 0.5) +
    scale_fill_manual(values = distinct_palette(n = length(unique(umap_dat$lineage)), pal = "brewerPlus", add = "lightgrey")) +
    scale_color_manual(values = distinct_palette(n = length(unique(umap_dat$lineage)), pal = "brewerPlus", add = "lightgrey")) +
    theme_classic(base_size = 20) +
    geom_label_repel(
        data = umap_dat %>% group_by(lineage) %>% summarize(V1 = median(V1), V2 = median(V2)),
        aes(label = lineage, color = lineage),
        size = 4, fill = "white", alpha = 0.7, show.legend = FALSE, max.overlaps = 100
    ) +
    theme(legend.position = "none") +
    labs(title = "UMAP of AE compressed CCLE Multi-omics")

The AE embedded multi-OMICS appears to increase the granularity of cell line clustering by lineage.

Conclusions

It appears that both RNA alone and multi-omics autoencoder embeddings preserve biological signals related to cell lineage. The multi-omics embedding seems to provide even richer structure, potentially capturing more nuanced cell states. But whether these embeddings lead to improved drug sensitivity prediction accuracy remains to be tested.

In the next post, I will compare ML prediction accuracy on independent test data using deep features vs. full omics data.

References

DepMap Portal: https://depmap.org/portal/
H2O Documentation: https://docs.h2o.ai/h2o/latest-stable/h2o-docs/data-science/deep-learning.html
Hands-On Machine Learning with R: https://bradleyboehmke.github.io/HOML/