# This module contains utility functions for zero-shot predictions using various protein language models import argparse from Bio import SeqIO import numpy as np import pandas as pd import scipy.stats as ss import torch from tqdm import tqdm from model.utils.other_utils import read_msa, greedy_select, msa_splicer, AAs def zero_shot_esm_dms(wt_seq, model_locations = ['esm1v_t33_650M_UR90S_1', 'esm1v_t33_650M_UR90S_2', 'esm1v_t33_650M_UR90S_3', 'esm1v_t33_650M_UR90S_4', 'esm1v_t33_650M_UR90S_5', 'esm2_t36_3B_UR50D'], scoring_strategy='wt-marginals', num_msa_seqs=400, **kwargs): """ Perform deep mutational scanning using ESM model. Args: wt_seq (str): Wild-type protein sequence scoring_strategy (str): 'wt-marginals' or 'masked-marginals' **kwargs: Additional arguments Returns: pandas.DataFrame: DataFrame containing mutation scores and statistics """ from esm import pretrained, MSATransformer # create list of all possible single point mutations in the wildtype sequence amino_acids = AAs[:-1] mutations = [] for i, residue in enumerate(wt_seq): for aa in amino_acids: if wt_seq[i] == aa: continue mutations.append(wt_seq[i] + str(i + 1) + aa) # Compute token probs for each model. model_probs = [] if torch.backends.mps.is_available(): device = "mps" elif torch.cuda.is_available(): device = "cuda:0" else: device = "cpu" for model_location in model_locations: model, alphabet = pretrained.load_model_and_alphabet(model_location) model.eval() model = model.to(device) batch_converter = alphabet.get_batch_converter() if isinstance(model, MSATransformer): assert kwargs['msa_file'] is not None, 'No MSA file provided.' msa = read_msa(kwargs['msa_file']) # Prep the MSA, making the appropriate mutations inputs = greedy_select(msa, num_seqs=num_msa_seqs) # can change this to pass more/fewer sequences #This splices the MSA to exclude gaps in the first sequence, due to MSATransformer context window #size limit of 1024. If your MSA width is less than 1024, then you don't need to do this data = [msa_splicer(inputs)] # Run the model, retrieve logits _, __, batch_tokens = batch_converter(data) all_token_probs = [] for i in tqdm(range(batch_tokens.size(2))): batch_tokens_masked = batch_tokens.clone() batch_tokens_masked[0, 0, i] = alphabet.mask_idx # mask out first sequence with torch.no_grad(): token_probs = torch.log_softmax( model(batch_tokens_masked.to(device))["logits"], dim=-1 ) all_token_probs.append(token_probs[:, 0, i]) # vocab size token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0) else: data = [ ('protein1', wt_seq), ] batch_labels, batch_strs, batch_tokens = batch_converter(data) if scoring_strategy == 'wt-marginals': with torch.no_grad(): token_probs = torch.log_softmax(model(batch_tokens.to(device))['logits'], dim=-1) elif scoring_strategy == 'masked-marginals': all_token_probs = [] for i in tqdm(range(batch_tokens.size(1))): batch_tokens_masked = batch_tokens.clone() batch_tokens_masked[0, i] = alphabet.mask_idx with torch.no_grad(): token_probs = torch.log_softmax( model(batch_tokens_masked.to(device))['logits'], dim=-1 ) all_token_probs.append(token_probs[:, i]) # vocab size token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0) else: raise ValueError(f'Invalid scoring strategy {scoring_strategy}') model_probs.append(token_probs.cpu().numpy()[0]) X = [] for model_prob in model_probs: X_sub = [] for mutation in mutations: wt, idx, mt = mutation[0], int(mutation[1:-1])-1, mutation[-1] assert wt_seq[idx] == wt, 'Wild-type residue does not match input sequence.' wt_encoded, mt_encoded = alphabet.tok_to_idx[wt], alphabet.tok_to_idx[mt] score = model_prob[idx + 1, mt_encoded] - model_prob[idx + 1, wt_encoded] if not np.isfinite(score): score = 0. X_sub.append(score) X.append(X_sub) # set up dataframe data = {'mutations': mutations} for i in range(len(X)): data[f'model_{i+1}_logratio'] = X[i] df = pd.DataFrame(data) # Calculate average log ratio across all models logratio_cols = [f'model_{i+1}_logratio' for i in range(len(X))] df['average_model_logratio'] = df[logratio_cols].mean(axis=1) # Calculate pass/fail for each model for i in range(len(X)): df[f'model_{i+1}_pass'] = df[f'model_{i+1}_logratio'].apply(lambda x: 1 if x > 0 else 0) # Sum up total passes pass_cols = [f'model_{i+1}_pass' for i in range(len(X))] df['total_model_pass'] = df[pass_cols].sum(axis=1) df.sort_values(by='average_model_logratio', ascending=False, inplace=True) df.sort_values(by='total_model_pass', ascending=False, inplace=True) df_ls = [] # sort dataframe by total_model_pass and then by average_model_logratio total_model_pass_list = list(set(df['total_model_pass'].values)) total_model_pass_list = total_model_pass_list[::-1] for model_pass_value in total_model_pass_list: subset = df[df['total_model_pass'] == model_pass_value].copy() subset.sort_values(by='average_model_logratio', ascending=False, inplace=True) df_ls.append(subset) df_sorted = pd.concat(df_ls) return df_sorted def zero_shot_esm_if_dms(wt_seq, pdb_file, chain_id = 'A', scoring_strategy='wt-marginals', **kwargs): """ Perform deep mutational scanning using ESM-IF (Inverse Folding) model. Args: wt_seq (str): Wild-type protein sequence pdb_file (str): Path to PDB file chain_id (str): Chain ID in the PDB file scoring_strategy (str): Currently not used, kept for consistency **kwargs: Additional arguments Returns: pandas.DataFrame: DataFrame containing mutation scores """ import torch_geometric import torch_sparse from torch_geometric.nn import MessagePassing import esm from esm import pretrained from esm.inverse_folding.util import CoordBatchConverter amino_acids = AAs[:-1] mutations = [] for i, residue in enumerate(wt_seq): for aa in amino_acids: if wt_seq[i] == aa: continue mutations.append(wt_seq[i] + str(i + 1) + aa) model_locations = ['esm_if1_gvp4_t16_142M_UR50'] model, alphabet = pretrained.load_model_and_alphabet(model_locations[0]) model = model.eval() structure = esm.inverse_folding.util.load_structure(pdb_file, chain_id) coords, native_seq = esm.inverse_folding.util.extract_coords_from_structure(structure) if native_seq == wt_seq: print(f"Native sequence from structure matches input sequence ({len(native_seq)} residues)") else: print(f"Warning: Native sequence from structure ({len(native_seq)} residues) does not match input sequence ({len(wt_seq)} residues)") device = next(model.parameters()).device batch_converter = CoordBatchConverter(alphabet) batch = [(coords, None, wt_seq)] coords, confidence, strs, tokens, padding_mask = batch_converter( batch, device=device) prev_output_tokens = tokens[:, :-1].to(device) target = tokens[:, 1:] logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens) # Average model scores and find scores for the mutations-of-interest. scores = logits.detach().numpy()[0] mutation_score = {} for pos in range(len(wt_seq)): wt = wt_seq[pos] for mt in alphabet.all_toks: mutation = f'{wt}{pos + 1}{mt}' mutation_score[mutation] = scores[alphabet.tok_to_idx[mt], pos] X = [] for mutation in mutations: wt = mutation[0]+mutation[1:-1]+mutation[0] score = mutation_score[mutation] - mutation_score[wt] if not np.isfinite(score): score = 0. X.append(score) df = pd.DataFrame({'mutations': mutations, 'logratio': X}) return df def zero_shot_esm( mutations, model_locations, sequence, scoring_strategy='wt-marginals', **kwargs ): """ Perform zero-shot prediction using ESM (Evolutionary Scale Modeling) model. Args: mutations (list): List of mutation sets model_locations (list): List of ESM model file paths sequence (str): Original protein sequence scoring_strategy (str): 'wt-marginals' or 'masked-marginals' **kwargs: Additional arguments (e.g., device) Returns: numpy.ndarray: Array of mutation scores """ from esm import pretrained # Compute token probs for each model. model_probs = [] for model_location in model_locations: model, alphabet = pretrained.load_model_and_alphabet(model_location) model.eval() model = model.to(kwargs['device']) batch_converter = alphabet.get_batch_converter() data = [ ('protein1', sequence), ] batch_labels, batch_strs, batch_tokens = batch_converter(data) if scoring_strategy == 'wt-marginals': with torch.no_grad(): token_probs = torch.log_softmax(model(batch_tokens.to(kwargs['device']))['logits'], dim=-1) elif scoring_strategy == 'masked-marginals': all_token_probs = [] for i in tqdm(range(batch_tokens.size(1))): batch_tokens_masked = batch_tokens.clone() batch_tokens_masked[0, i] = alphabet.mask_idx with torch.no_grad(): token_probs = torch.log_softmax( model(batch_tokens_masked.to(kwargs['device']))['logits'], dim=-1 ) all_token_probs.append(token_probs[:, i]) # vocab size token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0) else: raise ValueError(f'Invalid scoring strategy {scoring_strategy}') model_probs.append(token_probs.cpu().numpy()[0]) # Sum model scores and find scores for the mutations-of-interest. scores = np.sum(model_probs, axis=0) mutation_score = {} for pos in range(len(sequence)): wt = sequence[pos] for mt in alphabet.all_toks: mutation = f'{wt}{pos + 1}{mt}' mutation_score[mutation] = scores[pos + 1, alphabet.tok_to_idx[mt]] X = [] for mutation_set in mutations: score = np.mean([ mutation_score[mutation] for mutation in mutation_set ]) if not np.isfinite(score): score = 0. X.append(score) return np.array(X) def zero_shot_msa( mutations, sequence, **kwargs, ): """ Perform zero-shot prediction using MSA Transformer model. Args: mutations (list): List of mutation sets sequence (str): Original protein sequence **kwargs: Additional arguments (must include 'msa_file') Returns: numpy.ndarray: Array of mutation scores """ import esm import torch torch.set_grad_enabled(False) # Check to see if there is an MSA file in **kwargs. assert kwargs['msa_file'] is not None, 'No MSA file provided.' msa = read_msa(kwargs['msa_file']) # Instantiate the model msa_transformer, msa_transformer_alphabet = esm.pretrained.esm_msa1b_t12_100M_UR50S() msa_transformer = msa_transformer.eval() msa_transformer_batch_converter = msa_transformer_alphabet.get_batch_converter() # Prep the MSA, making the appropriate mutations inputs = greedy_select(msa, num_seqs=128) # can change this to pass more/fewer sequences #This splices the MSA to exclude gaps in the first sequence, due to MSATransformer context window #size limit of 1024. If your MSA width is less than 1024, then you don't need to do this inputs = [msa_splicer(inputs)] # Run the model, retrieve logits _, __, msa_transformer_batch_tokens = msa_transformer_batch_converter(inputs) msa_transformer_batch_tokens = msa_transformer_batch_tokens.to(next(msa_transformer.parameters()).device) predictions = msa_transformer.forward(msa_transformer_batch_tokens, repr_layers=[12]) logits = predictions['logits'][0][0] token_probs = torch.softmax(logits, dim=-1) print(token_probs.shape) print('pulling out mutations') # Pull specific logits out for the mutations-of-interest. mutation_score = {} for pos in range(len(sequence)): wt = sequence[pos] for mt in msa_transformer_alphabet.all_toks: mutation = f'{wt}{pos + 1}{mt}' mutation_score[mutation] = token_probs[pos + 1, msa_transformer_alphabet.tok_to_idx[mt]] X = [] for mutation_set in mutations: score = np.mean([ mutation_score[mutation] for mutation in mutation_set ]) if not np.isfinite(score): score = 0. X.append(score) return np.array(X) def zero_shot_esm_if( mutations, model_locations, sequence, pdb_file, chain_id, **kwargs ): """ Perform zero-shot prediction using ESM-IF (Inverse Folding) model. Args: mutations (list): List of mutation sets model_locations (list): List of ESM-IF model file paths sequence (str): Original protein sequence pdb_file (str): Path to PDB file chain_id (str): Chain ID in the PDB file **kwargs: Additional arguments Returns: numpy.ndarray: Array of mutation scores """ # Check that imports are correctly installed import torch_geometric import torch_sparse from torch_geometric.nn import MessagePassing import esm from esm import pretrained from esm.inverse_folding.util import CoordBatchConverter # If one of the above fails, run the following in your conda environment # import torch # def format_pytorch_version(version): # return version.split('+')[0] # TORCH_version = torch.__version__ # TORCH = format_pytorch_version(TORCH_version) # def format_cuda_version(version): # return 'cu' + version.replace('.', '') # CUDA_version = torch.version.cuda # CUDA = format_cuda_version(CUDA_version) # !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html # !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html # !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html # !pip install -q torch-spline-conv -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html # !pip install -q torch-geometric # # Install esm # !pip install -q git+https://github.com/facebookresearch/esm.git # # Install biotite # !pip install -q biotite # Compute token probs for each model. model, alphabet = pretrained.load_model_and_alphabet(model_locations[0]) model = model.eval() structure = esm.inverse_folding.util.load_structure(pdb_file, chain_id) coords, native_seq = esm.inverse_folding.util.extract_coords_from_structure(structure) device = next(model.parameters()).device batch_converter = CoordBatchConverter(alphabet) batch = [(coords, None, sequence)] coords, confidence, strs, tokens, padding_mask = batch_converter( batch, device=device) prev_output_tokens = tokens[:, :-1].to(device) target = tokens[:, 1:] logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens) # Average model scores and find scores for the mutations-of-interest. scores = logits.detach().numpy()[0] mutation_score = {} for pos in range(len(sequence)): wt = sequence[pos] for mt in alphabet.all_toks: mutation = f'{wt}{pos + 1}{mt}' mutation_score[mutation] = scores[alphabet.tok_to_idx[mt], pos] # logits are vocab x length (no padding) X = [] for mutation_set in mutations: score = np.mean([ mutation_score[mutation] for mutation in mutation_set ]) if not np.isfinite(score): score = 0. X.append(score) return np.array(X)