import argparse from pathlib import Path from biotite.sequence.io.fasta import FastaFile, get_sequences from Bio import pairwise2 import pandas as pd import shutil import os import numpy as np import pandas as pd import bokeh.io import bokeh.plotting import bokeh.palettes from bokeh.transform import factor_cmap from bokeh.io import export_svg import iqplot from colorcet import glasbey_category10 import subprocess bla_dict = { 'protein':'bla', 'pdbfile':'bla_1m40_a.pdb', 'chain':'A', 'dmsfile': 'dms_bla.csv', 'fitness_col': 'DMS_amp_2500_(b)', 'threshold' : 0.01 } calm1_dict = { 'protein':'CALM1', 'pdbfile':'calm1_5v03_r.pdb', 'chain':'R', 'dmsfile': 'dms_calm1.csv', 'fitness_col':'DMS', 'threshold': 1 } haeiiim_dict = { 'protein':'haeIIIM', 'pdbfile':'haeiiim_3ubt_b.pdb', 'chain':'B', 'dmsfile': 'dms_haeiiim.csv', 'fitness_col':'DMS_G3', 'threshold' : 0.01 } gal4_dict = { 'protein':'GAL4', 'pdbfile':'gal4_3coq_b.pdb', 'chain':'B', 'dmsfile': 'dms_gal4.csv', 'fitness_col':'DMS_nonsel_24', 'threshold': 1 } hras_dict = { 'protein':'HRAS', 'pdbfile':'hras_2ce2_x.pdb', 'chain':'X', 'dmsfile': 'dms_hras.csv', 'fitness_col':'DMS_unregulated', 'threshold': 1 } mapk1_dict = { 'protein':'MAPK1', 'pdbfile':'mapk1_4zzn_a.pdb', 'chain':'A', 'dmsfile': 'dms_mapk1.csv', 'fitness_col':'DMS_VRT', 'threshold': 1 } tpk1_dict = { 'protein':'TPK1', 'pdbfile':'tpk1_3s4y_a.pdb', 'chain':'A', 'dmsfile': 'dms_tpk1.csv', 'fitness_col':'DMS', 'threshold': 1 } tpmt_dict = { 'protein':'TPMT', 'pdbfile':'tpmt_2bzg_a.pdb', 'chain':'A', 'dmsfile': 'dms_tpmt.csv', 'fitness_col':'DMS', 'threshold': 1 } ube2i_dict = { 'protein':'UBE2I', 'pdbfile':'ube2i_5f6e_a.pdb', 'chain':'A', 'dmsfile': 'dms_tpmt.csv', 'fitness_col':'DMS', 'threshold': 1 } ubi4_dict = { 'protein':'UBI4', 'pdbfile':'ubi4_4q5e_b.pdb', 'chain':'B', 'dmsfile': 'dms_ubi4.csv', 'fitness_col':'DMS_limiting_(b)', 'threshold': 1 } studies = [ bla_dict, calm1_dict, gal4_dict, haeiiim_dict, hras_dict, tpmt_dict, tpk1_dict, mapk1_dict, ube2i_dict, ubi4_dict, ] def score_if(prot_dict): prefix = 'data/dms/' fasta_outpath = f"dms_libs/{prot_dict['pdbfile'][:-4]}_dmsLib.fasta" score_if_cmd = [ 'python', 'model/score_log_likelihoods.py', prefix + f"structures/{prot_dict['pdbfile']}", '--seqpath', prefix + fasta_outpath, '--outpath', f"output/dms/if_results/dms_{prot_dict['protein']}_if.csv" ,'--chain', prot_dict['chain'] ] subprocess.run(score_if_cmd, check=True) def get_stats(prot_dict, n, top_per): top_per = top_per/100 prefix = 'output/dms/' esm1v_outpath = 'dms_esm1vScored/' + f"dms_{prot_dict['protein']}" + '_maskMarginals.csv' if_results_df = pd.read_csv( prefix + f"if_results/dms_{prot_dict['protein']}_if.csv") dms_df = pd.read_csv(prefix + esm1v_outpath) unprofiled_df = dms_df[dms_df[prot_dict['fitness_col']].isna()] unprofiled_variants = unprofiled_df['variant'].to_list() dms_df = dms_df.dropna(subset = [prot_dict['fitness_col']]) pop_n = len(dms_df) if_results_df = if_results_df[~if_results_df['seqid'].isin(unprofiled_variants)] assert (len(if_results_df) == len(dms_df)) threshold = dms_df[prot_dict['fitness_col']].quantile(1-top_per) highFit_subset_variants = dms_df[dms_df[prot_dict['fitness_col']] >= threshold]['variant'] n_highFit = len(highFit_subset_variants) dms_df['esm1v'] = dms_df.loc[:, dms_df.columns.str.startswith('esm1v')].mean(axis=1) top_esm1v_df = dms_df.sort_values(by = 'esm1v', ascending = False)[:n] top_esm1v_variants = top_esm1v_df['variant'].to_list() top_if_df = if_results_df.sort_values(by = 'log_likelihood', ascending = False)[:n] top_if_variants = top_if_df['seqid'].to_list() esm1v_hits_df = top_esm1v_df[top_esm1v_df['variant'].isin(highFit_subset_variants)] if_hits_df = top_if_df[top_if_df['seqid'].isin(highFit_subset_variants)] n_esm1v_hits = len(esm1v_hits_df) n_if_hits = len(if_hits_df) if_hit_enrich = (n_if_hits / n) / top_per esm1v_hit_enrich = (n_esm1v_hits / n) / top_per if_hit_rate = (n_if_hits / n) esm1v_hit_rate = (n_esm1v_hits / n) return (prot_dict['protein'], pop_n, n_highFit, n_if_hits, n_esm1v_hits, if_hit_enrich, esm1v_hit_enrich, if_hit_rate, esm1v_hit_rate ) def plot_bars(melted_df): p = bokeh.plotting.figure( height=350, width=1100, y_axis_label="High Fitness \n Prediction Precision", x_axis_label = "Functional Percentile Threshold \n for High Fitness Classification", x_range=bokeh.models.FactorRange(*factors, group_padding = 1.2), tools="save", title='' ) p.output_backend = "svg" p.vbar( source=melted_df, x="cats", top="Hit Rate", width = 1, line_color='black', alpha = 'alpha', legend_field='legend_label', fill_color=bokeh.transform.factor_cmap( 'Method', palette=['#999999', '#43a2ca'], factors=list(melted_df['Method'].unique()), start=1, end=2 ) ) p.xgrid.grid_line_color = None p.ygrid.grid_line_color = None p.x_range.range_padding = 0.03 p.legend.location = "top_right" p.legend.spacing = 20 p.legend.label_text_font_size = "11pt" p.legend.label_text_color = "black" p.legend.orientation = "horizontal" p.xaxis.major_label_orientation = 1.2 p.xaxis.separator_line_alpha = 0 p.xaxis.group_text_font_size = '11pt' p.xaxis.subgroup_text_font_size = '0pt' p.xaxis.axis_label_text_font_size = '11pt' p.xaxis.major_label_text_font_size = '10pt' p.xaxis.axis_label_text_font_style = 'normal' p.xaxis.major_label_text_color = 'black' p.xaxis.axis_label_text_color = 'black' p.yaxis.axis_label_text_font_size = '11pt' p.yaxis.axis_label_text_font_style = 'normal' p.yaxis.axis_label_text_color = 'black' p.yaxis.major_label_text_font_size = '10pt' p.xaxis.group_text_color = 'dimgrey' return p def plot_ecdf_hits(prot_dict, n, top_per_lst): #use least stringent threshold to plot all hits top_per = max(top_per_lst)/100 prefix = 'output/dms/' esm1v_outpath = 'dms_esm1vScored/' + f"dms_{prot_dict['protein']}" + '_maskMarginals.csv' if_results_df = pd.read_csv( prefix + f"if_results/dms_{prot_dict['protein']}_if.csv") dms_df = pd.read_csv(prefix + esm1v_outpath) unprofiled_df = dms_df[dms_df[prot_dict['fitness_col']].isna()] unprofiled_variants = unprofiled_df['variant'].to_list() dms_df = dms_df.dropna(subset = [prot_dict['fitness_col']]) dms_df['percentile'] = dms_df[prot_dict['fitness_col']].rank(pct=True, method='average') dms_df['zscore'] = (dms_df[prot_dict['fitness_col']] - dms_df[prot_dict['fitness_col']].mean())/dms_df[prot_dict['fitness_col']].std() if_results_df = if_results_df[~if_results_df['seqid'].isin(unprofiled_variants)] assert (len(if_results_df) == len(dms_df)) threshold = dms_df[prot_dict['fitness_col']].quantile(1-top_per) highFit_subset_variants = dms_df[dms_df[prot_dict['fitness_col']] >= threshold]['variant'] n_highFit = len(highFit_subset_variants) dms_df['esm1v'] = dms_df.loc[:, dms_df.columns.str.startswith('esm1v')].mean(axis=1) top_esm1v_df = dms_df.sort_values(by = 'esm1v', ascending = False)[:n] top_esm1v_variants = top_esm1v_df['variant'].to_list() top_if_df = if_results_df.sort_values(by = 'log_likelihood', ascending = False)[:n] top_if_variants = top_if_df['seqid'].to_list() esm1v_hits_df = top_esm1v_df[top_esm1v_df['variant'].isin(highFit_subset_variants)] if_hits_df = top_if_df[top_if_df['seqid'].isin(highFit_subset_variants)] if_hits_df = pd.merge(if_hits_df, dms_df[['variant', 'zscore', 'percentile', prot_dict['fitness_col']]], left_on='seqid', right_on='variant', how='left') p_1v = iqplot.ecdf( data=dms_df, q=prot_dict['fitness_col'], style="staircase", palette = 'lightgrey', line_kwargs= {'line_width':6}, x_axis_label = prot_dict['protein']+' Fitness Measure', #title = prot_dict['protein'] ) p_if = iqplot.ecdf( data=dms_df, q=prot_dict['fitness_col'], style="staircase", palette = 'lightgrey', line_kwargs= {'line_width':6}, x_axis_label = None, #prot_dict['protein']+'Fitness Measure', title = prot_dict['protein'] ) colors = ['#1f78b4', '#33a02c', '#ff7f00',] for i, top_per in enumerate(top_per_lst): top_per = top_per/100 hline = bokeh.models.Span(location=1-top_per, dimension='width', line_color=colors[i], line_dash='dashed', line_width=4, ) p_1v.add_layout(hline) p_if.add_layout(hline) p_1v.circle( source = esm1v_hits_df, x = prot_dict['fitness_col'], y = 'percentile', color = 'dimgrey', size = 11, line_color = 'lightgrey', alpha = 0.8, ) p_if.circle( source = if_hits_df, x = prot_dict['fitness_col'], y = 'percentile', color = '#43a2ca', size = 11, line_color = 'white', alpha = 0.8, ) p_if.title.text_font_style = 'normal' p_if.title.text_color = 'black' p_if.title.align = 'center' p_if.title.text_font_size = '16pt' p_1v.xaxis.axis_label_text_font_size = '14pt' for p in [p_if, p_1v]: p.xaxis.axis_label_text_font_style = 'normal' p.yaxis.axis_label_text_font_style = 'normal' p.xaxis.axis_label_text_font_size = '14pt' p.yaxis.axis_label_text_font_size = '14pt' p.xaxis.axis_label_text_color = 'black' p.yaxis.axis_label_text_color = 'black' p.xaxis.major_label_text_font_size = '14pt' p.yaxis.major_label_text_font_size = '14pt' p.xaxis.major_label_text_color = 'black' p.yaxis.major_label_text_color = 'black' p.output_backend = "svg" return [p_if, p_1v] if __name__ == '__main__': for s in studies: score_if(s) shutil.copytree("data/dms/dms_esm1vScored", "output/dms/dms_esm1vScored", dirs_exist_ok=True) name, n_pop, n_popHits, n_if_hits, n_esm1v_hits, if_hit_enrich, esm1v_hit_enrich, if_hit_rate, esm1v_hit_rate, p_threshold = ([] for _ in range(10)) data_lists = [name, n_pop, n_popHits, n_if_hits, n_esm1v_hits, if_hit_enrich, esm1v_hit_enrich, if_hit_rate, esm1v_hit_rate, p_threshold] for s in studies: for p in [5, 10, 20 ]: outputs = get_stats(s, 10, p) for i in range(len(outputs)): data_lists[i].append(outputs[i]) data_lists[len(data_lists)-1].append(p) col_names = ['Protein', 'Total Library Size', 'Library Hits (Variants with Fitness >95th)', 'Inverse Folding Hits (in Top 10)', 'ESM1v Hits (in Top 10)', 'Inverse Folding Hit Enrichment', 'ESM1v Hit Enrichment','Inverse Folding Hit Rate', 'ESM1v Hit Rate','Percentile Threshold' ] results_df = pd.DataFrame({col_names[i]: data_lists[i] for i in range(len(data_lists))}) # Melt the dataframe melted_df = results_df.melt(id_vars=['Protein', 'Percentile Threshold'], value_vars=[ 'ESM1v Hit Rate', 'Inverse Folding Hit Rate',], var_name='Method', value_name='Hit Rate') melted_df['Method'] = melted_df['Method'].replace({ 'Inverse Folding Hit Rate': 'Inverse Folding', 'ESM1v Hit Rate': 'Language Model'}) melted_df['alpha'] = (1/melted_df['Percentile Threshold'].to_numpy())*3 + 0.4 melted_df['Percentile Threshold'] = melted_df['Percentile Threshold'].astype(str) melted_df['legend_label'] = ['Structure-Informed Language Model' if x=='Inverse Folding' else 'Language Model' for x in melted_df['Method']] melted_df['cats'] = melted_df.apply(lambda x: (x["Protein"], x["Method"], x['Percentile Threshold'],), axis = 1) factors = list(melted_df.cats) p = plot_bars(melted_df) comparePrecision_fname = 'output/dms/precision_comparison.html' bokeh.plotting.output_file(comparePrecision_fname) bokeh.io.show(p) ecdf_plots = [] for s in studies: ecdf_plots.extend(plot_ecdf_hits(s, 10, [20,10,5])) mid = int(len(ecdf_plots)/2) ecdf_fname = 'output/dms/fitness_ecdfs.html' bokeh.plotting.output_file(ecdf_fname) bokeh.io.show(bokeh.layouts.gridplot([ecdf_plots[:mid:2], ecdf_plots[1:mid:2], ecdf_plots[mid::2], ecdf_plots[mid+1::2]]) )