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

1. 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”.

2. 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.

3. 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:

import pandas as pd
import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.model_selection import GroupShuffleSplit
from Bio.Align import substitution_matrices
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.model_selection import train_test_split

See my previous post for source of data. First, let’s see what the input data looks like:

df.head()

Output:

   Peptide_ID           Peptide     MHC_ID  \
0      104653   VAPIEHIASMRRNYF  DRB1_1302   
1       37106   HDDKETSFIRNCARK  DRB1_0101   
2      118433   LIWVGINTRNMTMSM  DRB1_0101   
3       80770  GVTVIKNNMINNDLGP  DRB1_1501   
4       19888   PAPMLAAAAGWQTLS  DRB1_1101   

                                  MHC         Y  
0  QEFFIASGAAVDAIMESSFDYFDIDEATYHVGFT  0.501084  
1  QEFFIASGAAVDAIMWLFLECYDLQRATYHVGFT  0.441298  
2  QEFFIASGAAVDAIMWLFLECYDLQRATYHVGFT  0.217673  
3  QEFFIASGAAVDAIMWPRFDYFDIQAATYHVVFT  0.807811  
4  QEFFIASGAAVDAIMESSFDYFDFDRATYHVGFT  0.583271  

Here we have the peptide sequences, MHC pseudo-sequences, and their corresponding binding affinity values (Y).

Let’s perfrom peptide clustering using a 3-gram approach to find overlapping peptides without source information. We will use the CountVectorizer from sklearn to create a matrix of 3-gram counts for each unique peptide, and then compute the cosine similarity between the peptides based on this matrix. Finally, we will use hierarchical clustering to group similar peptides together.

vect = CountVectorizer(analyzer='char', ngram_range=(3, 3)) # sequence of 3 a.a.
pep_matrix = vect.fit_transform(df['Peptide'].unique()) # row: unique peptides, col: unique 3 a.a.
pep_sim = cosine_similarity(pep_matrix) # peptide X peptide similarity matrix
# use cosine dist when the orientation or pattern of the data is more important than the absolute scale
# often used in text, genetic sequence

# Cluster peptides that share many 3-mers (likely from the same protein)
pep_clusters = AgglomerativeClustering(
    n_clusters=None,
    distance_threshold=0.5, # Adjust: lower = more groups, stricter split
    metric='precomputed', # input is pre-computed dist
    linkage='complete'
).fit_predict(1 - pep_sim)

pep_map = pd.DataFrame({
    'Peptide': df['Peptide'].unique(),
    'pep_group': pep_clusters
})

df = df.merge(pep_map, on='Peptide')

print(pep_map.sort_values(by='pep_group').head(15))

Output:

                    Peptide  pep_group
1576         QYIKANAKFIGITE          0
702       YATFFIKANSKFIGITE          0
6842       MQYIKANSKFIGITEL          0
15556        QYQKANSKFIGITE          0
4185         QYIKANSKFIGITE          0
7940   SAMILAAYHPQQFIYAGSLS          1
16944  AAIGLSMAGSSAMILAAYHP          1
9611       RQIRMAKLLGRDPEQS          2
860        HGRQIRMAKLFGRDPE          2
6826       EGELHGRQIRMAKLLG          2
14166      HGRQIKMAKLLGRDPE          2
687         HGRQIRMAKLLGRDP          2
4640       HGRQIRMAKLLGRDPE          2
6532       HGRQIRMAKLLTRDPE          2
4614        QIRMAKLLGRDPEQS          2

We can see that peptides that are likely from the same protein (e.g., “QYIKANAKFIGITE”, “YATFFIKANSKFIGITE”, “MQYIKANSKFIGITEL”) are clustered together in the same group (group 0). This indicates that our clustering approach is effectively grouping similar peptides together based on their 3-gram composition.

Next, we will cluster MHC pseudo-sequences based on their BLOSUM62 amino acid similarity. We will define a function to calculate the BLOSUM62 similarity between two sequences, and then create a similarity matrix for the MHC pseudo-sequences. Finally, we will use hierarchical clustering to group similar MHCs together.

# Function to calculate BLOSUM62 similarity
def blosum62_similarity(seq1, seq2):
    matrix = substitution_matrices.load('BLOSUM62')
    score = 0.0

    # Align sequences by padding the shorter one (simple padding for score calculation)
    # Note: For rigorous alignment, a proper sequence alignment algorithm (e.g., Needleman-Wunsch) is needed.
    # This simplified approach assumes aligned positions are comparable.
    for i in range(min(len(seq1), len(seq2))):
        try:
            score += matrix[seq1[i], seq2[i]]
        except KeyError: # Handle cases where amino acid might not be in BLOSUM (e.g., 'X')
            score += 0 

    # Normalize score (for simplicity, divide by max possible score)
    # A more robust normalization might involve the self-similarity score.
    # Here a max possible similarity score for each aa is calculated and summed
    max_score1 = sum(matrix[aa, aa] for aa in seq1 if (aa, aa) in matrix.keys())
    max_score2 = sum(matrix[aa, aa] for aa in seq2 if (aa, aa) in matrix.keys())
    if max_score1 == 0 or max_score2 == 0: # Avoid division by zero
        return 0.0
    return score / max(max_score1, max_score2)

mhc_sequences = df['MHC'].unique()

# Create a similarity matrix for MHCs
n_mhc = len(mhc_sequences)
mhc_sim_matrix = np.zeros((n_mhc, n_mhc))
for i in range(n_mhc):
    for j in range(i, n_mhc):
        sim = blosum62_similarity(mhc_sequences[i], mhc_sequences[j])
        mhc_sim_matrix[i, j] = sim
        mhc_sim_matrix[j, i] = sim

# Cluster MHCs
mhc_clusters = AgglomerativeClustering(
    n_clusters=None,
    distance_threshold=0.2, # Adjust: lower = more groups, stricter split for MHCs
    metric='precomputed',
    linkage='complete'
).fit_predict(1 - mhc_sim_matrix)

mhc_map = pd.DataFrame({
    'MHC': mhc_sequences,
    'mhc_group': mhc_clusters
})

df = df.merge(mhc_map, on='MHC')

print(mhc_map.sort_values(by='mhc_group'))

Output:

                                   MHC  mhc_group
5   QEFFIASGAAVDAIMELSFEYYVLQKQNYHVVFT          0
15  QEFFIASGAAVDAIMERSYDYYVLQKRNYHVGFT          0
12  QEFFIASGAAVDAIMELSFEHYDLQKQNYHVGFT          0
8   QEFFIASGAAVDAIMESSYDYFDLQKRNYHVVFT          0
9   CNYHQGGGARVAHIMFFGLTYYDVGTETVHVAGI          1
..                                 ...        ...
64  YTYFLRRGGQTGHILHFPLIYYDYRTETVHKTPT         26
66  XXYHWTSGGQTGHGWALGSNYYDIRTETVHGVHT         27
70  QEFFIASGAAVDAIMESSFEYYDLQRATYHVGFT         28
62  QEFFIASGAAVDAIMESSFEYYDLQKRNYHVGFT         28
29  QEFFIASGAAVDAIMESGLEHFVIDRATYHAVFT         29

[75 rows x 2 columns]

We can see that MHC pseudo-sequences that are similar based on their BLOSUM62 scores are clustered together in the same group (e.g., “QEFFIASGAAVDAIMELSFEYYVLQKQNYHVVFT”, “QEFFIASGAAVDAIMERSYDYYVLQKRNYHVGFT”, “QEFFIASGAAVDAIMELSFEHYDLQKQNYHVGFT” are all in group 0). This indicates that our clustering approach is effectively grouping similar MHC pseudo-sequences together based on their amino acid composition and similarity.

Next, we will perform the regular C3 split. To ensure total isolation between the training and testing sets, we will create a ‘SuperGroup’ that combines both the peptide and MHC groups. This way, we can ensure that no similar peptide-MHC pairs are present in both the training and testing sets, thus minimizing data leakage.

# Perform the "Regular C3 split" first
# We create a 'SuperGroup' that combines both to ensure total isolation
df['super_group'] = df['pep_group'].astype(str) + "_" + df['mhc_group'].astype(str)

# Split train vs temp_test
gss = GroupShuffleSplit(n_splits=1, train_size=0.7, random_state=42)
train_indices, temp_test_indices = next(gss.split(df, groups=df['super_group']))

train_df = df.iloc[train_indices][['Peptide_ID', 'Peptide', 'MHC_ID', 'MHC', 'Y']]
temp_test_df = df.iloc[temp_test_indices]

# Split temp_text into test and valid
gss1 = GroupShuffleSplit(n_splits=1, train_size=0.66, random_state=42)
test_indices, valid_indices = next(gss1.split(temp_test_df, groups=temp_test_df['super_group']))

test_df = df.iloc[test_indices][['Peptide_ID', 'Peptide', 'MHC_ID', 'MHC', 'Y']]
valid_df = df.iloc[valid_indices][['Peptide_ID', 'Peptide', 'MHC_ID', 'MHC', 'Y']]

print(df)

Output:

        Peptide_ID               Peptide                 MHC_ID  \
0            29993       MSGPMQQLTQPLQQV  HLA-DPA10201-DPB11401   
1            50950  CGKYLFNWAVRTKLKLTPIA              DRB1_1302   
2            58992       YKRQLMNILGAVYRY  HLA-DPA10201-DPB10101   
3             8140       FLGCLVKEIPPRLLY  HLA-DQA10501-DQB10201   
4           111514       KTQIDQVESTAGSLQ  HLA-DPA10201-DPB10501   
...            ...                   ...                    ...   
134276       15592       RFFLPIFSEFVLLAT              DRB1_0405   
134277       21667      EVFFQRLGIASGRARY              DRB1_1302   
134278       62271       YLFAKDKSGPLQPGV  HLA-DQA10102-DQB10602   
134279       79102       GELQIVDKIDADFKI              DRB1_1302   
134280      100096       SAAPLRTITADTFRK              DRB1_0701   

                                       MHC         Y  pep_group  mhc_group  \
0       YAFFQFSGGAILNTLHLQFEYFDLEKVRVHLDVT  0.216969       6334          6   
1       QEFFIASGAAVDAIMESSFDYFDIDEATYHVGFT  0.559077        205          3   
2       YAFFQFSGGAILNTLYGQFEYFAIEKVRVHLDVT  0.423504        457          6   
3       YNYHQRXFATVLHSLYFGLSSFAIRKARVHLETT  0.328660       1876         21   
4       YAFFQFSGGAILNTLFGQFEYFEIEKVRMHLDVT  0.000000       1157          6   
...                                    ...       ...        ...        ...   
134276  QEFFIASGAAVDAIMEVHFDYYSLQRATYHVGFT  0.721843       6951          4   
134277  QEFFIASGAAVDAIMESSFDYFDIDEATYHVGFT  0.257618       3632          3   
134278  CNYHQGGGARVAHIMFFGLTYYDVGTETVHVAGI  0.141380       2158          1   
134279  QEFFIASGAAVDAIMESSFDYFDIDEATYHVGFT  0.409599       3568          3   
134280  QEFFIASGAAVDAIMWGYFELYVIDRQTVHVGFT  0.242001       1440         17   

       super_group  
0           6334_6  
1            205_3  
2            457_6  
3          1876_21  
4           1157_6  
...            ...  
134276      6951_4  
134277      3632_3  
134278      2158_1  
134279      3568_3  
134280     1440_17  

[134281 rows x 8 columns]

Now, let’s perform the strict C3 split. In this split, we will ensure that for every pair in the test set, both the peptide and MHC classes are entirely absent from the training set. This means that we will first split the MHC groups and then split the peptide groups independently, ensuring that there is no overlap in either component between the training and testing sets.

def strict_double_split(df, train_size=0.7):
    # 1. First, split the MHC groups
    gss_mhc = GroupShuffleSplit(n_splits=1, train_size=train_size, random_state=42)
    mhc_train_idx, mhc_test_idx = next(gss_mhc.split(df, groups=df['mhc_group']))
    
    mhc_train_alleles = df.iloc[mhc_train_idx]['mhc_group'].unique()
    mhc_test_alleles = df.iloc[mhc_test_idx]['mhc_group'].unique()
    
    # 2. Second, split the Peptides groups
    # This prevents the model from seeing different fragments of the same protein
    gss_pep = GroupShuffleSplit(n_splits=1, train_size=train_size, random_state=42)
    pep_train_idx, pep_test_idx = next(gss_pep.split(df, groups=df['pep_group']))
    
    pep_train_prots = df.iloc[pep_train_idx]['pep_group'].unique()
    pep_test_prots = df.iloc[pep_test_idx]['pep_group'].unique()
    
    # 3. Create the "Strict" Test Set: 
    # Only interactions where BOTH the MHC is new AND the Protein is new.
    train_df = df[df['mhc_group'].isin(mhc_train_alleles) & df['pep_group'].isin(pep_train_prots)].copy()
    test_df = df[df['mhc_group'].isin(mhc_test_alleles) & df['pep_group'].isin(pep_test_prots)].copy()
    
    return train_df, test_df

train_temp, test_strict = strict_double_split(df)

# Random split of data into train or valid data
train_strict, valid_strict = train_test_split(train_temp, test_size=0.125, random_state=42)

# Save data to disc
import os
save_path = '/data'
train_df.to_feather(os.path.join(save_path, 'train_dat_cold.feather'))
valid_df.to_feather(os.path.join(save_path, 'valid_dat_cold.feather'))
test_df.to_feather(os.path.join(save_path, 'test_dat_cold.feather'))
train_strict.to_feather(os.path.join(save_path, 'train_dat_strict.feather'))
valid_strict.to_feather(os.path.join(save_path, 'valid_dat_strict.feather'))
test_strict.to_feather(os.path.join(save_path, 'test_dat_strict.feather'))

train_df.shape, test_df.shape, valid_df.shape
train_strict.shape, test_strict.shape, valid_strict.shape

Output:

Regular split shape:
((93392, 5), (26531, 5), (14358, 5))

Strict split shape:
((59757, 8), (9588, 8), (8537, 8))

We can see that the regular C3 split results in a larger training and testing set compared to the strict C3 split. When we prepare the data for model training, we will need to downsample the regular C3 split to match the sample size of the strict C3 split to ensure a fair comparison between the two splitting strategies.

2. Embeddings Extraction and Model Training

Embedding extraction and model training procedures are similar to my previous post, so I will not go through the code in detail here. Let’s load the data and downsample the regular C3 split to match the sample size of the strict C3 split:

all_dat = {
    'train': pd.read_feather(os.path.join(load_path, 'train_dat_cold.feather')).sample(n=59757, random_state=42).reset_index(drop=True), 
    'valid': pd.read_feather(os.path.join(load_path, 'valid_dat_cold.feather')).sample(n=8537, random_state=42).reset_index(drop=True),
    'test': pd.read_feather(os.path.join(load_path, 'test_dat_cold.feather')).sample(n=9588, random_state=42).reset_index(drop=True)
}

# Shuffle samples
for keys, df in all_dat.items():
  all_dat[keys] = df.sample(frac=1, random_state=42).reset_index(drop=True)

For both data splits, we will extract ESM2 embeddings (facebook/esm2_t30_150M_UR50D) for the peptides and MHC pseudo-sequences, concatenate them, and then train a neural network model with convolutional layers. The model architecture and training procedure will be the same for both splits.

Here is the model architecture we will use:

# Define 1D CNN + attention + MLP with multiple inputs using functional API
import tensorflow as tf
tf.keras.backend.clear_session()

tf.random.set_seed(42)

def make_dense_block(n_neurons, dropout_rate=0.2):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(n_neurons, activation='relu', kernel_initializer='he_normal'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(dropout_rate)
    ])

# Define the input layers
embedding_input = tf.keras.layers.Input(shape=embedding_dict['train'].shape[1:], name='embedding_input')
peptide_length_input = tf.keras.layers.Input(shape=(1,), name='peptide_length_input')
mhc_length_input = tf.keras.layers.Input(shape=(1,), name='mhc_length_input')

# CNN branch for embeddings
x = tf.keras.layers.Conv1D(filters=128, kernel_size=5, padding='same', activation='relu')(embedding_input)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.Conv1D(filters=64, kernel_size=3, padding='same', activation='relu')(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.AveragePooling1D(pool_size=2)(x)

# Self-attention layer
attention_output = tf.keras.layers.Attention()([x, x])
flat_attention = tf.keras.layers.Flatten()(attention_output)

# Concatenate attention output with length inputs
combined_features = tf.keras.layers.Concatenate()([flat_attention, peptide_length_input, mhc_length_input])

# MLP branch
y = make_dense_block(128)(combined_features)
y = make_dense_block(128)(y)
y = make_dense_block(128)(y)
y = make_dense_block(128)(y)
y = make_dense_block(128)(y)
output = tf.keras.layers.Dense(1)(y)

# Create the functional model
cnn_model = tf.keras.Model(inputs=[embedding_input, peptide_length_input, mhc_length_input], outputs=[output])

cnn_model.summary()

 Model: "functional_5"

┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)        ┃ Output Shape      ┃    Param # ┃ Connected to      ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩
│ embedding_input     │ (None, 71, 640)   │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv1d (Conv1D)     │ (None, 71, 128)   │    409,728 │ embedding_input[… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ batch_normalization │ (None, 71, 128)   │        512 │ conv1d[0][0]      │
│ (BatchNormalizatio… │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout (Dropout)   │ (None, 71, 128)   │          0 │ batch_normalizat… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv1d_1 (Conv1D)   │ (None, 71, 64)    │     24,640 │ dropout[0][0]     │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ batch_normalizatio… │ (None, 71, 64)    │        256 │ conv1d_1[0][0]    │
│ (BatchNormalizatio… │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_1 (Dropout) │ (None, 71, 64)    │          0 │ batch_normalizat… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ average_pooling1d   │ (None, 35, 64)    │          0 │ dropout_1[0][0]   │
│ (AveragePooling1D)  │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ attention           │ (None, 35, 64)    │          0 │ average_pooling1… │
│ (Attention)         │                   │            │ average_pooling1… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ flatten (Flatten)   │ (None, 2240)      │          0 │ attention[0][0]   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ peptide_length_inp… │ (None, 1)         │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ mhc_length_input    │ (None, 1)         │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ concatenate         │ (None, 2242)      │          0 │ flatten[0][0],    │
│ (Concatenate)       │                   │            │ peptide_length_i… │
│                     │                   │            │ mhc_length_input… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ sequential          │ (None, 128)       │    287,616 │ concatenate[0][0] │
│ (Sequential)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ sequential_1        │ (None, 128)       │     17,024 │ sequential[0][0]  │
│ (Sequential)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ sequential_2        │ (None, 128)       │     17,024 │ sequential_1[0][… │
│ (Sequential)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ sequential_3        │ (None, 128)       │     17,024 │ sequential_2[0][… │
│ (Sequential)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ sequential_4        │ (None, 128)       │     17,024 │ sequential_3[0][… │
│ (Sequential)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dense_5 (Dense)     │ (None, 1)         │        129 │ sequential_4[0][… │
└─────────────────────┴───────────────────┴────────────┴───────────────────┘

 Total params: 790,977 (3.02 MB)

 Trainable params: 789,313 (3.01 MB)

 Non-trainable params: 1,664 (6.50 KB)

Training with max 100 epochs, with early stopping and performance scheduler callbacks:

# Performance scheduling of learning rate
lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5)

# Early stopping
early_stopping = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
cnn_model.compile(loss='mse', optimizer=optimizer, metrics=['RootMeanSquaredError', 'R2Score'])
fit_history = cnn_model.fit(train_tfds, epochs=100, validation_data=valid_tfds,
                            callbacks=[lr_scheduler, early_stopping])

Regular C3 split training metrics:

Strict C3 split training metrics:

3. Comparing Regular vs. Strict C3 split

We will now compare the performance of the model trained on the regular C3 split and the strict C3 split using an independent test set. We will evaluate the models using metrics such as score, root mean squared error (RMSE), and loss.

cnn_model.evaluate(X_test_inputs, y_test)

Regular C3 split test metrics:

300/300 ━━━━━━━━━━━━━━━━━━━━ 4s 8ms/step - R2Score: 0.6597 - RootMeanSquaredError: 0.1524 - loss: 0.0232
[0.023598162457346916, 0.15361693501472473, 0.6540257930755615]

Strict C3 split test metrics:

300/300 ━━━━━━━━━━━━━━━━━━━━ 5s 13ms/step - R2Score: 0.1770 - RootMeanSquaredError: 0.2391 - loss: 0.0572
[0.05780040845274925, 0.24041715264320374, 0.18053573369979858]

As we can see, the model trained on the regular C3 split performs significantly better on the test set compared to the model trained on the strict C3 split. The score is much higher and the RMSE is much lower for the regular C3 split. This suggests that although the model generalizes well on predicting binding peptides of seen MHC2 alleles, it struggles to generalize to unseen MHC2 alleles when using the strict C3 split.

Let’s take a look at the prediction plots:

Regular C3 split prediction plot for all MHC2:

Strict C3 split prediction plot for all MHC2:

Regular C3 split prediction plot separated by MHC2:

Strict C3 split prediction plot separated by MHC2:

As we can see, while the regular C3 split shows good performance across multiple MHC2 alleles, with some achieving an score > 0.8, the strict C3 split shows max score of around 0.2 for the best performing MHC2 allele.

Notably, although the total test sample numbers are the same between the two split, due to the strict nature of the strict C3 split, the variety of MHC2 alleles in the test set is much smaller compared to the regular C3 split. Thus, we are perhaps being a little “strict” in our evaluation, since the strict C3 split has not been evaluated on many MHC2 alleles.

This result suggests important limitations for this model: it predicts well on MHC2 alleles that are present in the training data, but it does not perform well on unseen or novel MHC2 alleles. Perhaps the model is very good at memorizing specific interactions between peptides and MHC2 alleles, but did not learn well the general rules of protein-protein interaction, which is also a significant challenge for many PPI prediction models utilizing pLLMs. Some potential solutions to this issue could include:

• Increasing the diversity of the training data to include a wider range of MHC2 alleles, which may help the model learn more generalizable patterns.
• Incorporating additional features or using more complex model architectures that can capture the underlying biology of peptide binding to MHC2 molecules, rather than relying solely on the embeddings from the pre-trained pLLM.
• Re-train a pLLM on strict data with held out samples, so that the embeddings themselves are less prone to data leakage and more generalizable to unseen data.

Conclusion

In this project, I explored the issue of data leakage in peptide-MHC2 binding prediction models using pre-trained protein language models (pLLMs). I compared two data splitting strategies, a regular C3 split and a strict C3 split, to evaluate the extent of data leakage and its impact on model performance. My findings showed that while the model trained on the regular C3 split performed well on the test set, it struggled to generalize to unseen MHC2 alleles when evaluated using the strict C3 split. This highlights the importance of using stringent 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.

References

• Park, Y., Marcotte, E. Flaws in evaluation schemes for pair-input computational predictions. Nat Methods 9, 1134–1136 (2012).
• Szymborski, Joseph, and Amin Emad. “Data for”A Flaw in Using Pre-trained pLMs in Protein-protein Interaction Inference Models””. bioRxiv, December 5, 2025.
• Lin Z, Akin H, Rao R, Hie B, Zhu Z, Lu W, Smetanin N, Verkuil R, Kabeli O, Shmueli Y, Dos Santos Costa A, Fazel-Zarandi M, Sercu T, Candido S, Rives A. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science. 2023 Mar 17;379(6637):1123-1130. doi: 10.1126/science.ade2574. Epub 2023 Mar 16. PMID: 36927031.

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:

1. 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.

2. 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.

3. 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.

4. 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:

all_dat['train'].head()

Output:

   Peptide_ID           Peptide     MHC_ID  \
0      104653   VAPIEHIASMRRNYF  DRB1_1302   
1       37106   HDDKETSFIRNCARK  DRB1_0101   
2      118433   LIWVGINTRNMTMSM  DRB1_0101   
3       80770  GVTVIKNNMINNDLGP  DRB1_1501   
4       19888   PAPMLAAAAGWQTLS  DRB1_1101   

                                  MHC         Y  
0  QEFFIASGAAVDAIMESSFDYFDIDEATYHVGFT  0.501084  
1  QEFFIASGAAVDAIMWLFLECYDLQRATYHVGFT  0.441298  
2  QEFFIASGAAVDAIMWLFLECYDLQRATYHVGFT  0.217673  
3  QEFFIASGAAVDAIMWPRFDYFDIQAATYHVVFT  0.807811  
4  QEFFIASGAAVDAIMESSFDYFDFDRATYHVGFT  0.583271  

all_dat is a dictionary containing the training, validation, and test datasets as dataframes, each with peptide and MHC2 sequences, and binding affinity Y.

The embeddings extraction code to retain the 2D structure of the embeddings:

from transformers import AutoTokenizer, EsmModel
import numpy as np
import torch
import math
from tqdm import tqdm

# Load pre-trained ESM2 model and tokenizer
model_checkpoint = "facebook/esm2_t30_150M_UR50D" # 640 dim
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = EsmModel.from_pretrained(model_checkpoint)

# Function to get embeddings
def extract_full_embedding(
    sequence: list[str],
    tokenizer: AutoTokenizer,
    model: EsmModel,
    device: torch.device | None = None,
    batch_size: int = 64,
    max_len: int = 50
    ) -> np.ndarray:

    """Extract full embeddings for peptide sequences from an LLM model.
       Batch iteration is required as this is memory intensive."""

    # Use GPU when available
    if not device:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Separate sequence list into batches
    n_batches = math.ceil(len(sequence) / batch_size)
    all_batch_embeddings = []

    model = model.to(device) # Move model to the target device
    model.eval() # Set model to evaluation mode

    steps = tqdm(range(n_batches))

    for i in steps:
        steps.set_description(f"Processing batch {i+1}/{n_batches}")
        start = i * batch_size
        end = (i + 1) * batch_size
        batch = sequence[start:end]

        # Tokenize peptide sequence and pad to max length
        inputs = tokenizer(batch, return_tensors="pt", padding='max_length', truncation=True, max_length=max_len)

        # Move input to the target device
        inputs = {name: tensor.to(device) for name, tensor in inputs.items()}

        # Forward pass through the model without gradient tracking to get embeddings
        with torch.no_grad():
          batch_mean_embeddings = model(**inputs).last_hidden_state.detach().cpu().numpy()
        all_batch_embeddings.append(batch_mean_embeddings)

    embeddings = np.concatenate(all_batch_embeddings, axis=0)
    # This will return a (N_samples, max_len, n_dim) array
    return embeddings

Since each peptide and MHC2 sequence can have different lengths, we need to pad them to a consistent length during the embedding extraction so that they can be concatenated later. Let’s get the max sequence lengths for peptides or MHC2 sequences so we can pad the embeddings accordingly.

length_pt = []
for df in all_dat.values():
  length_pt.append(df['Peptide'].apply(len).max())
max_len_pt = max(length_pt)

length_mhc = []
for df in all_dat.values():
  length_mhc.append(df['MHC'].apply(len).max())
max_len_mhc = max(length_mhc)

print(max_len_pt, max_len_mhc)

Output:

37 34

The maximum peptide length is 37, and the maximum MHC2 length is 34.

Now we can extract the full embeddings for both peptides and MHC2 sequences, and pad them each to their respective max lengths.

embedding_dict_pt = {}
for data, sequence in sequence_dict_pt.items():
    print(f"Extracting embeddings for {data}...")
    embedding_dict_pt[data] = extract_full_embedding(sequence, tokenizer, model, batch_size=1280, max_len=max_len_pt)
    
embedding_dict_mhc = {}
for data, sequence in sequence_dict_mhc.items():
    print(f"Extracting embeddings for {data}...")
    embedding_dict_mhc[data] = extract_full_embedding(sequence, tokenizer, model, batch_size=1280, max_len=max_len_mhc)

Concatenate the peptide and MHC2 embeddings along the sequence dimension (axis=1) to create the final input embeddings for our model.

embedding_dict = {}
for data_type in ['train', 'valid', 'test']:
  pt_embeddings = embedding_dict_pt[f'{data_type}_pt']
  mhc_embeddings = embedding_dict_mhc[f'{data_type}_mhc']
  # Concatenate along the sequence length axis
  embedding_dict[data_type] = np.concatenate([pt_embeddings, mhc_embeddings], axis=1)

# Print shapes to verify
for key, value in embedding_dict.items():
  print(f"Shape of {key} combined embeddings: {value.shape}")

Output:

Shape of train combined embeddings: (82388, 71, 640)
Shape of valid combined embeddings: (11702, 71, 640)
Shape of test combined embeddings: (23410, 71, 640)

The output here is: (sample size, sequence_length, embedding_dimension).

Finally, we need to determine the actual lengths of each peptide and MHC2 sequence (before padding) to provide as additional inputs to the model:

# Determine each peptide or MHC sequence length to add into the embeddings for training input
seq_length_pt = {}
for key, sequence in sequence_dict_pt.items():
    seq_length_pt[key] = [len(seq) for seq in sequence]

seq_length_mhc = {}
for key, sequence in sequence_dict_mhc.items():
    seq_length_mhc[key] = [len(seq) for seq in sequence]

2. Model Architecture

Make the TensorFlow datasets from these 2D embeddings.

# Convert data to Tensorflow dataset
import tensorflow as tf

# Embeddings (already concatenated along sequence length axis)
X_train_emb = embedding_dict['train']
X_valid_emb = embedding_dict['valid']
X_test_emb = embedding_dict['test']

# Target variable
y_train = all_dat['train']['Y'].values
y_valid = all_dat['valid']['Y'].values
y_test = all_dat['test']['Y'].values

# Sequence lengths (converting lists to numpy arrays and then to tf tensors)
X_train_pt_len = tf.convert_to_tensor(np.array(seq_length_pt['train_pt']), dtype=tf.float32)
X_train_mhc_len = tf.convert_to_tensor(np.array(seq_length_mhc['train_mhc']), dtype=tf.float32)
X_valid_pt_len = tf.convert_to_tensor(np.array(seq_length_pt['valid_pt']), dtype=tf.float32)
X_valid_mhc_len = tf.convert_to_tensor(np.array(seq_length_mhc['valid_mhc']), dtype=tf.float32)
X_test_pt_len = tf.convert_to_tensor(np.array(seq_length_pt['test_pt']), dtype=tf.float32)
X_test_mhc_len = tf.convert_to_tensor(np.array(seq_length_mhc['test_mhc']), dtype=tf.float32)

# Create tf.data.Dataset with multiple inputs
train_tfds = tf.data.Dataset.from_tensor_slices(((X_train_emb, X_train_pt_len, X_train_mhc_len), y_train))
# Shuffle, batch and prefetch the train tfds
train_tfds = train_tfds.shuffle(1024, seed=42).batch(32).prefetch(tf.data.AUTOTUNE)

valid_tfds = tf.data.Dataset.from_tensor_slices(((X_valid_emb, X_valid_pt_len, X_valid_mhc_len), y_valid))
valid_tfds = valid_tfds.batch(32).prefetch(tf.data.AUTOTUNE)

# For evaluation on X_test, we will need to ensure it's a tuple of tensors
X_test_inputs = (X_test_emb, X_test_pt_len, X_test_mhc_len)

# For X_train_inputs, we will need to ensure it's a tuple of tensors for R2 evaluation
X_train_inputs = (X_train_emb, X_train_pt_len, X_train_mhc_len)

Define the CNN + MLP model architecture. I’m using average pooling rather than max pooling as I find it tends to work better for this task, likely because it captures the overall presence of features rather than just the strongest activation.

# Define 1D CNN + MLP with multiple inputs using functional API
import tensorflow as tf
tf.keras.backend.clear_session()

tf.random.set_seed(42)

def make_dense_block(n_neurons, dropout_rate=0.2):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(n_neurons, activation='relu', kernel_initializer='he_normal'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(dropout_rate)
    ])

# Define the input layers
embedding_input = tf.keras.layers.Input(shape=embedding_dict['train'].shape[1:], name='embedding_input')
peptide_length_input = tf.keras.layers.Input(shape=(1,), name='peptide_length_input')
mhc_length_input = tf.keras.layers.Input(shape=(1,), name='mhc_length_input')

# CNN branch for embeddings
x = tf.keras.layers.Conv1D(filters=128, kernel_size=5, padding='same', activation='relu')(embedding_input)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.Conv1D(filters=64, kernel_size=3, padding='same', activation='relu')(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.AveragePooling1D(pool_size=2)(x)
x = tf.keras.layers.Flatten()(x)

# Concatenate CNN output with length inputs
combined_features = tf.keras.layers.Concatenate()([x, peptide_length_input, mhc_length_input])

# MLP branch
y = make_dense_block(128)(combined_features)
y = make_dense_block(128)(y)
y = make_dense_block(128)(y)
y = make_dense_block(128)(y)
y = make_dense_block(128)(y)
output = tf.keras.layers.Dense(1)(y)

# Create the functional model
cnn_model = tf.keras.Model(inputs=[embedding_input, peptide_length_input, mhc_length_input], outputs=[output])

cnn_model.summary()

As we can see below, under the CNN architecture, the filter effectively compresses the embedding dimension from 640 down to 128, and then to 64. We set the padding to ‘same’ to retain the original sequence length of 71 after convolution. But if set to ‘valid’, the sequence length would reduce after each convolution layer, depending on the kernel size. The average pooling layer then reduces the sequence length dimension according to the pool size of 2 (basically halves it), and finally, we flatten the CNN output before concatenating it with the peptide and MHC2 lengths. This results in a feature vector of size 64 * 35 = 2240 from the CNN branch, which is then concatenated with the two length inputs, giving a total of 2242 features fed into the dense layers. Ultimately, the pooling layers help to prevent overfitting and save computational resources by reducing the dimensionality of the feature maps.

Compare this with the MLP-only model architecture that we will compare against:

# Define MLP layers
# Define a helper function for the repetitive dense -> batch norm -> dropout block
def make_dense_block(n_neurons, dropout_rate=0.2):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(n_neurons, activation='relu', kernel_initializer='he_normal'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(dropout_rate)
    ])

tf.random.set_seed(42)

dnn_model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=X_train.shape[1:]),
    make_dense_block(512),
    make_dense_block(256),
    make_dense_block(128),
    make_dense_block(128),
    make_dense_block(128),
    tf.keras.layers.Dense(1)
])

3. Training and Evaluation

Set performance scheduling, early stopping, optimizer, compile model, and train the CNN + MLP model. The MLP-only model is trained with the same settings for comparison.

# Performance scheduling of learning rate
lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5)

# Early stopping
early_stopping = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)

# Define optimizer
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)

# Compile and train the model
cnn_model.compile(loss='mse', optimizer=optimizer, metrics=['RootMeanSquaredError', 'R2Score'])
fit_history = cnn_model.fit(train_tfds, epochs=100, validation_data=valid_tfds,
                            callbacks=[lr_scheduler, early_stopping])

Training CNN models is quite memory intensive. This specific data and model architecture required around 30-40 GB of GPU memory and 80-90 GB of system memory. With sufficient resources, the training completed in about 20 minutes per run on an NVIDIA A100 GPU.

Plot training metrics:

# Plot train and validation loss and RMSE across epochs
import matplotlib.pyplot as plt

fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(8, 8))

pd.DataFrame(fit_history.history)[['loss', 'val_loss']].plot(
    grid=True, xlabel="Epoch", ax=ax[0, 0],
    style=["r--", "b-"])

pd.DataFrame(fit_history.history)[['RootMeanSquaredError', 'val_RootMeanSquaredError']].plot(
    grid=True, xlabel="Epoch", ax=ax[0, 1],
    style=["r--", "b-"])

pd.DataFrame(fit_history.history)[['R2Score', 'val_R2Score']].plot(
    grid=True, xlabel="Epoch", ax=ax[1, 0],
    style=["r--", "b-"])

pd.DataFrame(fit_history.history)[['learning_rate']].plot(
    grid=True, xlabel="Epoch", ax=ax[1, 1],
    style=["g-"])

ax[0, 0].set_ylabel('Loss')
ax[0, 0].set_title('Loss Over Epochs')
ax[0, 0].legend(['Training Loss', 'Validation Loss'])

ax[0, 1].set_ylabel('RMSE')
ax[0, 1].set_title('RMSE Over Epochs')
ax[0, 1].legend(['Training RMSE', 'Validation RMSE'])

ax[1, 0].set_ylabel('R2')
ax[1, 0].set_title('R2 Over Epochs')
ax[1, 0].legend(['Training R2', 'Validation R2'])

ax[1, 1].set_ylabel('Learning Rate')
ax[1, 1].set_title('Learning Rate Over Epochs')

plt.tight_layout()
plt.show()

The CNN + MLP model trained for the full 100 epochs, while the MLP-only model stopped early at epoch 80. Both models showed good convergence without overfitting.

Evaluate the model on the test set:

cnn_model.evaluate(X_test_inputs, y_test)

Test Results:
- R2Score: 0.6219
- RootMeanSquaredError: 0.1606
- loss: 0.0258

For comparison, the MLP-only model achieved the following test results:
- R2Score: 0.5452
- RootMeanSquaredError: 0.1775
- loss: 0.0315

This indicates that, in this specific setting, incorporating CNN layers improved the model’s performance in predicting protein binding affinity.

Let’s plot a comparison of MLP-only vs CNN + MLP prediction accuracy on the test set, stratified by MHC2 alleles. Higher R2Score indicates better prediction performance.

From the comparison plot, we can see that the CNN + MLP model generally provides better predictions across various MHC2 alleles compared to the MLP-only model.

Let’s look at one allele example on how well the CNN + MLP model predicts the affinity of HLA-DPA10201-DPB10101:

As we can see, the = 0.8, which is an improvement over the MLP-only model of 0.72.

It is interesting to note that the MHC allele sample number does not correlate with the prediction performance, indicating that the model has learned meaningful patterns rather than just memorizing frequent alleles.

4. Saliency Map and In Silico Mutagenesis Analysis

To compute the saliency maps, we will calculate the gradients of the model output with respect to the input embeddings. This will help us identify which amino acid positions in the peptide and MHC2 sequences are most influential for the binding affinity predictions.

We will first randomly select 5 high-affinity samples from the test set for saliency analysis:

# Select 5 samples from the test data where Y > 0.9
high_affinity_data = all_dat['test'][all_dat['test']['Y'] > 0.9].sample(5, random_state=42)

# Get the indices of these samples
sample_indices = high_affinity_data.index.values

print(f"Selected {len(high_affinity_data)} samples with Y > 0.9 for saliency analysis.")
print(high_affinity_data[['Peptide', 'MHC_ID', 'MHC', 'Y']])

Output:

Selected 5 samples with Y > 0.9 for saliency analysis.
               Peptide                 MHC_ID  \
21199  EKKYYAATQFEPLAA  HLA-DPA10301-DPB10402   
5997   AFLIGANYLGKPKEQ              DRB1_0101   
22257  DKRLAAYLMLMRSPS              DRB1_1501   
8396   SQVNPITLTAALLLL              DRB1_0701   
17575  NDKFTVFEGAFNKAI              DRB5_0101   

                                      MHC         Y  
21199  YMFFMFSGGAISNTLFGQFEYFDIEKVRMHLGMT  0.915458  
5997   QEFFIASGAAVDAIMWLFLECYDLQRATYHVGFT  1.000000  
22257  QEFFIASGAAVDAIMWPRFDYFDIQAATYHVVFT  1.000000  
8396   QEFFIASGAAVDAIMWGYFELYVIDRQTVHVGFT  0.956843  
17575  QEFFIASGAAVDAIMQDYFHDYDFDRATYHVGFT  1.000000  

Now, we will extract the embeddings for these high-affinity samples:

# Get embeddings for these high affinity data
high_affinity_peptide = high_affinity_data['Peptide'].tolist()
high_affinity_mhc = high_affinity_data['MHC'].tolist()

max_len_pt = 37
max_len_mhc = 34

peptide_embeddings = extract_full_embedding(high_affinity_peptide, tokenizer, model, batch_size=1280, max_len=max_len_pt)
mhc_embeddings = extract_full_embedding(high_affinity_mhc, tokenizer, model, batch_size=1280, max_len=max_len_mhc)

# Concatenate the embeddings
X_high_affinity_embeddings = np.concatenate([peptide_embeddings, mhc_embeddings], axis=1)

# Extract corresponding sequence lengths
X_high_affinity_pt_len = high_affinity_data['Peptide'].apply(len).values
X_high_affinity_mhc_len = high_affinity_data['MHC'].apply(len).values

Define a helper function to compute saliency maps:

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

def calculate_saliency_map(model: tf.keras.Model,
                           embedding_data: np.ndarray,
                           peptide_len_data: np.ndarray,
                           mhc_len_data: np.ndarray,
                           target_idx: int = 0) -> np.ndarray:

    """Calculates saliency map for a given input using the trained model."""

    # Convert numpy arrays to TensorFlow tensors and ensure float32
    input_embeddings_tensor = tf.convert_to_tensor(embedding_data, dtype=tf.float32)
    peptide_len_tensor = tf.convert_to_tensor(peptide_len_data, dtype=tf.float32)
    mhc_len_tensor = tf.convert_to_tensor(mhc_len_data, dtype=tf.float32)

    with tf.GradientTape() as tape:
        tape.watch(input_embeddings_tensor) # Watch only the embeddings for saliency
        
        # Pass all inputs to the model for prediction of Y
        predictions = model((input_embeddings_tensor, peptide_len_tensor, mhc_len_tensor))
        
        # Select the target output if there are multiple outputs or a specific class
        target_output = predictions[:, target_idx]

    # Calculate gradients of the target output with respect to the input embeddings
    gradients = tape.gradient(target_output, input_embeddings_tensor)

    # Take the mean across the embedding dimension to get a single saliency score per amino acid position
    saliency = tf.reduce_mean(gradients, axis=-1) # Take mean along last axis: embedding_dim

    return saliency.numpy().flatten()

# Package inputs for the model's predict method
X_high_affinity_inputs = (X_high_affinity_embeddings, X_high_affinity_pt_len, X_high_affinity_mhc_len)

# Predictions for selected samples
y_pred_high_affinity = cnn_model.predict(X_high_affinity_inputs)

Compute and plot the saliency maps for each selected high-affinity sample:

fig, axes = plt.subplots(nrows=len(high_affinity_data), ncols=2, figsize=(18, len(high_affinity_data) * 3))
fig.suptitle('Saliency Maps for High Affinity Peptide-MHC Pairs (Y > 0.9)', fontsize=16)

# Ensure axes is always 2D even for a single row
if len(high_affinity_data) == 1:
    axes = np.array([axes])

for i in range(len(high_affinity_data)):
    sample_info = high_affinity_data.iloc[i]
    peptide_sequence = sample_info['Peptide']
    mhc_sequence = sample_info['MHC']
    mhc_id = sample_info['MHC_ID']
    y_true_sample = sample_info['Y']
    y_pred_sample = y_pred_high_affinity[i][0] # Access pre-calculated prediction

    # Get one sample's embeddings and lengths at a time
    current_embedding_data = X_high_affinity_embeddings[i:i+1]
    current_peptide_len_data = X_high_affinity_pt_len[i:i+1]
    current_mhc_len_data = X_high_affinity_mhc_len[i:i+1]

    # Calculate saliency map using the updated function signature
    # Normalize with a constant to enhance visualization
    beta = 1e5
    saliency_scores = calculate_saliency_map(cnn_model, 
                                             current_embedding_data, 
                                             current_peptide_len_data, 
                                             current_mhc_len_data) * beta

    peptide_len = len(peptide_sequence)
    mhc_len = len(mhc_sequence)

    # Correctly slice saliency scores based on actual sequence lengths
    peptide_saliency = saliency_scores[:peptide_len]
    mhc_saliency = saliency_scores[max_len_pt : max_len_pt + mhc_len]

    annot_settings = {"rotation": 90}

    # Plot for Peptide Saliency
    sns.heatmap(
        [peptide_saliency],
        cmap='coolwarm', 
        annot=True,
        annot_kws=annot_settings,
        fmt=".2f",
        xticklabels=list(peptide_sequence),
        yticklabels=['Peptide'],
        cbar_kws={'label': 'Saliency Score'},
        ax=axes[i, 0]
    )
    axes[i, 0].set_title(f'Sample {i+1} - Peptide: {peptide_sequence}\nTrue Y: {y_true_sample:.2f}, Pred Y: {y_pred_sample:.2f}')
    axes[i, 0].tick_params(axis='x', rotation=0)

    # Plot for MHC Saliency
    sns.heatmap(
        [mhc_saliency],
        cmap='coolwarm', 
        annot=True,
        annot_kws=annot_settings,
        fmt=".2f",
        xticklabels=list(mhc_sequence),
        yticklabels=['MHC'],
        cbar_kws={'label': 'Saliency Score'},
        ax=axes[i, 1]
    )
    axes[i, 1].set_title(f'Sample {i+1} - MHC: {mhc_sequence}\nMHC ID: {mhc_id}')
    axes[i, 1].tick_params(axis='x', rotation=0)

plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust layout to make space for suptitle
plt.show()

Positive saliency scores (red) indicate that the amino acids are positively correlated with a higher binding affinity, while negative scores (blue) suggest that those amino acids contribute to lower binding affinity.

Next, let’s perform In Silico Mutagenesis (ISM) analysis to see how single amino acid mutations affect the predicted binding affinity. Instead of mutating every position, which will be computationally expensive, we will focus on mutating only the most salient positions (highest saliency score) in the peptide sequences for each sample.

First, generate a dictionary of the most disruptive mutations for each amino acid:

# A dictionary that suggests the most disruptive mutagenesis for each a.a.
disruptive_mutations = {
    'A': 'W', # Small nonpolar -> Bulkiest nonpolar/aromatic
    'R': 'D', # Positively charged -> Negatively charged
    'N': 'F', # Polar uncharged -> Hydrophobic aromatic
    'D': 'R', # Negatively charged -> Positively charged
    'C': 'S', # Disulfide bond former -> Similar size but non-disulfide/polar
    'E': 'K', # Negatively charged -> Positively charged
    'Q': 'W', # Polar uncharged -> Bulkiest nonpolar/aromatic
    'G': 'P', # Flexible -> Conformational constraint
    'H': 'D', # Positively charged -> Negatively charged (pH depending)
    'I': 'E', # Hydrophobic -> Negatively charged hydrophilic
    'L': 'K', # Hydrophobic -> Positively charged hydrophilic
    'K': 'D', # Positively charged -> Negatively charged
    'M': 'E', # Nonpolar/sulfur -> Negatively charged hydrophilic
    'F': 'E', # Hydrophobic aromatic -> Negatively charged hydrophilic
    'P': 'G', # Conformational constraint -> Flexible
    'S': 'L', # Polar uncharged -> Hydrophobic
    'T': 'I', # Polar uncharged -> Hydrophobic
    'W': 'A', # Bulkiest nonpolar -> Smallest nonpolar
    'Y': 'D', # Aromatic/polar -> Negatively charged
    'V': 'E'  # Hydrophobic -> Negatively charged hydrophilic
}

Perform ISM on the top salient positions in the peptide sequences:

# Find aa with the highest saliency score across all samples
max_scores_aa_pos_pt = {}
for i in range(len(high_affinity_data)):
  sample_info = high_affinity_data.iloc[i]
  peptide_sequence = sample_info['Peptide']
  current_embedding_data = X_high_affinity_embeddings[i:i+1]
  current_peptide_len_data = X_high_affinity_pt_len[i:i+1]
  current_mhc_len_data = X_high_affinity_mhc_len[i:i+1]
  saliency = calculate_saliency_map(cnn_model, current_embedding_data, current_peptide_len_data, current_mhc_len_data)
  max_score_pos = saliency[:current_peptide_len_data[0]].argmax().tolist()
  max_score_aa = peptide_sequence[max_score_pos]
  max_scores_aa_pos_pt[sample_info['Peptide_ID'].tolist()] = {max_score_aa: max_score_pos}

print("Identified highest saliency amino acids and their positions:")
display(max_scores_aa_pos_pt)

Output:

Identified highest saliency amino acids and their positions:
{60550: {'L': 12},
 117646: {'L': 2},
 51178: {'P': 13},
 73939: {'A': 10},
 26350: {'F': 3}}

mutated_peptide_predictions = []

for i, (peptide_id, saliency_info) in enumerate(max_scores_aa_pos_pt.items()):
    # Get original sample info
    sample_row = high_affinity_data[high_affinity_data['Peptide_ID'] == peptide_id].iloc[0]
    original_peptide = sample_row['Peptide']
    original_mhc = sample_row['MHC']
    original_y_pred = y_pred_high_affinity[high_affinity_data['Peptide_ID'] == peptide_id].flatten()[0]

    # Get the AA with highest saliency and its position
    saliency_aa = list(saliency_info.keys())[0]
    saliency_pos = list(saliency_info.values())[0]

    # Determine the disruptive mutation
    disruptive_aa = disruptive_mutations.get(saliency_aa, 'A') # Default to Alanine if not in dict

    # Create mutated peptide sequence
    mutated_peptide_sequence_list = list(original_peptide)
    mutated_peptide_sequence_list[saliency_pos] = disruptive_aa
    mutated_peptide = ''.join(mutated_peptide_sequence_list)

    # --- Extract embeddings for the mutated peptide ---
    # Ensure original MHC embeddings are reused to isolate peptide effect
    original_mhc_embeddings = extract_full_embedding([original_mhc], tokenizer, model, batch_size=1, max_len=max_len_mhc)

    # Extract embeddings for the mutated peptide
    mutated_peptide_embeddings = extract_full_embedding([mutated_peptide], tokenizer, model, batch_size=1, max_len=max_len_pt)

    # Concatenate mutated peptide embeddings with original MHC embeddings
    X_mutated_combined_embeddings = np.concatenate([mutated_peptide_embeddings, original_mhc_embeddings], axis=1)

    # Get sequence lengths for mutated peptide and original MHC
    mutated_peptide_len = np.array([len(mutated_peptide)])
    original_mhc_len = np.array([len(original_mhc)])

    # Package inputs for prediction
    X_mutated_inputs = (
        X_mutated_combined_embeddings,
        tf.convert_to_tensor(mutated_peptide_len, dtype=tf.float32),
        tf.convert_to_tensor(original_mhc_len, dtype=tf.float32)
    )

    # Predict binding affinity for the mutated peptide
    mutated_y_pred = cnn_model.predict(X_mutated_inputs).flatten()[0]

    mutated_peptide_predictions.append({
        'Peptide_ID': peptide_id,
        'Original_Peptide': original_peptide,
        'Mutated_Peptide': mutated_peptide,
        'Saliency_AA': saliency_aa,
        'Saliency_Pos': saliency_pos + 1,
        'Disruptive_AA': disruptive_aa,
        'Original_Pred_Y': original_y_pred,
        'Mutated_Pred_Y': mutated_y_pred,
        'Change_in_Y': mutated_y_pred - original_y_pred
    })

mutated_predictions_df = pd.DataFrame(mutated_peptide_predictions)

print(mutated_predictions_df)

Output:

   Peptide_ID Original_Peptide  Mutated_Peptide Saliency_AA  Saliency_Pos  \
0       60550  EKKYYAATQFEPLAA  EKKYYAATQFEPKAA           L            13   
1      117646  AFLIGANYLGKPKEQ  AFKIGANYLGKPKEQ           L             3   
2       51178  DKRLAAYLMLMRSPS  DKRLAAYLMLMRSGS           P            14   
3       73939  SQVNPITLTAALLLL  SQVNPITLTAWLLLL           A            11   
4       26350  NDKFTVFEGAFNKAI  NDKETVFEGAFNKAI           F             4   

  Disruptive_AA  Original_Pred_Y  Mutated_Pred_Y  Change_in_Y  
0             K         0.879057        0.721061    -0.157996  
1             K         0.782684        0.729118    -0.053567  
2             G         0.652823        0.643233    -0.009590  
3             W         0.665965        0.581599    -0.084366  
4             E         0.678189        0.355576    -0.322613  

Plot the changes in predicted binding affinity due to the mutations:

import matplotlib.pyplot as plt
import seaborn as sns
from adjustText import adjust_text

plt.figure(figsize=(8, 6))
sns.scatterplot(
    data=mutated_predictions_df,
    x='Original_Pred_Y',
    y='Mutated_Pred_Y',
    hue='Change_in_Y', 
    s=200,
    palette='coolwarm', 
    legend='full'
)

# Add a diagonal line for reference (where Y_original == Y_mutated)
plt.plot([0, 1], [0, 1], 'k--', lw=2, alpha=0.5, label='No Change')

# Annotate points with Peptide_ID or a combination of info
texts = []
for idx, row in mutated_predictions_df.iterrows():
    label = f"{row['Original_Peptide']}\n{row['Saliency_AA']}{row['Saliency_Pos']}->{row['Disruptive_AA']}"
    texts.append(plt.text(row['Original_Pred_Y'], row['Mutated_Pred_Y'], label, fontsize=10))

adjust_text(texts, arrowprops=dict(arrowstyle="-", color='gray', lw=1))

plt.title('Predicted Binding Affinity: Original vs. Mutated Peptides')
plt.xlabel('Original Predicted Y')
plt.ylabel('Mutated Predicted Y')
plt.grid(True, linestyle='--', alpha=0.6)
plt.xlim(0.3, 1)
plt.ylim(0.3, 1)
plt.axhline(y=0.5, color='gray', linestyle=':', lw=1, alpha=0.7)
plt.axvline(x=0.5, color='gray', linestyle=':', lw=1, alpha=0.7)
plt.legend(title='Change in Y')
plt.show()

As we can see from the ISM analysis, mutating the most salient amino acid positions in the peptide sequences generally leads to a decrease in predicted binding affinity, as indicated by the points falling below the diagonal line. Some mutations result in significant drops in predicted affinity, while others are less impactful. However, we are only mutating a single position here, and multiple mutations may have a more pronounced effect. This analysis provides a hypothesis for experimental validation of key residues involved in peptide-MHC binding.

Conclusion

Incorporating convolutional neural networks (CNNs) into the protein binding prediction model improved performance compared to a MLP-only approach. By retaining the 2D structure of the ESM2 embeddings and applying 1D convolutional layers, the model was able to capture local sequence patterns that are important for binding affinity prediction. The CNN + MLP model achieved higher R2 scores and lower RMSE on the test set, demonstrating its effectiveness.

It is important to note that further hyperparameter tuning for either model could yield different results. Additionally, exploring other architectures, such as attention mechanisms, could provide additional insights into modeling protein binding affinity.

References

• ESM2
• Hands-on Machine Learning with Scikit-Learn, Keras, and TensorFlow
• Deep Learning for Biology

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

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

Summary Figure

1. Data Selection and Splitting

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

# Access the MHC2_IEDB_Jensen data from TDC API and obtain splitted data
from tdc.multi_pred import PeptideMHC

data = PeptideMHC(name = "MHC2_IEDB_Jensen")
split = data.get_split(method = "random", seed = 524, frac = [0.7, 0.1, 0.2])
train_df, valid_df, test_df = split['train'], split['valid'], split['test']

# Save the files to drive
import os
save_path = './data'
os.makedirs(save_path, exist_ok=True)

train_df.to_feather(os.path.join(save_path, 'train_dat.feather'))
valid_df.to_feather(os.path.join(save_path, 'valid_dat.feather'))
test_df.to_feather(os.path.join(save_path, 'test_dat.feather'))

2. Exploratory Data Analysis

Let’s examine the structure and distribution of our data.

# how many data points in each data set
print(f"Number of data points in the training set: {len(train_df)}")
print(f"Number of data points in the validation set: {len(valid_df)}")
print(f"Number of data points in the test set: {len(test_df)}")

Output:

Number of data points in the training set: 93997
Number of data points in the validation set: 13428
Number of data points in the test set: 26856

A quick look at the training data reveals the structure: Peptide sequence, MHC sequence (pseudo-sequence), and the binding affinity Y.

# quick look at the train data structure
train_df.head()

# look at the distribution of Y in training data
import matplotlib.pyplot as plt
import seaborn as sns

plt.figure(figsize=(10, 6))
sns.histplot(train_df['Y'], bins=50, kde=True)
plt.title('Distribution of Y (Binding Affinity) in Training Data')
plt.xlabel('Binding Affinity (Y)')
plt.ylabel('Frequency')
plt.grid(True, linestyle='--', alpha=0.7)
plt.show()

# How many unique MHC_ID and peptides are there?
unique_mhc = train_df['MHC_ID'].nunique()
unique_peptides = train_df['Peptide'].nunique()
print(f'Number of unique MHC_IDs in training set: {unique_mhc}')
print(f'Number of unique peptides in training set: {unique_peptides}')

Number of unique MHC_IDs in training set: 79
Number of unique peptides in training set: 14597

# plot barplot of instances per MHC2 molecule
plt.figure(figsize=(14, 6))
train_df["MHC_ID"].value_counts().sort_values(ascending = False).plot(kind = 'bar')

3. Extract and Merge LLM Embeddings

To represent the protein sequences for our neural network, we will use ESM-2 (Evolutionary Scale Modeling), a state-of-the-art protein language model. We’ll use the 150M parameter version (esm2_t30_150M_UR50D), which has 640 embedding dimensions.

# Only keep MHC_ID that has >= 1000 instances in train data, so that the model has sufficient training data
n_data = train_df["MHC_ID"].value_counts()
keep_mhc_id = n_data[n_data >= 1000].index.tolist()
train_df = train_df[train_df["MHC_ID"].isin(keep_mhc_id)].reset_index(drop=True)
valid_df = valid_df[valid_df["MHC_ID"].isin(keep_mhc_id)].reset_index(drop=True)
test_df = test_df[test_df["MHC_ID"].isin(keep_mhc_id)].reset_index(drop=True)

# How many data points in each data set
print(f"Number of data points in the training set: {len(train_df)}")
print(f"Number of data points in the validation set: {len(valid_df)}")
print(f"Number of data points in the test set: {len(test_df)}")

Output:

Number of data points in the training set: 82388
Number of data points in the validation set: 11702
Number of data points in the test set: 23410

from transformers import AutoTokenizer, EsmModel

# model checkpoints can be seen here: https://github.com/facebookresearch/esm#available-models-and-datasets-
# the one we use has medium complexity and has 640 dimensions

model_checkpoint = "facebook/esm2_t30_150M_UR50D"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = EsmModel.from_pretrained(model_checkpoint)

A note to LLM model selection for this specific task: I’ve tried larger ESM2 model with 1280 or 2560 embedding dimensions, and they did not improve the accuracy of binding prediction. As peptides and MHC lengths are quite short, it is possible that the 640 dimension model has already saturated learnable information for this task.

# Write a function to extract mean embeddings for protein sequences
# GPU will significantly speed up the process
# This is memory intensive, so we will batch the iteration
import numpy as np
from transformers import AutoTokenizer, EsmModel
import torch
import math
from tqdm import tqdm
import pandas as pd

def extract_mean_embedding(
    sequence: list[str],
    tokenizer: AutoTokenizer,
    model: EsmModel,
    device: torch.device | None = None,
    batch_size: int = 64
    ) -> pd.DataFrame:

    """Extract mean embeddings for peptide sequences from a LLM model."""

    # Use GPU when available
    if not device:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Separate sequence list into batches
    n_batches = math.ceil(len(sequence) / batch_size)
    all_batch_embeddings = []

    model = model.to(device) # Move model to the target device
    model.eval() # Set model to evaluation mode

    steps = tqdm(range(n_batches))

    for i in steps:
        steps.set_description(f"Processing batch {i+1}/{n_batches}")
        start = i * batch_size
        end = (i + 1) * batch_size
        batch = sequence[start:end]

        # Tokenize peptide sequence and pad to equal length
        inputs = tokenizer(batch, return_tensors="pt", padding=True)

        # Move input to the target device
        inputs = {name: tensor.to(device) for name, tensor in inputs.items()}

        # Forward pass through the model without gradient tracking to get mean embeddings
        with torch.no_grad():
          batch_mean_embeddings = model(**inputs).last_hidden_state.mean(dim=1).detach().cpu().numpy()
        all_batch_embeddings.append(batch_mean_embeddings)

    embeddings = pd.DataFrame(np.vstack(all_batch_embeddings))
    embeddings.columns = [f"me_{i + 1}" for i in range(embeddings.shape[1])]
    return embeddings

The ESM2 model returns a sequence of embeddings for each token in the input sequence. We take the mean of these embeddings to get a single embedding for the entire sequence. It is also possible to keep the sequence of embeddings for each token, and use other methods to aggregate the embeddings, such as max pooling or attention pooling, or train a 1D Convolutional Neural Network (CNN) to learn the optimal way to aggregate the embeddings. In this case, we will keep the mean embeddings for each sequence.

# Put all peptide/MHC data into a dictionary
sequence_dict = {
    "train_pt": train_df['Peptide'].tolist(),
    "valid_pt": valid_df['Peptide'].tolist(),
    "test_pt": test_df['Peptide'].tolist(),
    "train_mhc": train_df['MHC'].tolist(),
    "valid_mhc": valid_df['MHC'].tolist(),
    "test_mhc": test_df['MHC'].tolist()
}

# Extract embeddings for all peptide/MHC data
embedding_dict = {}
for data, sequence in sequence_dict.items():
    print(f"Extracting embeddings for {data}...")
    embedding_dict[data] = extract_mean_embedding(sequence, tokenizer, model, batch_size=2560)

There are other, more complex ways to combine embeddings from two interacting entities, but for this task, we will use the simple approach of concatenating the embeddings.

# Modify column names to add peptide/MHC prefix
for key, me in embedding_dict.items():
  me.columns = [f"{key.split('_')[1]}_{col}" for col in me.columns]

# Concatenate peptide and MHC mean embeddings and name columns
train_me = pd.concat([train_df[['Peptide', 'MHC_ID', 'Y']],
                      embedding_dict['train_pt'],
                      embedding_dict['train_mhc']], axis=1)
valid_me = pd.concat([valid_df[['Peptide', 'MHC_ID', 'Y']],
                      embedding_dict['valid_pt'],
                      embedding_dict['valid_mhc']], axis=1)
test_me = pd.concat([test_df[['Peptide', 'MHC_ID', 'Y']],
                     embedding_dict['test_pt'],
                     embedding_dict['test_mhc']], axis=1)

# Save data as feather files
model_cp = str(model.name_or_path).replace('/', '_')
df_list = [(train_me, 'train_me'), (valid_me, 'valid_me'), (test_me, 'test_me')]
save_path = './data'
for df, df_name in df_list:
  df.to_feather(os.path.join(save_path, f'{model_cp}_{df_name}.feather'))

Let’s visualize the embeddings using t-SNE to see if there are any patterns in the data.

# Plot t-SNE of sampled peptide-MHC mean embeddings
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE

sampled_index = train_me.sample(n=5000, random_state=1).index # sample 5000 data points for visualization
X_embedded = TSNE(n_components=2, learning_rate='auto',
                  init='random', perplexity=30).fit_transform(train_me.loc[sampled_index])

plt.figure(figsize=(10, 8))
plt.scatter(X_embedded[:, 0], X_embedded[:, 1], alpha=0.5, s=2)
plt.title('t-SNE of Peptide-MHC Embeddings')
plt.show()

In the above t-SNE plot, we can see some cluster structures correlating with higher Y, indicating that the embeddings encode information relating to peptide-MHC affinity.

4. Neural Network Training

We will build a MLP that takes the concatenated peptide and MHC embeddings as input, and outputs the predicted affinity.

# Convert data to Tensorflow dataset
import tensorflow as tf
import pandas as pd

# Compile all feature columns (peptide and MHC embeddings) into a single NumPy array for each dataset
X_train = train_me.filter(regex=f"^pt|^mhc").values
y_train = train_me['Y'].values

X_valid = valid_me.filter(regex=f"^pt|^mhc").values
y_valid = valid_me['Y'].values

X_test = test_me.filter(regex=f"^pt|^mhc").values
y_test = test_me['Y'].values

# Create tf.data.Dataset from these concatenated arrays
train_tfds = tf.data.Dataset.from_tensor_slices((X_train, y_train))

# Shuffle, batch and prefetch the tfds
train_tfds = train_tfds.shuffle(1024, seed=42).batch(32).prefetch(1)

valid_tfds = tf.data.Dataset.from_tensor_slices((X_valid, y_valid))
valid_tfds = valid_tfds.batch(32).prefetch(1)

# Define MLP layers
# Define a helper function for the repetitive dense -> batch norm -> dropout block
def make_dense_block(n_neurons, dropout_rate=0.2):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(n_neurons, activation='relu', kernel_initializer='he_normal'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(dropout_rate)
    ])

tf.random.set_seed(42)
norm_layer = tf.keras.layers.Normalization(input_shape=X_train.shape[1:])

dnn_model = tf.keras.Sequential([
    norm_layer,
    make_dense_block(300),
    make_dense_block(300),
    make_dense_block(300),
    make_dense_block(300),
    make_dense_block(300),
    tf.keras.layers.Dense(1)
])
dnn_model.summary()

# Performance scheduling of learning rate
lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5)

# Early stopping
early_stopping = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
dnn_model.compile(loss='mse', optimizer=optimizer, metrics=['RootMeanSquaredError', 'R2Score'])

# Adapt the normalization layer using only the feature tensors from the dataset
norm_layer.adapt(train_tfds.map(lambda x, y: x))
tf.random.set_seed(42)
fit_history = dnn_model.fit(train_tfds, epochs=50, validation_data=valid_tfds,
                            callbacks=[lr_scheduler, early_stopping])

# Plot train and validation loss and RMSE across epochs
import matplotlib.pyplot as plt

fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(8, 8))

pd.DataFrame(fit_history.history)[['loss', 'val_loss']].plot(
    grid=True, xlabel="Epoch", ax=ax[0, 0],
    style=["r--", "b-"])

pd.DataFrame(fit_history.history)[['RootMeanSquaredError', 'val_RootMeanSquaredError']].plot(
    grid=True, xlabel="Epoch", ax=ax[0, 1],
    style=["r--", "b-"])

pd.DataFrame(fit_history.history)[['R2Score', 'val_R2Score']].plot(
    grid=True, xlabel="Epoch", ax=ax[1, 0],
    style=["r--", "b-"])

pd.DataFrame(fit_history.history)[['learning_rate']].plot(
    grid=True, xlabel="Epoch", ax=ax[1, 1],
    style=["g-"])

ax[0, 0].set_ylabel('Loss')
ax[0, 0].set_title('Loss Over Epochs')
ax[0, 0].legend(['Training Loss', 'Validation Loss'])

ax[0, 1].set_ylabel('RMSE')
ax[0, 1].set_title('RMSE Over Epochs')
ax[0, 1].legend(['Training RMSE', 'Validation RMSE'])

ax[1, 0].set_ylabel('R2')
ax[1, 0].set_title('R2 Over Epochs')
ax[1, 0].legend(['Training R2', 'Validation R2'])

ax[1, 1].set_ylabel('Learning Rate')
ax[1, 1].set_title('Learning Rate Over Epochs')

plt.tight_layout()
plt.show()

5. Model Evaluation

Finally, we evaluate the model on the training and test set.

# Evaluate R-squared of predicted vs actual binding affinity for each MHC type
from sklearn.metrics import r2_score

y_pred = dnn_model.predict(X_train)
y_true = y_train

mhc_id = train_me['MHC_ID'].values
r2_df = pd.DataFrame({'MHC_ID': mhc_id, 'Y': y_true, 'Y_pred': y_pred.flatten()})

# remove MHC_ID <= 10 instances
r2_df = r2_df.groupby('MHC_ID').filter(lambda x: len(x) > 10)

# Calculate R2 and MHC instances
mhc_n = r2_df['MHC_ID'].value_counts()
r2_by_mhc = r2_df.groupby('MHC_ID').apply(lambda x: r2_score(x['Y'], x['Y_pred']))

# Join statistics by MHC_ID
r2_df = pd.concat([mhc_n, r2_by_mhc], axis=1)
r2_df.columns = ['n', 'r2']
r2_df.sort_values(by='r2', ascending=False, inplace=True)

# Plot MHC instance vs prediction accuracy scatter plot
plt.figure(figsize=(8, 6))
plt.scatter(r2_df['n'], r2_df['r2'], alpha=0.5)

from adjustText import adjust_text
texts = [plt.text(r2_df['n'].iloc[i], r2_df['r2'].iloc[i], r2_df.index[i]) for i in range(len(r2_df))]
adjust_text(
  texts, expand=(1.1, 1.1), arrowprops=dict(arrowstyle="->", color="grey")
)

plt.xlabel('Number of Instances')
plt.ylabel('R-squared')
plt.title('Number of MHC Instances vs. R-squared')
plt.grid(True, linestyle='--', alpha=0.7)
plt.show()

# Evaluate performance on test data
dnn_model.evaluate(X_test, y_test)

# Plot scatter plot of predicted vs actual on test data
import matplotlib.pyplot as plt

y_pred = dnn_model.predict(X_test)
y_true = y_test

plt.figure(figsize=(8, 6))
plt.scatter(y_true, y_pred, alpha=0.5)
plt.plot([min(y_true), max(y_true)], [min(y_pred), max(y_pred)], 'k--', lw=2)
plt.xlabel('True Values')
plt.ylabel('Predictions')
plt.title('True Values vs. Predictions')

# Add R-squared on the plot
from sklearn.metrics import r2_score
r2 = r2_score(y_true, y_pred)
plt.text(0.05, 0.95, f'R-squared: {r2:.2f}', transform=plt.gca().transAxes, fontsize=12, verticalalignment='top')

plt.show()

# Evaluate R-squared of predicted vs actual binding affinity for each MHC type (test data)
y_pred = dnn_model.predict(X_test)
y_true = y_test

# ...code abbreviated for brevity...

# Plot HLA-DPA10201-DPB10101 and peptide binding predicted vs actual scatter plot
import matplotlib.pyplot as plt

y_pred = dnn_model.predict(X_test)
y_true = y_test

mhc_id = test_me['MHC_ID'].values
target_mhc = 'HLA-DPA10201-DPB10101'
target_index = np.where(mhc_id == target_mhc)[0]

y_pred = y_pred[target_index]
y_true = y_true[target_index]

plt.figure(figsize=(8, 6))
plt.scatter(y_true, y_pred, alpha=0.5)
plt.plot([min(y_true), max(y_true)], [min(y_pred), max(y_pred)], 'k--', lw=2)
plt.xlabel('True Values')
plt.ylabel('Predictions')
plt.title('True Values vs. Predictions (HLA-DPA10201-DPB10101)')

# Add R-squared on the plot
from sklearn.metrics import r2_score
r2 = r2_score(y_true, y_pred)
plt.text(0.05, 0.95, f'R-squared: {r2:.2f}', transform=plt.gca().transAxes, fontsize=12, verticalalignment='top')

plt.show()

Conclusion

By leveraging ESM2 embeddings and MLP, we built a pipeline to predict peptide-MHC II binding. This exercise showcases the potential of combining LLMs 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.

References

• ESM2
• TensorFlow
• Keras
• Hands-on Machine Learning with Scikit-Learn, Keras, and TensorFlow
• Deep Learning for Biology

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

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")

# 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

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")

# 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 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")

# 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()

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

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)
    )

# 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

# 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)

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)
    )

# 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"
    )

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.")
}

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.

Data Loading

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

1. Sample Info: Metadata about the cell lines (lineage, disease subtype, etc.).
2. Chronos Data: CRISPR knockout scores representing gene dependency (lower score = higher dependency).
3. 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.

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

Sample Data Info:
<class 'pandas.core.frame.DataFrame'>
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

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

Chronos Data Info:
<class 'pandas.core.frame.DataFrame'>
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

# 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:
<class 'pandas.core.frame.DataFrame'>
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:

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()

# 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()

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.

# 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)

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

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}

# 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()

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

References

• Scikit-Learn Documentation
• Pandas Documentation
• Matplotlib Documentation
• DepMap Portal
• Hands-on Machine Learning with Scikit-Learn, Keras, and TensorFlow (3rd Edition)

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:

1. Random Forest (RF): A tree-based ensemble method that can capture non-linear relationships.
   
2. 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.

# 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.

# 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)

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

# 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)

# 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")

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.

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

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:

1. 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.

2. 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.

3. 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.

4. 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/

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.

# 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()

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

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

# 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()

Aim 2: Multi-omics Autoencoder

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

# 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

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.

# 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)"
    )

# 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")