import numpy as np import pandas as pd import bokeh.io import bokeh.plotting import bokeh.palettes from bokeh.transform import factor_cmap import datashader import holoviews as hv import holoviews.operation.datashader hv.extension("bokeh") import warnings # Suppress FutureWarning messages warnings.simplefilter(action='ignore',) cr9114_dict = { 'ab_name' : 'CR9114', 'files' : { 'ESM-1v Ab-Ag': 'output/ab_mutagenesis_expts/cr9114/cr9114_esm1vAbAg_exp_data_maskMargLabeled.csv', 'ESM-1v Ab only': 'output/ab_mutagenesis_expts/cr9114/cr9114_esm1vbothchains_exp_data_maskMargLabeled.csv', 'ESM-1v Ab VH only': 'output/ab_mutagenesis_expts/cr9114/cr9114_exp_data_maskMargLabeled.csv', }, 'dms_df' : pd.read_csv('data/ab_mutagenesis_expts/cr9114/cr9114_exp_data.csv', dtype = {'genotype': str}, ).rename(columns={'h1_mean': 'H1', 'h3_mean' : 'H3'}), 'ag_columns': ['H1', 'H3'], 'expt_type': 'Combinatorial Mutagenesis for Affinity', 'palette': bokeh.palettes.Spectral6 } cr6261_dict = { 'ab_name' : 'CR6261', 'files' : { 'ESM-1v Ab-Ag': 'output/ab_mutagenesis_expts/cr6261/cr6261_esm1vAbAg_exp_data_maskMargLabeled.csv', 'ESM-1v Ab only': 'output/ab_mutagenesis_expts/cr6261/cr6261_esm1vbothchains_exp_data_maskMargLabeled.csv', 'ESM-1v Ab VH only': 'output/ab_mutagenesis_expts/cr6261/cr6261_exp_data_maskMargLabeled.csv', }, 'dms_df' : pd.read_csv('data/ab_mutagenesis_expts/cr6261/cr6261_exp_data.csv', dtype={'genotype': str}, ).rename(columns={'h1_mean': 'H1', 'h9_mean' : 'H9'}), 'ag_columns': ['H1', 'H9'], 'expt_type': 'Combinatorial Mutagenesis for Affinity', 'palette': bokeh.palettes.Dark2_6, } g6LC_dict = { 'ab_name' : 'g6', 'files' : { 'LM Ab-Ag': 'output/ab_mutagenesis_expts/g6/g6Lc_esm1vAbAg_exp_data_maskMargLabeled.csv', 'LM Ab only': 'output/ab_mutagenesis_expts/g6/g6Lc_esm1vbothchains_exp_data_maskMargLabeled.csv', 'LM Ab VH/VL only': 'output/ab_mutagenesis_expts/g6/g6Lc_exp_data_maskMargLabeled.csv', }, 'dms_df' : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv'), 'ag_columns': ['norm_binding'], 'expt_type': 'Deep Mutational Scan for Binding', 'palette': bokeh.palettes.Pastel1_6, 'chain': 'VL' } g6HC_dict = { 'ab_name' : 'g6', 'files' : { 'LM Ab-Ag': 'output/ab_mutagenesis_expts/g6/g6Hc_esm1vAbAg_exp_data_maskMargLabeled.csv', 'LM Ab only': 'output/ab_mutagenesis_expts/g6/g6Hc_esm1vbothchains_exp_data_maskMargLabeled.csv', 'LM Ab VH/VL only': 'output/ab_mutagenesis_expts/g6/g6Hc_exp_data_maskMargLabeled.csv', }, 'dms_df' : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv'), 'ag_columns': ['norm_binding'], 'expt_type': 'Deep Mutational Scan for Binding', 'palette': bokeh.palettes.Pastel1_4, 'chain': 'VH' } def apply_mask_and_average(scores, genotype): #scores here should already be a in log space if '1' in genotype: masked_list = [] for s, mask_char in zip(scores, genotype): if mask_char == '1': masked_list.append(s) multi_avg = np.mean(masked_list) else: #if genotype is all zeros = wt multi_avg = 0 return multi_avg def transform_single_to_multi( dms_df, singleMuts_df, condition, sort_col, ascending ): multi_scores = [] #use input parameter as method bc key is kwarg in pd.sort_values #method should be either esm1v or abysis or AbLang #for cr sorts the single mutation dataframe in residue order, ie binary '10000' before '00010'. ascending = False #for g6 sorts the single mutation dataframe in position order, ie residue index. ascending = True singleMuts_sorted = singleMuts_df.sort_values(sort_col, key=lambda x: x.astype(int), ascending=ascending) singles_scores = singleMuts_sorted[condition].to_list() for g in dms_df['genotype']: multi_scores.append(apply_mask_and_average(singles_scores, g )) return multi_scores #get melted correlations matrix of structure input x target antigen, for a given antibody def get_corr(ab_name, files, dms_df, ag_columns, dropLOQ = False ): #retreive correlations from abysis and esm1v scores for key, filepath in files.items(): singleMuts_df = pd.read_csv(filepath) # Get the columns for all esm1v models esm_columns = [col for col in singleMuts_df.columns if col.startswith('esm')] # Calculate the average for each esm column esm_avg_values = singleMuts_df[esm_columns].mean(axis=1) singleMuts_df[key] = esm_avg_values if ab_name == 'g6': #g6 data set is only single mutations dms_df[key] = esm_avg_values else: #cr datasets are multiple mutations dms_df[key] = transform_single_to_multi(dms_df, singleMuts_df, key, 'genotype', ascending = False) conditions = list(files.keys()) if not dropLOQ: correlations = dms_df[ag_columns + conditions].corr(method='spearman') correlations = correlations.drop(conditions, axis= 1) correlations = correlations.drop(ag_columns, axis= 0) # Melt the correlations DataFrame and rename the index column melted_correlations = ( correlations .reset_index() .melt(id_vars='index', var_name='Target', value_name='Correlation') .rename(columns={'index': 'Input'}) ) melted_correlations['Method'] = ['LM' for i in melted_correlations['Input']] return melted_correlations #if we wish to drop the samples on the lower limit of quantitation, drop samples from each ag group individually and compute correlation with IF log_like else: all_ag_melted = pd.DataFrame({}) for ag in ag_columns: filt_df = dms_df[dms_df[ag] > min(dms_df[ag])] correlations = filt_df[[ag] + conditions].corr(method='spearman') correlations = correlations.drop(conditions, axis= 1) correlations = correlations.drop(ag, axis= 0) # Melt the correlations DataFrame and rename the index column melted_correlations = ( correlations .reset_index() .melt(id_vars='index', var_name='Target', value_name='Correlation') ).rename(columns={'index': 'Input'}) all_ag_melted = pd.concat([all_ag_melted, melted_correlations], ignore_index= True) all_ag_melted['Method'] = ['LM' for i in all_ag_melted['Input']] return all_ag_melted #get melted correlations matrix of structure input x target antigen, for a given antibody def get_g6_corr(g6Hc_dict, g6LC_dict, dropLOQ = False ): ab_name = 'g6' ag_columns = g6LC_dict['ag_columns'] vh_and_vl_dms_df = pd.DataFrame({}) for d in [g6HC_dict, g6LC_dict]: files = d['files'] dms_df = d['dms_df'] #retreive correlations from abysis, ablang, and esm1v scores for key, filepath in files.items(): singleMuts_df = pd.read_csv(filepath) # Get the columns for all esm1v models esm_columns = [col for col in singleMuts_df.columns if col.startswith('esm')] # Calculate the average for each esm column esm_avg_values = singleMuts_df[esm_columns].mean(axis=1) singleMuts_df[key] = esm_avg_values if ab_name == 'g6': #g6 data set is only single mutations dms_df[key] = esm_avg_values vh_and_vl_dms_df = pd.concat([vh_and_vl_dms_df, dms_df], ignore_index= True) conditions = list(files.keys()) # if not dropLOQ: correlations = vh_and_vl_dms_df[ag_columns + conditions].corr(method='spearman') correlations = correlations.drop(conditions, axis= 1) correlations = correlations.drop(ag_columns, axis= 0) # Melt the correlations DataFrame and rename the index column melted_correlations = ( correlations .reset_index() .melt(id_vars='index', var_name='Target', value_name='Correlation') .rename(columns={'index': 'Input'}) ) melted_correlations['Method'] = ['LM' for i in melted_correlations['Input']] melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations) return melted_correlations #if we wish to drop the samples on the lower limit of quantitation, drop samples from each ag group individually and compute correlation with IF log_like else: all_ag_melted = pd.DataFrame({}) for ag in ag_columns: filt_df = vh_and_vl_dms_df[vh_and_vl_dms_df[ag] > min(vh_and_vl_dms_df[ag])] correlations = filt_df[[ag] + conditions].corr(method='spearman') correlations = correlations.drop(conditions, axis= 1) correlations = correlations.drop(ag, axis= 0) # Melt the correlations DataFrame and rename the index column melted_correlations = ( correlations .reset_index() .melt(id_vars='index', var_name='Target', value_name='Correlation') ).rename(columns={'index': 'Input'}) all_ag_melted = pd.concat([all_ag_melted, melted_correlations], ignore_index= True) all_ag_melted['Method'] = ['LM' for i in melted_correlations['Input']] melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations) return all_ag_melted #plot correlation bars for cr antibodies: cr6261 and cr9114 def plot_hbar(title, melted_correlations, palette = bokeh.palettes.Spectral6, compare = 'model'): if compare == 'model': melted_correlations.replace('Ab-Ag','InverseFolding', inplace = True) color_by = 'Input' elif compare == 'IF': color_by = 'Method' else: color_by = 'Target' melted_correlations['cats'] = melted_correlations.apply(lambda x: (x["Target"], x["Input"]), axis = 1) factors = list(melted_correlations.cats)[::-1] p = bokeh.plotting.figure( height=340, width=440, x_axis_label="Spearman Correlation", x_range=[0, 1], y_range=bokeh.models.FactorRange(*factors), tools="save", title = title ) p.hbar( source=melted_correlations, y="cats", right="Correlation", height=0.6, line_color = 'black', legend_field = color_by, # use the palette to colormap based on the the x[1:2] values # fill_color=factor_cmap( # color_by, # palette= palette, # factors = list(melted_correlations[color_by].unique()), # start=1, # end=3) fill_color=bokeh.palettes.Dark2_6[1] ) #labels_df = melted_correlations[melted_correlations['Input'].isin(['Ab-Ag','ESM-1v', 'abYsis', 'InverseFold'])] labels_df = melted_correlations labels_df['corr_str'] = labels_df['Correlation'].apply(lambda x: round(x, 2)).astype(str) labels_source = bokeh.models.ColumnDataSource(labels_df) labels = bokeh.models.LabelSet(x='Correlation', y='cats', text='corr_str',text_font_size = "10px", x_offset=12, y_offset=-5, source=labels_source, render_mode='canvas') p.ygrid.grid_line_color = None p.y_range.range_padding = 0.1 p.add_layout(labels) p.legend.visible = False p.output_backend = "svg" return p if __name__ == '__main__': datasets = [cr9114_dict, cr6261_dict, (g6HC_dict, g6LC_dict) ] #compute with and without points on lower limit of quantitation ignored for dropLOQ in [False]: all_corr_plots = [] for d in datasets: if type(d) is tuple: g6Hc, g6Lc = d title = g6Hc['ab_name'] + ', ' + g6Hc['expt_type'] g6_combined = get_g6_corr(g6Hc, g6Lc, dropLOQ= dropLOQ) all_corr_plots.append(plot_hbar(title, g6_combined, g6Lc['palette'], 'IF')) else: corr_df = get_corr(d['ab_name'], d['files'], d['dms_df'], d['ag_columns'], dropLOQ= dropLOQ) title = d['ab_name'] + ', ' + d['expt_type'] all_corr_plots.append(plot_hbar(title, corr_df, d['palette'], 'Target')) all_corr_fname = f"output/ab_mutagenesis_expts/esm1v_benchmarks{'_dropLOQ' if dropLOQ else ''}.html" bokeh.plotting.output_file(all_corr_fname) bokeh.io.show(bokeh.layouts.gridplot(all_corr_plots, ncols = len(all_corr_plots)))