onconpc/onconpc-visualization

1from absl import app2from absl import flags3import pandas as pd4import pickle5import utils_training6import utils7import os8import numpy as np9 10"""11Author: Intae Moon12In this script, we train and evaluate the performance of the XGBoost model on the OncoTree-based cancer types.13One can choose to use k-fold cross validation or train on the entire dataset and test on the held out set.14"""15 16flags.DEFINE_string('config', None, 'Cancer centers to incorporate for feature processing (genie, profile_dfci, or both)')17flags.DEFINE_integer('k_fold', 10, 'Number of folds in k-fold cross-validation')18flags.DEFINE_boolean('use_held_out_set', False, 'Whether to use a held-out set')19flags.DEFINE_boolean('filter_out_non_informatives', False, 'Whether to filter out non-informative features and samples')20flags.DEFINE_string('save_model_name', None, 'Name of the model to save')21flags.DEFINE_string('index_profile_samples_by', 'SAMPLE_ACCESSION_NBR', 'Index profile samples by DFCI_MRN or SAMPLE_ACCESSION_NBR')22# DATA_PATH = '../data/internal_training_data/'23DATA_PATH = '../data/'24 25def main(argv=None):26	# Custom check for flags that you still require to be non-default27	FLAGS = flags.FLAGS28	# Training configs:29	config = FLAGS.config30	k_fold = FLAGS.k_fold31	use_held_out_set = FLAGS.use_held_out_set32	save_model_name = FLAGS.save_model_name33	index_profile_samples_by = FLAGS.index_profile_samples_by34	filter_out_non_informatives = FLAGS.filter_out_non_informatives35	if config not in ['genie', 'profile_dfci', 'both']:36		raise ValueError("config must be one of genie, profile_dfci, or both")37	if config == 'genie' and use_held_out_set == True:38		raise ValueError("Cannot use held out set when training on genie data only")39	# load xgb parameters40	hyperparam_xgb = pd.read_csv(os.path.join(DATA_PATH, 'xgb_hyperparams'), sep = '\t')  #xgb_params_to_run # params_w_results_Mar_1st41	hyperparam_xgb.drop(columns = 'Unnamed: 0', inplace = True)42	# Get the best performing hyperparameters from the previous experiments43	# this can be extended to try multiple hyperparameters44	params_xgb = hyperparam_xgb.iloc[0:1]45 46	# ====================================================================================================47	# Loading data and labels, and processing them for training48	# ====================================================================================================49	# Load trainable data and labels50	# Tab separated feature data for CKPs51	feature_data_name = os.path.join(DATA_PATH, 'features_genie_') # Replace with your feature file path52	# Tab separated label data for CKPs53	label_data_name = os.path.join(DATA_PATH, 'labels_genie_')  # Replace with your label file path54	if config in ['profile_dfci', 'both']:55		# Tab separated feature data for CUPs56		feature_data_name_cup = os.path.join(DATA_PATH, 'features_combined_cup_onco_tree_based_dev')57		features_cup_df = pd.read_csv(feature_data_name_cup, sep = '\t')58		features_cup_df.set_index('Unnamed: 0', inplace = True)59	features_ckp_df = pd.read_csv(feature_data_name, sep = '\t')60	features_ckp_df.set_index(features_ckp_df.columns[0], inplace = True)61	labels_ckp_df = pd.read_csv(label_data_name, sep = '\t')62	labels_ckp_df.set_index(labels_ckp_df.columns[0], inplace = True)63 64	if config in ['profile_dfci', 'both']:65		if index_profile_samples_by == 'DFCI_MRN':66			# Load ../data/internal_training_data/onconpc_sample_id_to_dfci_mrn.pkl67			with open(os.path.join(DATA_PATH, 'onconpc_sample_id_to_dfci_mrn.pkl'), "rb") as fp:   # Unpickling68				onconpc_sample_id_to_dfci_mrn = pickle.load(fp)69			# index CKP data by DFCI_MRN70			features_ckp_df.index = utils_training.get_new_indices(features_ckp_df.index, onconpc_sample_id_to_dfci_mrn)71			labels_ckp_df.index = utils_training.get_new_indices(labels_ckp_df.index, onconpc_sample_id_to_dfci_mrn)72			features_cup_df.index = utils_training.get_new_indices(features_cup_df.index, onconpc_sample_id_to_dfci_mrn)73			# Set index name74			features_ckp_df.index.name = 'DFCI_MRN_FOR_PROFILE'75			labels_ckp_df.index.name = 'DFCI_MRN_FOR_PROFILE'76			features_cup_df.index.name = 'DFCI_MRN_FOR_PROFILE'77		elif index_profile_samples_by == 'SAMPLE_ACCESSION_NBR':78			# Load ../data/internal_training_data/onconpc_sample_id_to_sample_accession_nbr.pkl79			with open(os.path.join(DATA_PATH, 'onconpc_sample_id_to_sample_accession_nbr.pkl'), "rb") as fp:   # Unpickling80				onconpc_sample_id_to_sample_accession_nbr = pickle.load(fp)81			# index CKP data by SAMPLE_ACCESSION_NBR82			features_ckp_df.index = utils_training.get_new_indices(features_ckp_df.index, onconpc_sample_id_to_sample_accession_nbr)83			labels_ckp_df.index = utils_training.get_new_indices(labels_ckp_df.index, onconpc_sample_id_to_sample_accession_nbr)84			features_cup_df.index = utils_training.get_new_indices(features_cup_df.index, onconpc_sample_id_to_sample_accession_nbr)85			# Set index name86			features_ckp_df.index.name = 'SAMPLE_ACCESSION_NBR_FOR_PROFILE'87			labels_ckp_df.index.name = 'SAMPLE_ACCESSION_NBR_FOR_PROFILE'88			features_cup_df.index.name = 'SAMPLE_ACCESSION_NBR_FOR_PROFILE'89 90		if use_held_out_set:91			heldout_ckps_preds_df = pd.read_csv(os.path.join(DATA_PATH, 'heldout_ckps_preds_onco_tree_based_dev'), sep = '\t')92			heldout_ckps_preds_df.set_index(heldout_ckps_preds_df.columns[0], inplace = True)93			held_out_indices = utils_training.get_new_indices([idx[:-3] for idx in heldout_ckps_preds_df.index],94														onconpc_sample_id_to_sample_accession_nbr)95			# Exclude held out samples from training96			indices_to_choose = set(features_ckp_df.index) - set(held_out_indices)97			features_ckp_df = features_ckp_df.loc[indices_to_choose]98			labels_ckp_df = labels_ckp_df.loc[indices_to_choose]99	feature_group_to_features_dict = utils.partition_feature_names_by_group(list(features_ckp_df.columns))100	if filter_out_non_informatives:101		if config in ['profile_dfci', 'both']:102			(features_ckp_final_df,103		labels_ckp_final_df,104		features_cup_final_df) = utils_training.filter_out_low_freq_feats_and_samples(features_ckp_df,105																						labels_ckp_df,106																						features_cup_df, 107																						feature_group_to_features_dict)108		else:109			(features_ckp_final_df,110		labels_ckp_final_df,111		_) = utils_training.filter_out_low_freq_feats_and_samples(features_ckp_df,112															labels_ckp_df,113															features_ckp_df, # placeholder for cup_df114															feature_group_to_features_dict)115	else:116		features_ckp_final_df = features_ckp_df117		labels_ckp_final_df = labels_ckp_df118		if config in ['profile_dfci', 'both']:119			features_cup_final_df = features_cup_df120 121	# ====================================================================================================122	# Model Training and Evaluation123	# ====================================================================================================124	# Cancer types to consider125	cancer_types = list(np.unique(labels_ckp_final_df.cancer_type.values))126	# Check duplicate indices127	if len(features_ckp_final_df) == len(set(features_ckp_final_df.index)):128		print('\n')129		print('No duplicate indices')130	else:131		print('\n')132		print('Duplicate indices detected (most likely due to DFCI_MRN indexing)')133	if k_fold > 0:134		(k_fold_to_performance_report_dict,135	pred_probs_on_val_total_df) = utils_training.perform_k_fold(features_ckp_final_df,136																labels_ckp_final_df,137																cancer_types,138																params_xgb,139																k_fold=k_fold,140																save_model_name=save_model_name)141		# Store k_fold_to_performance_report_dict and pred_probs_on_val_total_df142		with open(os.path.join(DATA_PATH, f'k_fold_to_performance_report_dict_{save_model_name}.pkl'), 'wb') as f:143			pickle.dump(k_fold_to_performance_report_dict, f)144		pred_probs_on_val_total_df.to_csv(os.path.join(DATA_PATH, f'pred_probs_on_val_total_df_{save_model_name}.csv'))145	else:146		X_train = features_ckp_final_df.copy()147		y_train = labels_ckp_final_df.loc[X_train.index]['cancer_label']148		# Standardize Age based on train data149		age_mean = X_train['Age'].mean()150		age_std = X_train['Age'].std()151		X_train['Age'] = (X_train['Age'] - age_mean) / age_std152		xg_clf = utils_training.fit_and_evaluate_model(X_train.values,153												 y_train.values,154												 None,155												 None,156												 params_xgb,157												 cancer_types)158		# Evaluate the model performance based on different maximum prediction probability cut-offs.159		if save_model_name is not None:160			if not os.path.exists('../models'):161				os.makedirs('../models/models')162			xg_clf.save_model(f'../models/xgboost_{save_model_name}_trained_on_all_ckps.json')163	return 164 165if __name__ == '__main__':166	flags.mark_flags_as_required(['k_fold', 'config', 'use_held_out_set', 'save_model_name'])167	app.run(main)168