作者，Evil Genius
而我们需要实现的分析，即VDJ聚类与motif分析
分析会在课上讲到，报名链接在2024年单细胞空间系列课程
完整脚本如下（封装版）
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from keras.callbacks import EarlyStopping
from keras.callbacks import TensorBoard
import tensorflow as tf
import tensorflow_addons as tfa
#import tensorflow_datasets as tfds
import csv
import os
import argparse
#import other python files
import config_script
import model_architecture
import enc_decode
import plot_latent_space
import cluster_performance_functions
# parse the command line
if __name__ == "__main__":
# initialize Argument Parser
parser = argparse.ArgumentParser() #help="Debug", nargs='?', type=int, const=1, default=7
# TCR input parameters
parser.add_argument("--input_type", help="beta_chain or paired_chain", type=str)
parser.add_argument("--encoding_type", help="onehot or ESM", type=str) ## one-hot or ESM embeddings
parser.add_argument("--esm_type", help="ESM1v or ESM2", type=str, default='ESM1v')
parser.add_argument("--chain_type", help="TCR or ab_VC or ab_V", type=str, default='ab_V')
# model parameters
parser.add_argument("--embedding_space", type=int) ## embedding dimension for TCR
parser.add_argument("--epochs", type=int, default=100)
parser.add_argument("--patience", type = int, default = 15)
parser.add_argument("--batch_size", type=int, default=64)
parser.add_argument("--learning_rate", type=float, default=0.0001)
parser.add_argument("--triplet_loss", help="hard or semihard", type=str, default='semihard')
#plot embeddings
parser.add_argument("--plot_embedding", type = bool, default = False)
######################################## --> started here
# we then read the arguments from the comand line
args = parser.parse_args()
INPUT_TYPE = args.input_type #beta_chain or paired_chain
ENCODING_TYPE = args.encoding_type
ESM_TYPE = args.esm_type
CHAIN_TYPE = args.chain_type
EMBEDDING_SPACE = args.embedding_space
EPOCHS = args.epochs
PATIENCE = args.patience
BATCH_SIZE = args.batch_size
LRATE = args.learning_rate
LOSS = args.triplet_loss # 'semihard' or 'hard'
PLOT_EMB = args.plot_embedding
# we then read the arguments from the comand line
# DATA_SIZE = args.data_size
#triplet hyperparameters
LOSS = args.triplet_loss #'semihard' or 'hard'
PLOT_EMB = args.plot_embedding
########################################################################################
############################ Open TCRpMHC and encode them ##############################
# open the files and assign some constant variables
TCRpMHC = pd.read_csv(config_script.TCRpMHC_ESM2_data, header = 0)
epi_list = TCRpMHC['Epitope'].unique().tolist()
epi_color = plot_latent_space.palette_31[:len(epi_list)]
epi_color_dict = dict(zip(epi_list, epi_color))
TCRpMHC['color'] = TCRpMHC['Epitope'].map(epi_color_dict)
TCRpMHC['CDR123ab'] = TCRpMHC['CDR1a'] + TCRpMHC['CDR2a'] + TCRpMHC['CDR3a'] + \
TCRpMHC['CDR1b'] + TCRpMHC['CDR2b'] + TCRpMHC['CDR3b']
TCRpMHC['CDR123b'] = TCRpMHC['CDR1b'] + TCRpMHC['CDR2b'] + TCRpMHC['CDR3b']
if ENCODING_TYPE == 'onehot':
model_input_dict = config_script.model_input_dict
model_input_key = INPUT_TYPE
df_columns = model_input_dict[model_input_key][0] #get columns to encode
max_seq_len = model_input_dict[model_input_key][1] #get max length for padding
print ('df_columns: ', df_columns, 'max_seq_len: ', max_seq_len)
# make a dist of one-hot encoded amino acids for data encoding:
AA_list = list("ACDEFGHIKLMNPQRSTVWY_")
AA_one_hot_dict = {char : l for char, l in zip(AA_list, np.eye(len(AA_list), k=0))}
if ENCODING_TYPE == 'ESM':
if ESM_TYPE == 'ESM1v':
if CHAIN_TYPE == 'TCR':
embeddings = np.loadtxt(config_script.esm_embed_dir + 'albeta_VC_esm1v.txt', dtype = float)
if CHAIN_TYPE == 'ab_VC':
if INPUT_TYPE == 'beta_chain':
embeddings = np.loadtxt(config_script.esm_embed_dir + 'beta_VC_esm1v.txt', dtype = float)
if INPUT_TYPE == 'paired_chain':
embeddings_b = np.loadtxt(config_script.esm_embed_dir + 'beta_VC_esm1v.txt', dtype = float)
embeddings_a = np.loadtxt(config_script.esm_embed_dir + 'alpha_VC_esm1v.txt', dtype = float)
embeddings = np.concatenate((embeddings_a, embeddings_b), axis=1)
if CHAIN_TYPE == 'ab_V':
if INPUT_TYPE == 'beta_chain':
embeddings = np.loadtxt(config_script.esm_embed_dir + 'beta_V_esm1v.txt', dtype = float)
if INPUT_TYPE == 'paired_chain':
embeddings_a = np.loadtxt(config_script.esm_embed_dir + 'alpha_V_esm1v.txt', dtype = float)
embeddings_b = np.loadtxt(config_script.esm_embed_dir + 'beta_V_esm1v.txt', dtype = float)
embeddings = np.concatenate((embeddings_a, embeddings_b), axis=1)
if ESM_TYPE == 'ESM2':
if CHAIN_TYPE == 'TCR':
embeddings = np.loadtxt(config_script.esm_embed_dir + 'albeta_VC_esm2.txt', dtype = float)
if CHAIN_TYPE == 'ab_VC':
if INPUT_TYPE == 'beta_chain':
embeddings = np.loadtxt(config_script.esm_embed_dir + 'beta_VC_esm2.txt', dtype = float)
if INPUT_TYPE == 'paired_chain':
embeddings_a = np.loadtxt(config_script.esm_embed_dir + 'alpha_VC_esm2.txt', dtype = float)
embeddings_b = np.loadtxt(config_script.esm_embed_dir + 'beta_VC_esm2.txt', dtype = float)
embeddings = np.concatenate((embeddings_a, embeddings_b), axis=1)
if CHAIN_TYPE == 'ab_V':
if INPUT_TYPE == 'beta_chain':
embeddings = np.loadtxt(config_script.esm_embed_dir + 'beta_V_esm2.txt', dtype = float)
if INPUT_TYPE == 'paired_chain':
embeddings_a = np.loadtxt(config_script.esm_embed_dir + 'alpha_V_esm2.txt', dtype = float)
embeddings_b = np.loadtxt(config_script.esm_embed_dir + 'beta_V_esm2.txt', dtype = float)
embeddings = np.concatenate((embeddings_a, embeddings_b), axis=1)
# K-fold Cross Validation model evaluation #when model parameters are already settled
#folds_list = TCRpMHC['part'].unique().tolist()
tf.keras.utils.set_random_seed(1234) ## set ALL random seeds for the program
import random
all_results = pd.DataFrame()
test_val_folds = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 0]]
for test_val in test_val_folds: #len(test_val) range(6)
#test_val = random.sample(folds_list, 2)
test_val_fold_str = str(test_val[0]) + str(test_val[1])
#test_val_fold_str = str(fold[0]) + str(fold[1])
val = TCRpMHC[TCRpMHC['part']==test_val[0]]
test = TCRpMHC[TCRpMHC['part']==test_val[1]]
train = TCRpMHC[~TCRpMHC['part'].isin(test_val)]
#test = TCRpMHC[TCRpMHC['part']==5] #6
#encode the data
if ENCODING_TYPE == 'onehot':
X_train = enc_decode.CDR_one_hot_encoding(train, df_columns, AA_list, AA_one_hot_dict, max_seq_len)
X_val = enc_decode.CDR_one_hot_encoding(val, df_columns, AA_list, AA_one_hot_dict, max_seq_len)
X_test = enc_decode.CDR_one_hot_encoding(test, df_columns, AA_list, AA_one_hot_dict, max_seq_len)
#X_all_data = enc_decode.CDR_one_hot_encoding(TCRpMHC, df_columns, AA_list, AA_one_hot_dict, max_seq_len)
print ('Seqs encoded!')
if ENCODING_TYPE == 'ESM':
#1) get indexes
train_idx = train.index.tolist()
val_idx = val.index.tolist()
test_idx = test.index.tolist()
#2) get embeddings
X_train = embeddings[train_idx]
X_val = embeddings[val_idx]
X_test = embeddings[test_idx]
#3) change the dimension to train CNN
dim_2 = int(len(X_train[0])/20) #get the dimension for reshaping
X_train = np.reshape(X_train, (-1, dim_2, 20)) #it's either 64*20 or 128*20, depending on TCR input type
X_val = np.reshape(X_val, (-1, dim_2, 20))
X_test = np.reshape(X_test, (-1, dim_2, 20))
y_train = train['label'].values
y_val = val['label'].values
y_test = test['label'].values
#save the Epi label for plotting
y_train_epi = train['Epitope'].tolist()
y_val_epi = val['Epitope'].tolist()
y_test_epi = test['Epitope'].tolist()
print ('X_train: ', X_train.shape, \
'X_val: ', X_val.shape, \
'X_test: ', X_test.shape)
#turn array into tensors
train_triplet = tf.data.Dataset.from_tensor_slices((X_train, y_train))
train_triplet = train_triplet.shuffle(1024).batch(BATCH_SIZE)
val_triplet = tf.data.Dataset.from_tensor_slices((X_val, y_val))
val_triplet = val_triplet.batch(BATCH_SIZE)
######################################################################################
############################ Set up & train a model #################################
if LOSS == 'hard':
loss = tfa.losses.TripletHardLoss()
if LOSS == 'semihard':
loss = tfa.losses.TripletSemiHardLoss()
tf.keras.backend.clear_session()
triplet = model_architecture.Triplet_CNN_naive(X_train.shape[1:], embedding_units = EMBEDDING_SPACE)
triplet.compile(optimizer=tf.keras.optimizers.Adam(LRATE), loss=loss) #tfa.losses.TripletSemiHardLoss()
triplet.summary()
early_stop = EarlyStopping(monitor='val_loss', mode='min', patience = PATIENCE, \
verbose = 1, restore_best_weights = True)
print ('Starting training...')
history = triplet.fit(
train_triplet,
epochs=EPOCHS,
validation_data = val_triplet,
callbacks = [early_stop])
print ('Model trained!')
# Extract the best loss value (from early stopping)
loss_hist = history.history['val_loss']
best_val_loss = np.min(loss_hist)
best_epochs = np.argmin(loss_hist) + 1
if ENCODING_TYPE == 'onehot':
TRIPLET_NAME =f'triplet_{ENCODING_TYPE}_{INPUT_TYPE}_{EMBEDDING_SPACE}d_{LRATE}lr_{BATCH_SIZE}btch_{EPOCHS}ep_{PATIENCE}ptence_fold{test_val_fold_str}'
if ENCODING_TYPE == 'ESM':
TRIPLET_NAME =f'triplet_{INPUT_TYPE}_{ESM_TYPE}_{CHAIN_TYPE}_{EMBEDDING_SPACE}d_{LRATE}lr_{BATCH_SIZE}btch_{EPOCHS}ep_{PATIENCE}ptence_fold{test_val_fold_str}'
modelpath = config_script.triplet_out_model + TRIPLET_NAME
if not os.path.exists(modelpath):
os.makedirs(modelpath)
# save the model
triplet.save(config_script.triplet_out_model + TRIPLET_NAME)
#print ('\n')
print ('Model name: ', TRIPLET_NAME)
print ('best loss: ', best_val_loss, 'best_epochs: ', best_epochs)
print ('Triplet model saved!')
############################## extract clustering performance ###############################
############################## ################################ #############################
if INPUT_TYPE == 'beta_chain':
seq_column = 'CDR123b'
if INPUT_TYPE == 'paired_chain':
seq_column = 'CDR123ab'
results_val = triplet(X_val)
val_emb_numpy = results_val.numpy()
results_test = triplet(X_test)
test_emb_numpy = results_test.numpy()
epsilon_threshold = [0.2, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.7, 0.85, 0.95]
for epsilon in epsilon_threshold:
result_dbscan_val = cluster_performance_functions.get_performance_DBSCAN(val, TRIPLET_NAME, \
val_emb_numpy, seq_column, epsilon)
result_dbscan_val['split'] = 'val'
result_dbscan_val['fold'] = test_val[0]
result_dbscan_test = cluster_performance_functions.get_performance_DBSCAN(test, TRIPLET_NAME, \
test_emb_numpy, seq_column, epsilon)
result_dbscan_test['split'] = 'test'
result_dbscan_test['fold'] = test_val[1]
all_results = pd.concat([all_results, result_dbscan_val, result_dbscan_test], ignore_index = True)
all_results.to_csv(config_script.triplet_analysis_dir + TRIPLET_NAME + '_WITH_FOLDS_clstr_metrics.csv')
"""
########################################################################################
############################## Plot model embeddings ###################################
triplet = keras.models.load_model(config_script.triplet_out_model + TRIPLET_NAME)
if PLOT_EMB == True:
results_val = triplet(X_val)
results_test = triplet(X_test)
results_train = triplet(X_train)
umap_pca_test = plot_latent_space.umap_pca_df(results_test, y_test_epi, random_seed = 42)
plot_latent_space.plot_umap_pca_latent(umap_pca_test, 'Epitope', color_dict = epi_color_dict, \
fig_title = f'Epitope distribution by {TRIPLET_MODE} model with {INPUT_TYPE} triplet {EMBEDDING_SPACE}embedding ', \
legend_title = 'Epitope', node_size = 70, \
save = True, figname = f'Epitope_{TRIPLET_NAME}.jpg')
生活很好，有你更好
TCR clustering by contrastive learning on antigen specificity
扫一扫
6509
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
