onconpc/onconpc-visualization

1import os2import random3from typing import List, Mapping, Optional, Tuple, Union4 5import numpy as np6import pandas as pd7from sklearn.metrics import classification_report8from xgboost import XGBClassifier9 10"""11Author: Intae Moon12The following functions are used for training and evaluating XGBoost-based OncoNPC model.13"""14 15def get_new_indices(curr_indices: List[Union[str, float]],16					mapping_dict: Mapping[float, Union[float,str]],17					non_profile_prefix: Optional[str]='GENIE') -> List[str]:18	"""19	Get new indices based on the mapping dictionary.20	Args:21		curr_indices: list of indices to be mapped22		mapping_dict: dictionary of mapping23		non_profile_prefix: prefix for non-profile data24	Returns:25		new_indices: list of new indices26	"""27	new_indices = []28	for idx in curr_indices:29		if non_profile_prefix in str(idx):30			new_indices.append(idx)31		else:32			new_indices.append(str(mapping_dict[float(idx)]))33	return new_indices34 35def standardize_feat_names(curr_feat_names: List[str]) -> List[str]:36	"""37	Standardize feature names.38	Args:39		curr_feat_names: list of feature names40	Returns:41		new_feat_names: list of standardized feature names42	"""43	new_feat_names = []44	for feat in curr_feat_names:45		if '_mut' in feat:46			new_feat_names.append(feat.replace('_mut', '_MUT'))47		elif 'AGE' in feat or 'Age' in feat:48			new_feat_names.append('Age')49		elif 'GENDER' in feat or 'Sex' in feat:50			new_feat_names.append('Sex')51		elif 'SBS' in feat:52			new_feat_names.append(feat)53		else:   54			new_feat_names.append(feat + '_CNA')55	return new_feat_names56 57def filter_by_threshold(df: pd.DataFrame,58						threshold_per_sample: int,59						threshold_per_feature: int) -> Tuple[List[str], set]:60	"""61	Filter out features and samples based on threshold.62	"""63	binary_df = df != 064	feature_sum = binary_df.sum()65	sample_sum = binary_df.sum(axis=1)66	features_to_exclude = feature_sum.index[feature_sum < threshold_per_feature]67	samples_to_exclude = sample_sum.index[sample_sum < threshold_per_sample]68	return features_to_exclude, set(samples_to_exclude)69 70def categorize_samples_by_center(samples_to_exclude, labels_ckp):71	"""72	Categorize samples by cancer ceneter.73	"""74	profile_excluded = []75	msk_excluded = []76	vicc_excluded = []77	for sample_id in samples_to_exclude:78		sample_str = str(sample_id)79		if 'MSK' in sample_str:80			msk_excluded.append(sample_id)81		elif 'VICC' in sample_str:82			vicc_excluded.append(sample_id)83		else:84			profile_excluded.append(sample_id)85	return labels_ckp.loc[profile_excluded], labels_ckp.loc[vicc_excluded], labels_ckp.loc[msk_excluded]86 87def filter_out_low_freq_feats_and_samples(data_ckp: pd.DataFrame,88										  labels_ckp: pd.DataFrame,89										  data_cup: pd.DataFrame,90										  feature_group_to_features_dict: Mapping[str, List[str]],91										  threshold_per_sample: int=3,92										  threshold_per_feature: int=50) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:93	"""94	Filter out low frequency features and samples.95 96	Args:97		data_ckp: CKP data98		labels_ckp: CKP labels99		data_cup: CUP-CKP data100		feature_group_to_features_dict: dictionary of feature groups to features101		threshold_per_sample: threshold per sample102		threshold_per_feature: threshold per feature103	Returns:104		data_ckp_filtered: filtered CKP data105		labels_ckp_filtered: filtered CKP labels106		data_cup_fitlered: filtered CUP-CKP data107	"""108	mutation_features_to_exclude, mutation_samples_to_exclude = filter_by_threshold(109		data_ckp[feature_group_to_features_dict['mutation']], threshold_per_sample, threshold_per_feature110	)111	112	cna_features_to_exclude, cna_samples_to_exclude = filter_by_threshold(113		data_ckp[feature_group_to_features_dict['cna']], threshold_per_sample, threshold_per_feature114	)115	116	samples_to_exclude = mutation_samples_to_exclude & cna_samples_to_exclude117	118	profile_labels, vicc_labels, msk_labels = categorize_samples_by_center(samples_to_exclude, labels_ckp)119	120	total_excluded_samples = len(profile_labels) + len(vicc_labels) + len(msk_labels)121	print('\n')122	print('Filtering out low frequency features and samples...\n')123	print(f'Total excluded number of patients : {total_excluded_samples}\n')124	print('Profile:\n', pd.value_counts(profile_labels.cancer_type.values, sort=True))125	print('VICC:\n', pd.value_counts(vicc_labels.cancer_type.values, sort=True))126	print('MSK:\n', pd.value_counts(msk_labels.cancer_type.values, sort=True))127	128	features_to_drop = set(mutation_features_to_exclude) | set(cna_features_to_exclude)129	130	print('Dropping features...\n', features_to_drop)131	data_ckp_filtered = data_ckp.drop(columns=features_to_drop)132	data_cup_fitlered = data_cup.drop(columns=features_to_drop)133	134	print('Dropping samples...\n')135	data_ckp_filtered = data_ckp.drop(index=samples_to_exclude)136	labels_ckp_filtered = labels_ckp.drop(index=samples_to_exclude)137	138	print('Samples labels and features matching :')139	print('CKP labels:', all(labels_ckp.index == data_ckp.index))140	print('CUP-CKP features:', all(data_ckp.columns == data_cup.columns))141	return data_ckp_filtered, labels_ckp_filtered, data_cup_fitlered142 143def get_cancer_to_num_val_samples(labels: pd.DataFrame,144								  cancer_types: List[str],145								  k_fold: int=10) -> Mapping[str, int]:146	"""147	Get number of validation samples per cancer type.148	Args:149		labels: labels dataframe with cancer types150		cancer_types: list of cancer types to consider151		k_fold: number of folds152	Returns:153		cancer_to_num_val_samples_dict: dictionary of cancer type to number of validation samples154	"""155	test_frac = 1/k_fold156	cancer_to_num_val_samples_dict = {}157	for cancer in cancer_types:158		labels_cancer = labels.loc[labels['cancer_type'] == cancer]159		cancer_to_num_val_samples_dict[cancer] = int(np.floor(len(labels_cancer) * test_frac))160	return cancer_to_num_val_samples_dict161 162def get_sample_indices_and_labels_based_on_cut_off(pred_probs: np.ndarray,163												   p_max_cut_off: float) -> Tuple[List[int], List[int]]:164	"""165	Get sample indices and labels based on p max cut-off.166	Args:167		pred_probs: predicted probabilities168		p_max_cut_off: cut-off169	Returns:170		indices: list of indices171		max_prob_labels: list of labels172	"""173	indices, max_prob_labels = [], []174	for num_idx, max_prob_idx in enumerate(np.argmax(pred_probs, axis=1)):175		if pred_probs[num_idx][max_prob_idx] > p_max_cut_off:176			indices.append(num_idx)177			max_prob_labels.append(max_prob_idx)178	return indices, max_prob_labels179 180def fit_and_evaluate_model(X_train: np.ndarray,181						   y_train: np.ndarray,182						   X_test: Optional[np.ndarray],183						   y_test: Optional[np.ndarray],184						   params_xgb: Mapping[str, Union[str, int, float]],185						   cancer_types: List[str]) -> Tuple[XGBClassifier, np.ndarray, Mapping[str, Union[str, int, float]]]:186	"""187	Fit and evaluate XGBoost model.188	189	Args:190		X_train: training data191		y_train: training labels192		X_test: test data; optional193		y_test: test labels; optional194		params_xgb: XGBoost parameters195		cancer_types: list of cancer types196	Returns:197		xg_clf: XGBoost model198		pred_probs_on_test: predicted probabilities on test data199		performance_report: performance report200	"""201	xg_clf = XGBClassifier(202		tree_method='hist', 203		n_estimators=int(params_xgb['n_estimators']),204		max_depth=int(params_xgb['max_depth']),205		scale_pos_weight=int(params_xgb['scale_pos_weight']),206		learning_rate=float(params_xgb['learning_rate']),207		verbosity=0208	)209	xg_clf.fit(X_train, y_train, verbose=False)210	if X_test is not None and y_test is not None:211		pred_probs_on_test = xg_clf.predict_proba(X_test)212		performance_report = classification_report(y_test, xg_clf.predict(X_test), target_names=cancer_types, output_dict=True)213		return xg_clf, pred_probs_on_test, performance_report214	else:215		return xg_clf216 217def perform_k_fold(data: pd.DataFrame,218				   labels: pd.DataFrame,219				   cancer_types: List[str],220				   params_xgb: Mapping[str, Union[str, int, float]],221				   k_fold: int=10,222				   save_model_name: Optional[str]=None,223				   p_max_cut_offs: Optional[List[float]]=[0.0, 0.5, 0.7, 0.9]224				   ) -> Tuple[Mapping[int, pd.DataFrame], pd.DataFrame]:	225	"""226	Performs k-fold cross validation.227 228	Args:229		data: training and validation data230		labels: training and validation labels231		cancer_types: list of cancer types232		params_xgb: XGBoost parameters233		k_fold: number of folds234		save_model_name: name of the model to save235		p_max_cut_offs: list of max prediction probability cut-offs236	Returns:237		k_fold_to_performance_report_dict: dictionary of k-fold to performance report238		pred_probs_on_val_total_df: dataframe of predicted probabilities on validation data239	"""240	print('\n')241	print('Chosen parameters for XGBoost:\n', params_xgb)242	print('\n')243	cancer_to_num_val_samples_dict = get_cancer_to_num_val_samples(labels, cancer_types, k_fold)244	k_fold_to_performance_report_dict = {}245	total_val_sampled_so_far = []246	for k in range(k_fold):247		print('\n')248		print(f'k = {k}')249		if k == k_fold - 1:250			# At the last fold, use the remaining samples as validation data.251			val_labels_sampled = labels.loc[list(set(labels.index) - set(total_val_sampled_so_far))]252		else:253			val_labels_sampled = pd.DataFrame()254			# Update the validation labels to sample from.255			val_labels_to_sample_from = labels.loc[list(set(labels.index) - set(total_val_sampled_so_far))]256			for cancer in cancer_types:257				# Sample validation data from each cancer type.258				# this ensures that the validation data is balanced wrt cancer type.259				# Set random seed260				np.random.seed(k)261				random.seed(k)262				val_labels_sampled_curr = val_labels_to_sample_from.loc[val_labels_to_sample_from['cancer_type'] == cancer].sample(263					n=cancer_to_num_val_samples_dict[cancer], replace=False)264				val_labels_sampled = pd.concat([val_labels_sampled, val_labels_sampled_curr])265			# Update the total validation samples sampled so far.266			total_val_sampled_so_far.extend(list(val_labels_sampled.index))267 268		y_val = val_labels_sampled['cancer_label']269		X_val = data.loc[val_labels_sampled.index]270		X_train = data.loc[list(set(data.index) - set(X_val.index))]271		y_train = labels.loc[X_train.index]['cancer_label']272		# Standardize Age based on train data273		age_mean = X_train['Age'].mean()274		age_std = X_train['Age'].std()275		X_train['Age'] = (X_train['Age'] - age_mean) / age_std276		X_val['Age'] = (X_val['Age'] - age_mean) / age_std277		xg_clf, pred_probs_on_val, performance_report = fit_and_evaluate_model(X_train.values,278																		 y_train.values,279																		 X_val.values,280																		 y_val.values,281																		 params_xgb,282																		 cancer_types)283		# Evaluate the model performance based on different maximum prediction probability cut-offs.284		p_max_cut_off_to_performance_report_dict = {}285		for p_max_cut_off in p_max_cut_offs:286			indices, max_prob_labels = get_sample_indices_and_labels_based_on_cut_off(pred_probs_on_val, p_max_cut_off)287			unique_labels = list(set(y_val[indices]))288			report_cut_off = classification_report(y_val[indices], max_prob_labels,289										  target_names=cancer_types, labels=unique_labels, output_dict=True)290			print(pd.DataFrame(report_cut_off))291			p_max_cut_off_to_performance_report_dict[p_max_cut_off] = pd.DataFrame(report_cut_off)292		k_fold_to_performance_report_dict[k] = p_max_cut_off_to_performance_report_dict293		if save_model_name is not None:294			if not os.path.exists('../models'):295				os.makedirs('../models/models')296			xg_clf.save_model(f'../models/xgboost_{save_model_name}_{k}.json')297 298		# Get prediction probabilities on val data.299		pred_probs_on_val_df = pd.DataFrame(pred_probs_on_val, columns=cancer_types, index=X_val.index)300		pred_probs_on_val_df['predicted_cancer'] = [cancer_types[max_idx] for max_idx in pred_probs_on_val.argmax(axis=1)]301		pred_probs_on_val_df['prediction_prob'] = [pred_probs_on_val[num_idx][max_idx]302											 for num_idx, max_idx in enumerate(pred_probs_on_val.argmax(axis=1))]303		if k == 0:304			pred_probs_on_val_total_df = pred_probs_on_val_df305		else:306			pred_probs_on_val_total_df = pd.concat([pred_probs_on_val_total_df, pred_probs_on_val_df])307	return k_fold_to_performance_report_dict, pred_probs_on_val_total_df308