onconpc/onconpc-visualization

1import glob2from typing import List, Mapping, Optional, Any, Tuple3import collections4import os5import pickle6from matplotlib.patches import Patch7from rpy2.robjects import conversion, default_converter8 9import matplotlib.pyplot as plt10import numpy as np11import pandas as pd12import shap13import rpy2.robjects as robjects14from rpy2.robjects import pandas2ri15from rpy2.robjects.conversion import localconverter16import xgboost as xgb17from adjustText import adjust_text18from matplotlib.patches import Patch19 20"""21Author: Intae Moon22This script contains utility functions for predictions and interpretations.23"""24 25def get_xgboost_cancer_type_preds(xgb_model: xgb.sklearn.XGBClassifier,26								  features_test_df: pd.DataFrame,27								  cancer_types: List[str]) -> pd.DataFrame:28	"""Returns cancer type predictions for test set using XGBoost model.29	30	Args:31		xgb_model: XGBoost model.32		features_test_df: Test set features.33		cancer_types: List of cancer types.34	Returns:35		pd.DataFrame containing cancer type predictions and prediction probabilities.36	"""37	ckp_test_pred_probs = xgb_model.predict_proba(features_test_df.values)38	ckp_test_preds = ckp_test_pred_probs.argmax(axis=1)39	max_posteriors = [pred_dist[max_idx] for pred_dist, max_idx in40					  zip(ckp_test_pred_probs, ckp_test_preds)]41	ckp_test_preds_df = pd.DataFrame(ckp_test_pred_probs,42									 index=features_test_df.index,43									 columns=cancer_types)44	ckp_test_preds_df['max_posterior'] = max_posteriors45	ckp_test_preds_df['cancer_type'] = [cancer_types[max_idx] for max_idx in ckp_test_preds]46	return ckp_test_preds_df47 48def get_xgboost_latest_cancer_type_preds(xgb_model: xgb.core.Booster,49										 features_test_df: pd.DataFrame,50										 cancer_types: List[str]) -> pd.DataFrame:51	"""Returns cancer type predictions for test set using XGBoost model.52	53	Args:54		xgb_model: XGBoost model.55		features_test_df: Test set features.56		cancer_types: List of cancer types.57	Returns:58		pd.DataFrame containing cancer type predictions and prediction probabilities.59	"""60	dtest = xgb.DMatrix(features_test_df.values)61	ckp_test_pred_probs = xgb_model.predict(dtest, output_margin=True)62 63	ckp_test_pred_probs = np.exp(ckp_test_pred_probs)64	ckp_test_pred_probs /= ckp_test_pred_probs.sum(axis=1, keepdims=True)65	ckp_test_preds = ckp_test_pred_probs.argmax(axis=1)66	max_posteriors = [pred_dist[max_idx] for pred_dist, max_idx in67					  zip(ckp_test_pred_probs, ckp_test_preds)]68	ckp_test_preds_df = pd.DataFrame(ckp_test_pred_probs,69									 index=features_test_df.index,70									 columns=cancer_types)71	ckp_test_preds_df['max_posterior'] = max_posteriors72	ckp_test_preds_df['cancer_type'] = [cancer_types[max_idx] for max_idx in ckp_test_preds]73	return ckp_test_preds_df74 75def obtain_shap_values(model: xgb.sklearn.XGBClassifier,76					   data: pd.DataFrame) -> np.ndarray:77	"""Returns SHAP values for predictions based on data.78	79	Args:80		model: XGBoost model.81		data: Data to obtain SHAP values for.82	Returns:83		Numpy array containing SHAP values.84	"""85	# Get SHAP values using the model in byte array.86	mybooster = model.get_booster()87	model_bytearray = mybooster.save_raw()[4:]88	def in_bytearray(self=None):89		return model_bytearray90	mybooster.save_raw = in_bytearray91	shap_ex = shap.TreeExplainer(mybooster)92	return shap_ex.shap_values(data)93 94def obtain_shap_values_with_latest_xgboost(model: xgb.core.Booster,95										   data: pd.DataFrame) -> np.ndarray:96	"""Returns SHAP values for predictions based on data.97	98	Args:99		model: XGBoost model.100		data: Data to obtain SHAP values for.101	Returns:102		Numpy array containing SHAP values.103	"""104	# Directly use the Booster with the TreeExplainer105	shap_ex = shap.TreeExplainer(model)106	return shap_ex.shap_values(data)107 108def partition_feature_names_by_group(fature_names: List[str]):109	"""Partitions feature names into groups.110	111	Args:112		feature_names: List of feature names.113	Returns:114		Dictionary mapping feature groups to feature names and vice versa.115	"""116	feature_group_to_features_dict = collections.defaultdict(list)117	feature_to_feature_group_dict = {}118	for feat in fature_names:119		if 'SBS' in feat:120			feature_group_to_features_dict['signature'].append(feat)121			feature_to_feature_group_dict[feat] = 'signature'122		elif feat in ['Age', 'Sex']:123			feature_group_to_features_dict['clinical'].append(feat)124			feature_to_feature_group_dict[feat] = 'clinical'125		elif 'CNA' in feat:126			feature_group_to_features_dict['cna'].append(feat)127			feature_to_feature_group_dict[feat] = 'cna'128		else:129			feature_group_to_features_dict['mutation'].append(feat)130			feature_to_feature_group_dict[feat] = 'mutation'131	return feature_group_to_features_dict, feature_to_feature_group_dict132 133 134def get_color(feature_name: str, feature_group_dict: Mapping[str, List[str]]) -> str:135    """136    Determines the color for a given feature based on its group.137 138    Args:139        feature_name: Name of the feature.140        feature_group_dict: Dictionary mapping feature groups to feature names.141 142    Returns:143        A string representing the color associated with the feature's group.144    """145    color_mapping = {'mutation': '#984ea3', 'cna': '#f781bf', 'signature': '#4daf4a', 'clinical': '#377eb8'}146    for group, features in feature_group_dict.items():147        if feature_name in features:148            return color_mapping[group]149    return 'black'  # Default color if not found150 151def get_individual_pred_interpretation(shap_pred_sample_df: pd.DataFrame,152                                               feature_sample_df: pd.DataFrame,153                                               feature_group_to_features_dict: dict,154											   feature_to_feature_group_dict: dict,155                                               sample_info: str = None,156                                               filename: str = None,157                                               filepath: str = './others_prediction_explanation',158											   cohort_age_stats_path: str = './data/combined_cohort_age_stats.pkl',159                                               top_feature_num: int = 10,160                                               save_plot: bool = False):161    """162    Dynamic version of the function for individual prediction interpretation for a given tumor sample.163 164    Args:165        shap_pred_sample_df: DataFrame containing SHAP values for a given tumor sample.166        feature_sample_df: DataFrame containing feature values for a given tumor sample.167        feature_group_to_features_dict: Dictionary mapping feature groups to feature names.168        sample_info: Sample information to be displayed.169        filename: Filename to save the figure.170        top_feature_num: Number of top features to display.171    """172    # Initialize plot173    plt.rcParams.update({'font.size': 15, "font.family": "Arial"})174    fig, ax = plt.subplots()175 176	# remove top and right lines for bar graph177    ax.spines['right'].set_visible(False)178    ax.spines['top'].set_visible(False)179 180    # Sorting and selecting top features based on SHAP values181    all_features = sum(feature_group_to_features_dict.values(), [])182    sorted_features = sorted(all_features, key=lambda x: abs(shap_pred_sample_df.loc[x]), reverse=True)183    top_features = sorted(sorted_features[:top_feature_num], key=lambda x: shap_pred_sample_df.loc[x], reverse=False)184 185    # Preparing data for the bar chart186    top_feats_df = pd.DataFrame({187        'feat_name': top_features,188        'SHAP_val': shap_pred_sample_df.loc[top_features],189        'feat_val': feature_sample_df.loc[top_features],190        'color': [get_color(feat, feature_group_to_features_dict) for feat in top_features]191    })192 193    # Creating the bar chart194    ax.barh(top_feats_df['feat_name'], top_feats_df['SHAP_val'], color=top_feats_df['color'])195    ax.set_xlabel('SHAP Values')196    ax.set_title(sample_info)197    ax.set_yticks([])198    199    with open(cohort_age_stats_path, "rb") as fp:200        combined_cohort_age_stats = pickle.load(fp)201 202    # Dynamic positioning of feature names and values203    left_margin = ax.get_xlim()[0] * 1.1 -.1 # Calculate the left margin dynamically204    for i, (name, value) in enumerate(zip(top_feats_df['feat_name'], top_feats_df['feat_val'])):205        value_text = f'{int(value)}' if feature_to_feature_group_dict[name] == 'mutation' else f'{value:.2f}'206        # Feature name and value207        if name == 'Sex':208            value_text = 'Male' if value == 1.0 else 'Female'209        if name == 'Age':210            value_text = int(value * combined_cohort_age_stats['Std_mean'] + combined_cohort_age_stats['Age_mean'])211        if name[-3:] == 'CNA':212            value_text = f'{int(value)}'213 214        ax.text(left_margin, i, f'{name}: {value_text}', ha='right', va='center', fontsize=10)215 216    ax.text(left_margin, top_feature_num, f'feature: value', ha='right', va='center', fontsize=10)217 218    # Adding legend219    legend_elements = [Patch(facecolor=color, label=label) for label, color in zip(['Somatic Mut.', 'CNA events', 'Mutation Sig.', 'Age/Sex'], ['#984ea3', '#f781bf', '#4daf4a', '#377eb8'])]220    ax.legend(handles=legend_elements, title='Feature Groups')221 222    plt.tight_layout()223    if save_plot:224        plt.savefig(f'{filepath}/{filename}.svg')225        plt.savefig(f'{filepath}/{filename}.pdf')226    plt.show()227 228    return f'{filepath}/{filename}.svg'229 230def get_top_n_pred_and_shap(231		preds_df: pd.DataFrame,232		shaps: np.asarray,		233		features: List[str],		234		cancer_types: List[str],235		n: int=3) -> Mapping[str, Mapping[str, Any]]:236	"""237	Get top n predictions for each sample in the given DataFrame.238	Args:239		preds_df: DataFrame containing predictions.240		shaps: Array of SHAP values.241		features: List of features.242		n: Number of top predictions to return.243		cancer_types: List of cancer types.244	Returns:245		Dictionary containing top n predictions and SHAP values for each sample.246	"""247	# Filter only relevant columns (cancer types)248	df_chosen = preds_df[cancer_types]249	# Initialize the dictionary for SHAP values250	sample_id_to_top_n_pred_and_shaps_df = {}251	for sample_id, row in df_chosen.iterrows():252		# Getting the top n predictions sorted by their probabilities253		sorted_top_3 = row.nlargest(n).sort_values(ascending=False)254		# Get row index for sample_id255		row_idx = preds_df.index.get_loc(sample_id)256		# Store the top n predictions and their probabilities257		shap_dfs = []258		for pred_cancer in sorted_top_3.index:259			pred_cancer_idx = cancer_types.index(pred_cancer)260			shap_pred_cancer_df = pd.DataFrame(shaps[pred_cancer_idx][row_idx],261										   index=features,262										   columns=[pred_cancer]).T263			shap_dfs.append(shap_pred_cancer_df)264		shap_total_df = pd.concat(shap_dfs)265		sample_id_to_top_n_pred_and_shaps_df[sample_id] = {266			'top_n_predictions': sorted_top_3.to_dict(),267			'shap_values_df': shap_total_df268		}269	return sample_id_to_top_n_pred_and_shaps_df270 271def get_onconpc_prediction_explanations(query_ids: List[str], 272										preds_df: pd.DataFrame, 273										shaps: np.array, 274										df_features_genie: pd.DataFrame, 275										cancer_types_to_consider: List[str],276										filepath: str='./others_prediction_explanation',277										cohort_age_stats_path: str='./data/combined_cohort_age_stats.pkl',278										save_plot: bool=False,279										) -> List[Mapping[str, Any]]:280	"""281	Get OncoNPC predictions and generate SHAP-based explanation plots for multiple query IDs.282 283	Args:284		query_ids: List of IDs of the tumor samples to query.285		preds_df: DataFrame containing predictions.286		shaps: Array of SHAP values.287		df_features_genie: DataFrame containing features.288		cancer_types_to_consider: List of cancer types considered in the prediction.289		filepath: Path to save the explanation plots.290		save_plot: Whether to save the explanation plots.291	Returns:292		List of dictionaries containing prediction details and explanation plots for each query ID.293	"""294	results_dict = {}295	for query_id in query_ids:296		# Get OncoNPC prediction297		pred_prob = preds_df.at[query_id, 'max_posterior']298		pred_cancer = preds_df.at[query_id, 'cancer_type']299		pred_cancer_idx = cancer_types_to_consider.index(pred_cancer)300 301		# Get SHAP-based explanation for the prediction302		feature_sample_df = df_features_genie.loc[query_id]303		shap_pred_cancer_df = pd.DataFrame(shaps[pred_cancer_idx],304										   index=df_features_genie.index,305										   columns=df_features_genie.columns)306		shap_pred_sample_df = shap_pred_cancer_df.loc[query_id]307 308		# Information and plot generation309		sample_info = f'SAMPLE_ID: {query_id}\nPrediction: {pred_cancer}\nPrediction probability: {pred_prob:.3f}'310		feature_group_to_features_dict, feature_to_feature_group_dict = partition_feature_names_by_group(df_features_genie.columns)311		full_filename = get_individual_pred_interpretation(shap_pred_sample_df,312													 feature_sample_df,313													 feature_group_to_features_dict,314													 feature_to_feature_group_dict,315													 sample_info=sample_info,316													 filename=str(query_id),317													 filepath=filepath,318													 top_feature_num=10,319													 cohort_age_stats_path=cohort_age_stats_path,320													 save_plot=save_plot)321		# Store the results322		results_dict[query_id] = {323			'pred_prob': pred_prob,324			'pred_cancer': pred_cancer,325			'explanation_plot': full_filename326		}327	return results_dict328 329def get_onconpc_features_from_raw_data(df_patients_chosen: pd.DataFrame, 330									   df_samples_chosen: pd.DataFrame, 331									   df_mutations_chosen: pd.DataFrame, 332									   df_cna_chosen: pd.DataFrame, 333									   features_onconpc_path: str='data/features_onconpc.pkl',334									   combined_cohort_age_stats_path: str='data/combined_cohort_age_stats.pkl',335									   mut_sig_weights_filepath: str='data/mutation_signatures/sigProfiler*.csv'336									   ) -> Tuple[pd.DataFrame, pd.DataFrame]:337	"""338	Process raw GENIE data for OncoNPC prediction inference.339 340	Args:341		df_patients_chosen: DataFrame containing chosen patient data.342		df_samples_chosen: DataFrame containing chosen sample data.343		df_mutations_chosen: DataFrame containing chosen mutation data.344		df_cna_chosen: DataFrame containing chosen copy number alteration data.345		features_onconpc_path: Path to OncoNPC features file.346		combined_cohort_age_stats_path: Path to combined cohort age statistics file.347		weights_filepath: Path to mutation signatures weights file.348 349	Returns:350		Processed DataFrame ready for OncoNPC prediction inference.351	"""352	# Re-shape df_cna_chosen353	df_cna_reshaped = df_cna_chosen.set_index(['Hugo_Symbol']).T.copy()354	# Merge patient and sample data355	df_patients_merged = pd.merge(df_patients_chosen, df_samples_chosen, how='right', on='PATIENT_ID')356 357	# Process mutation data358	df_trinuc_feats_genie = get_snv_in_trinuc_context(df_mutations_chosen,359												   sample_id_col='Tumor_Sample_Barcode',360												   chromosome_col='Chromosome',361												   start_pos_col='Start_Position',362												   ref_allele_col='Reference_Allele',363												   alt_allele_col='Tumor_Seq_Allele2',364												   config='genie')365 366	# Obtain mutation signatures367	df_mut_sigs_genie = obtain_mutation_signatures(df_trinuc_feats_genie, weights_filepath=mut_sig_weights_filepath)368	# Preprocess features and labels369	df_features_genie, df_labels_genie = pre_process_features_genie(df_mutations_chosen, 370																 df_cna_reshaped, 371																 df_mut_sigs_genie, 372																 df_patients_merged)373	# Zero-pad missing features374	with open(features_onconpc_path, "rb") as fp:375		features_onconpc = pickle.load(fp)376	df_features_genie_final = zero_pad_missing_features(df_features_genie, features_onconpc)377	# Standardize feature names378	df_features_genie_final.columns = standardize_feat_names(df_features_genie_final.columns)379	# Standardize Age feature380	with open(combined_cohort_age_stats_path, "rb") as fp:381		combined_cohort_age_stats = pickle.load(fp)382	df_features_genie_final['Age'] = (df_features_genie_final['Age'] - combined_cohort_age_stats['Age_mean']) / combined_cohort_age_stats['Std_mean']383	return df_features_genie_final, df_labels_genie384 385def get_snv_in_trinuc_context(df_mutations: pd.DataFrame,386							  sample_id_col: str,387							  chromosome_col: str,388							  start_pos_col: str,389							  ref_allele_col: str,390							  alt_allele_col: str,391							  config: Optional[str]='genie') -> pd.DataFrame:392	"""393	Processes the given mutations DataFrame to get SNV in trinucleotide context.394 395	Args:396		df_mutations: DataFrame containing mutation data.397		sample_id_col: Column name for sample IDs.398		chromosome_col: Column name for chromosome.399		start_pos_col: Column name for start position.400		ref_allele_col: Column name for reference allele.401		alt_allele_col: Column name for alternative allele.402	Returns:403		DataFrame with SNV in trinucleotide context.404	"""405	# Specify the columns to be chosen406	columns_chosen = [sample_id_col, chromosome_col, start_pos_col, ref_allele_col, alt_allele_col]407 408	# Process the DataFrame409	df_mutations_chosen_trinuc = df_mutations[columns_chosen].copy()410	if config == 'genie':411		df_mutations_chosen_trinuc[chromosome_col] = 'chr' + df_mutations_chosen_trinuc[chromosome_col].astype(str)412	df_mutations_chosen_trinuc = df_mutations_chosen_trinuc[~df_mutations_chosen_trinuc[chromosome_col].str.contains('GL|chrMT')]413	# Change the data types for R processing414	df_mutations_chosen_trinuc[sample_id_col] = df_mutations_chosen_trinuc[sample_id_col].astype('str')415	# R function code416 417	r_code = '''418	function(mutationData, sample_id, chr, pos, ref, alt) {419		library(deconstructSigs)420		sigs.input <- mut.to.sigs.input(mut.ref = mutationData, 421										sample.id = sample_id,422										chr = chr, 423										pos = pos, 424										ref = ref, 425										alt = alt)426 427		# Filter samples with low mutations428		sigs.input <- sigs.input[rowSums(sigs.input) >= 1,]429		return(sigs.input)430	}431	'''432	# Convert the DataFrame for R processing433	with conversion.localconverter(default_converter):434 435		with localconverter(robjects.default_converter + pandas2ri.converter):436			r_df_mutations = robjects.conversion.py2rpy(df_mutations_chosen_trinuc)437 438	# Load and call the R function439		r_function = robjects.r(r_code)440	with conversion.localconverter(default_converter):441 442		with localconverter(robjects.default_converter + pandas2ri.converter):443			df_trinuc_feats = robjects.conversion.rpy2py(r_function(r_df_mutations, 444																sample_id_col, 445																chromosome_col, 446																start_pos_col, 447																ref_allele_col, 448																alt_allele_col))449	if config=='profile_dfci':450		df_trinuc_feats.index = df_trinuc_feats.index.astype('float')451	return df_trinuc_feats452 453def obtain_mutation_signatures(df_trinuc_feats: pd.DataFrame,454							   weights_filepath: str='./data/mutation_signatures/sigProfiler*.csv'455							   ) -> pd.DataFrame:456	"""Transforms tri-nucleotide features into mutation signature based features.457 458	Args:459		df_trinuc_feats: DataFrame containing tri-nucleotide features.460	Returns:461		mut_sig_based_df: DataFrame containing mutation signature based features.462	"""463	# Load Mutational Signatures COSMIC data.464	file_names = glob.glob(weights_filepath)465	# Check if the directory exists.466	if len(file_names) == 0:467		raise ValueError('No mutation signatures data.')468	file_names_key = [file_name.split('_')[-1][:-4] for file_name in file_names]469	sbs_conversion_df_dict = {}	470	for i in range(len(file_names)):471		df = pd.read_csv(file_names[i])472		df.dropna(inplace = True)473		# process base names s.t. they compatible with actual feature names474		new_name_list = []475		for change, subtype in zip(df.Type.values, df.Subtype.values):476			new_name = subtype[0] + '[' + change + ']' + subtype[2]477			new_name_list.append(new_name)478 479		df['Type_new'] = new_name_list480		df.set_index('Type_new', inplace = True)481		# check if column contains human-based experiment results482		merged_colum_names = ''.join(df.columns)483		if 'GRCh38' in merged_colum_names:484			df['values_oi'] = df[file_names_key[i] + '_GRCh38'].values485		elif 'GRCh37' in merged_colum_names:486			df['values_oi'] = df[file_names_key[i] + '_GRCh37'].values487		else:488			raise ValueError('No human-based experimental results.')489		sum_vals = np.sum(df.values_oi.values)490		if abs(sum_vals - 1.0) > 0.1:491			raise ValueError('Values do not sum up to 1.')492		sbs_conversion_df_dict[file_names_key[i]] = df493	mut_sig_based_df = pd.DataFrame([], index = df_trinuc_feats.index)494	for key, mut_df in sbs_conversion_df_dict.items():495		common_substitute_feats = list(set(mut_df.index) & set(df_trinuc_feats.columns))496		if len(common_substitute_feats) != 96:497			raise ValueError('Number of substitution features does not sum up to 96.')498		# Get mutation signature values499		mut_sig_vals = np.matmul(df_trinuc_feats[common_substitute_feats].values,500								 mut_df.loc[common_substitute_feats].values_oi.values)501		mut_sig_based_df[key] = mut_sig_vals502	return mut_sig_based_df503 504def pre_process_features_genie(df_mutations: pd.DataFrame,505							   df_cna: pd.DataFrame,506							   df_mutation_signatures: pd.DataFrame,507							   df_patients: pd.DataFrame,508							   cancer_types: Optional[List[str]]=None, 509							   id_column: str = 'SAMPLE_ID',510							   cancer_type_column: str='CANCER_TYPE') -> Tuple[pd.DataFrame, pd.DataFrame]:511	"""Pre-process genetics data to create feature df for GENIE.512	513	Args:514		df_mutations: DataFrame containing mutation data.515		df_cna: DataFrame containing CNA data.516		df_mutation_signatures: DataFrame containing mutation signature data.517		df_patients: DataFrame containing patient data.518		cancer_types: List of cancer types.519		id_column: Column name for sample IDs.520		cancer_type_column: Column name for cancer types.521	Returns:522		df_features_merged_final: DataFrame containing merged features.523		df_labels_final: DataFrame containing labels.524	"""525	# Get mutation features.526	sample_ids = df_patients[id_column].values527	df_mutations_chosen = df_mutations.loc[df_mutations.Tumor_Sample_Barcode.isin(sample_ids)]528	mutation_gene_names = np.unique(df_mutations_chosen.Hugo_Symbol.values)529	df_mutation_feature = pd.DataFrame(0, columns=mutation_gene_names, index=sample_ids)530	for entry in df_mutations_chosen[['Tumor_Sample_Barcode', 'Hugo_Symbol']].values:531		sample_id = entry[0]532		gene = entry[1]533		if gene in df_mutation_feature.columns:534			# Count the number of mutation frequency in each gene.535			df_mutation_feature.at[sample_id, gene] = df_mutation_feature.at[sample_id, gene] + 1536	# Add 'mut' at the end of each mutation feature537	df_mutation_feature.columns = [gene + '_mut' for gene in mutation_gene_names]538	539	# Get CNA features.540	df_cna_chosen = df_cna.loc[list(set(df_cna.index) & set(sample_ids))]541	df_cna_chosen.fillna(0, inplace = True)542 543	# Obtain age/sex features.544	df_patients_id_indexed = df_patients.set_index('SAMPLE_ID')545	df_sex_age_feature = df_patients_id_indexed[['SEX', 'AGE_AT_SEQ_REPORT']].copy()546	sex_list = []547	age_nan_indices = []548	age_list = []549	for sample_id, (sex, age) in zip(df_sex_age_feature.index, df_sex_age_feature.values):550		if sex == 'Male':551			sex_list.append(1)552		elif sex == 'Female':553			sex_list.append(-1)554		else: # Not reported555			sex_list.append(0)556		if pd.isnull(age):557			age_nan_indices.append(sample_id)558			age_list.append(0) # append zero for now559		elif '>89' in str(age):560			age_list.append(int(age[1:]))561		elif '<18' in str(age):562			age_list.append(int(age[1:]))563		else:564			age_list.append(int(age))565	df_sex_age_feature['SEX'] = sex_list566	df_sex_age_feature['AGE_AT_SEQ_REPORT'] = age_list567	df_sex_age_feature.columns = ['Sex', 'Age']568	# Merge all feature dfs using common sample IDs.569	common_ids = list(set(df_mutation_feature.index) & set(df_cna_chosen.index) & set(df_mutation_signatures.index)570				  & set(df_sex_age_feature.index))571	dfs_list = [df_mutation_feature, df_cna_chosen, df_mutation_signatures, df_sex_age_feature]572	df_features_merged = merge_dfs(dfs_list, common_ids)573	# Exclude age NaN indices.574	df_features_merged_final = df_features_merged.loc[list(set(df_features_merged.index) - set(age_nan_indices))]575	df_labels = df_patients_id_indexed.loc[df_features_merged_final.index][cancer_type_column]576	if cancer_types is not None:577		df_labels_final = pd.DataFrame([cancer_types.index(val) for val in df_labels.values],578									columns = ['cancer_label'], index = df_labels.index)579		df_labels_final['cancer_type'] = df_labels.values	580		return df_features_merged_final, df_labels_final581	else:582		return df_features_merged_final, df_labels583 584def pre_process_features_dfci(df_mutations: pd.DataFrame,585							  df_cna: pd.DataFrame,586							  df_mutation_signatures: pd.DataFrame,587							  df_patients: pd.DataFrame,588							  cancer_types: List[str], 589							  cup_samples_ids: Optional[List[str]]=None,590							  id_column: str = 'UNIQUE_SAMPLE_ID'591							  ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:592	"""Pre-process genetics data to create feature df for DFCI data.593	594	Args:595		df_mutations: DataFrame containing mutation data.596		df_cna: DataFrame containing CNA data.597		df_mutation_signatures: DataFrame containing mutation signature data.598		df_patients: DataFrame containing patient data.599		cancer_types: List of cancer types.600		cup_samples_ids: List of CUP sample IDs.601		id_column: Column name for sample IDs.602	Returns:603		df_features_merged_ckps: DataFrame containing merged features for CKP samples.604		df_labels_final: DataFrame containing labels for CKP samples.605		df_features_merged_cups: DataFrame containing merged features for CUP samples.606	"""607	# Get mutation features.608	sample_ids = df_patients[id_column].values609	df_mutations_chosen = df_mutations.loc[df_mutations[id_column].isin(sample_ids)]610	mutation_gene_names = np.unique(df_mutations_chosen.CANONICAL_GENE.values)611	df_mutation_feature = pd.DataFrame(0, columns=mutation_gene_names, index=sample_ids)612	for entry in df_mutations_chosen[['UNIQUE_SAMPLE_ID', 'CANONICAL_GENE']].values:613		sample_id = entry[0]; gene = entry[1];614		if gene in df_mutation_feature.columns:615			# Count the number of mutation frequency in each gene.616			df_mutation_feature.at[sample_id, gene] = df_mutation_feature.at[sample_id, gene] + 1617	# Add 'mut' at the end of each mutation feature618	df_mutation_feature.columns = [gene + '_mut' for gene in mutation_gene_names]619	# Get CNA features.620	df_cna_chosen = df_cna.loc[list(set(df_cna.index) & set(sample_ids))]621	df_cna_chosen.fillna(0, inplace = True)622 623	# Obtain age/sex features.624	df_patients_id_indexed = df_patients.set_index(id_column)625	df_sex_age_feature = df_patients_id_indexed[['Sex', 'Age']].copy()626	sex_list = []627	for sample_id, (sex, age) in zip(df_sex_age_feature.index, df_sex_age_feature.values):628		if sex == 'MALE':629			sex_list.append(1)630		elif sex == 'FEMALE':631			sex_list.append(-1)632		else: # Not reported633			sex_list.append(0)634	df_sex_age_feature['Sex'] = sex_list635	df_sex_age_feature.columns = ['Sex', 'Age']636	# Merge all features using the same sample IDs.637	common_ids = list(set(df_mutation_feature.index) & set(df_cna_chosen.index) &638				  set(df_mutation_signatures.index) & set(df_sex_age_feature.index))639	dfs_list = [df_mutation_feature, df_cna_chosen, df_mutation_signatures, df_sex_age_feature]640	df_features_merged = merge_dfs(dfs_list, common_ids)641	# Make gene names consistent across GENIE and PROFILE.642	profile_old_new_gene_mapping = {}643	profile_old_new_gene_mapping['C17ORF70'] = 'FAAP100' 644	profile_old_new_gene_mapping['C17orf70_mut'] = 'FAAP100_mut'645	profile_old_new_gene_mapping['C19ORF40'] = 'FAAP24' 646	profile_old_new_gene_mapping['C19orf40_mut'] = 'FAAP24_mut'647	profile_old_new_gene_mapping['LOC96610_mut'] = 'BMS1P20_mut'648	profile_old_new_gene_mapping['LOC729991-MEF2B_mut'] = 'MEF2BNB-MEF2B_mut'649	profile_old_new_gene_mapping['C1orf86_mut'] = 'FAAP20_mut'650	profile_old_new_gene_mapping['C1ORF86'] = 'FAAP20'651	profile_old_new_gene_mapping['GNB2L1_mut'] = 'RACK1_mut'652	profile_old_new_gene_mapping['GNB2L1'] = 'RACK1'653	df_features_merged.rename(columns=profile_old_new_gene_mapping,654								   inplace=True)655	df_features_merged_ckps = (df_features_merged656							   .loc[list(set(df_features_merged.index) - set(cup_samples_ids))])657	df_features_merged_cups = (df_features_merged658							   .loc[list(set(df_features_merged.index) & set(cup_samples_ids))])659	df_labels = df_patients_id_indexed.loc[df_features_merged_ckps.index]['CANCER_TYPE']660	df_labels_final = pd.DataFrame([cancer_types.index(val) for val in df_labels.values],661								   columns=['cancer_label'], index=df_labels.index)662	df_labels_final['cancer_type'] = df_labels.values663	return df_features_merged_ckps, df_labels_final, df_features_merged_cups664 665def zero_pad_missing_features(df: pd.DataFrame,666							  feature_list: List[str]) -> pd.DataFrame:667	"""668	Zero pad missing features.669	Args:670		df: DataFrame containing features.671		feature_list: List of features to include.672	Returns:673		df: DataFrame containing zero-padded features.674	"""675	# Create a dictionary for new columns to add676	new_columns = {feature: [0] * len(df) for feature in feature_list if feature not in df.columns}677 678	# Concatenate new columns with the original DataFrame679	df = pd.concat([df, pd.DataFrame(new_columns, index=df.index)], axis=1)680	return df[feature_list]681 682def merge_dfs(df_list: pd.DataFrame, ids_common: List[str]) -> pd.DataFrame:683	"""Merges a list of dataframes using common sample IDs.684	685	Args:686		df_list: List of dataframes to merge.687		ids_common: Common sample IDs.688	Returns:689		df_merged_to_return: Merged dataframe.690	"""691	for idx, df in enumerate(df_list):692		if idx == 0:693			df_merged_to_return = df.loc[ids_common]694		else:695			df_merged_to_return = pd.merge(df_merged_to_return, df.loc[ids_common], how = 'left', left_index = True, right_index = True)696	return df_merged_to_return697 698def standardize_feat_names(curr_feat_names: List[str]) -> List[str]:699	"""700	Standardize feature names.701	Args:702		curr_feat_names: list of feature names703	Returns:704		new_feat_names: list of standardized feature names705	"""706	new_feat_names = []707	for feat in curr_feat_names:708		if '_mut' in feat:709			new_feat_names.append(feat.replace('_mut', ''))710		elif 'AGE' in feat or 'Age' in feat:711			new_feat_names.append('Age')712		elif 'GENDER' in feat or 'Sex' in feat:713			new_feat_names.append('Sex')714		elif 'SBS' in feat:715			new_feat_names.append(feat)716		else:   717			new_feat_names.append(feat + ' CNA')718	return new_feat_names