import numpy as np import pandas as pd from collections import defaultdict import random import os import torch from Bio.Align.Applications import ClustalOmegaCommandline def filter_msa(msa_data, num_sequences_kept=3): """ Helper function to filter an input MSA msa_data (obtained via process_msa_data) and keep only num_sequences_kept aligned sequences. If the MSA already has fewer sequences than num_sequences_kept, we keep the MSA as is. If filtering, we always keep the first sequence of the MSA (ie. the wild type) by default. Sampling is done without replacement. """ if len(list(msa_data.keys())) <= num_sequences_kept: return msa_data filtered_msa = {} wt_name = next(iter(msa_data)) filtered_msa[wt_name] = msa_data[wt_name] del msa_data[wt_name] sequence_names = list(msa_data.keys()) sequence_names_sampled = random.sample(sequence_names,k=num_sequences_kept-1) for seq in sequence_names_sampled: filtered_msa[seq] = msa_data[seq] return filtered_msa def process_msa_data(MSA_data_file): """ Helper function that takes as input a path to a MSA file (expects a2m format) and returns a dict mapping sequence ID to the corresponding AA sequence. """ msa_data = defaultdict(str) sequence_name = "" with open(MSA_data_file, "r") as msa_file: for i, line in enumerate(msa_file): line = line.rstrip() if line.startswith(">"): sequence_name = line else: msa_data[sequence_name] += line.upper() return msa_data def get_one_hot_sequences_dict(msa_data,MSA_start,MSA_end,vocab): vocab_size = len(vocab.keys()) num_sequences_msa = len(msa_data.keys()) one_hots = np.zeros((num_sequences_msa,MSA_end-MSA_start,vocab_size)) for i,seq_name in enumerate(msa_data.keys()): sequence = msa_data[seq_name] for j,letter in enumerate(sequence): if letter in vocab: k = vocab[letter] one_hots[i,j,k] = 1.0 return one_hots def one_hot(sequence_string,vocab): one_hots = np.zeros((len(sequence_string),len(vocab.keys()))) for j,letter in enumerate(sequence_string): if letter in vocab: k = vocab[letter] one_hots[j,k] = 1.0 return one_hots.flatten() def get_msa_prior(MSA_data_file, MSA_weight_file_name, MSA_start, MSA_end, len_target_seq, vocab, retrieval_aggregation_mode="aggregate_substitution", filter_MSA=True, verbose=False): """ Function to enable retrieval inference mode, via computation of (weighted) pseudocounts of AAs at each position of the retrieved MSA. MSA_data_file: (string) path to MSA file (expects a2m format). MSA_weight_file_name: (string) path to sequence weights in MSA. MSA_start: (int) Sequence position that the MSA starts at (1-indexing). MSA_end: (int) Sequence position that the MSA ends at (1-indexing). len_target_seq: (int) Full length of sequence to be scored. vocab: (dict) Vocabulary of the tokenizer. retrieval_aggregation_mode: (string) Mode for retrieval inference (aggregate_substitution Vs aggregate_indel). If None, places a uniform prior over each token. filter_MSA: (bool) Whether to filter out sequences with very low hamming similarity (< 0.2) to the reference sequence in the MSA (first sequence). verbose: (bool) Whether to print to the console processing details along the way. """ msa_data = process_msa_data(MSA_data_file) vocab_size = len(vocab.keys()) if verbose: print("Target seq len is {}, MSA length is {}, start position is {}, end position is {} and vocab size is {}".format(len_target_seq,MSA_end-MSA_start,MSA_start,MSA_end,vocab_size)) if filter_MSA: if verbose: print("Num sequences in MSA pre filtering: {}".format(len(msa_data.keys()))) list_sequence_names = list(msa_data.keys()) focus_sequence_name = list(msa_data.keys())[0] ref_sequence_hot = one_hot(msa_data[focus_sequence_name],vocab) for sequence_name in list_sequence_names: seq_hot = one_hot(msa_data[sequence_name],vocab) hamming_similarity_seq_ref = np.dot(ref_sequence_hot,seq_hot) / np.dot(ref_sequence_hot,ref_sequence_hot) if hamming_similarity_seq_ref < 0.2: del msa_data[sequence_name] if verbose: print("Num sequences in MSA post filtering: {}".format(len(msa_data.keys()))) if MSA_weight_file_name is not None: if verbose: print("Using weights in {} for sequences in MSA.".format(MSA_weight_file_name)) assert os.path.exists(MSA_weight_file_name), "Weights file not located on disk." MSA_EVE = MSA_processing( MSA_location=MSA_data_file, use_weights=True, weights_location=MSA_weight_file_name ) #We scan through all sequences to see if we have a weight for them as per EVE pre-processing. We drop them otherwise. dropped_sequences=0 list_sequence_names = list(msa_data.keys()) MSA_weight=[] for sequence_name in list_sequence_names: if sequence_name not in MSA_EVE.seq_name_to_sequence: dropped_sequences +=1 del msa_data[sequence_name] else: MSA_weight.append(MSA_EVE.seq_name_to_weight[sequence_name]) if verbose: print("Dropped {} sequences from MSA due to absent sequence weights".format(dropped_sequences)) else: MSA_weight = [1] * len(list(msa_data.keys())) if retrieval_aggregation_mode=="aggregate_substitution" or retrieval_aggregation_mode=="aggregate_indel": one_hots = get_one_hot_sequences_dict(msa_data,MSA_start,MSA_end,vocab) MSA_weight = np.expand_dims(np.array(MSA_weight),axis=(1,2)) base_rate = 1e-5 base_rates = np.ones_like(one_hots) * base_rate weighted_one_hots = (one_hots + base_rates) * MSA_weight MSA_weight_norm_counts = weighted_one_hots.sum(axis=-1).sum(axis=0) MSA_weight_norm_counts = np.tile(MSA_weight_norm_counts.reshape(-1,1), (1,vocab_size)) one_hots_avg = weighted_one_hots.sum(axis=0) / MSA_weight_norm_counts msa_prior = np.zeros((len_target_seq,vocab_size)) msa_prior[MSA_start:MSA_end,:]=one_hots_avg else: msa_prior = np.ones((len_target_seq,vocab_size)) / vocab_size if verbose: for idx, position in enumerate(msa_prior): if len(position)!=25: print("Size error") if not round(position.sum(),2)==1.0: print("Position at index {} does not add up to 1: {}".format(idx, position.sum())) return msa_prior def update_retrieved_MSA_log_prior_indel(model, MSA_log_prior, MSA_start, MSA_end, mutated_sequence): """ Function to process MSA when scoring indels. To identify positions to add / remove in the retrieved MSA, we append and align the sequence to be scored to the original MSA for that protein family with Clustal Omega. If the original MSA is relatively deep (over 100k sequences), we sample (by default) 100k rows at random from that MSA to speed computations. MSA sampling is performed only once (for the first sequence to be scored). Subsequent scoring use the same MSA sample. """ if not os.path.isdir(model.MSA_folder + os.sep + "Sampled"): os.mkdir(model.MSA_folder + os.sep + "Sampled") sampled_MSA_location = model.MSA_folder + os.sep + "Sampled" + os.sep + "Sampled_" + model.MSA_filename.split(os.sep)[-1] if not os.path.exists(sampled_MSA_location): msa_data = process_msa_data(model.MSA_filename) msa_data_sampled = filter_msa(msa_data, num_sequences_kept=100000) #If MSA has less than 100k sequences, the sample is identical to original MSA with open(sampled_MSA_location, 'w') as sampled_write_location: for index, key in enumerate(msa_data_sampled): key_name = ">REFERENCE_SEQUENCE" if index==0 else key msa_data_sampled[key] = msa_data_sampled[key].upper() msa_data_sampled[key] = msa_data_sampled[key].replace(".","-") sampled_write_location.write(key_name+"\n"+"\n".join([msa_data_sampled[key][i:i+80] for i in range(0, len(msa_data_sampled[key]), 80)])+"\n") seq_to_align_location = model.MSA_folder + os.sep + "Sampled" + os.sep + "Seq_to_align_" + model.MSA_filename.split(os.sep)[-1] sequence_text_split = [mutated_sequence[i:i+80] for i in range(0, len(mutated_sequence), 80)] sequence_text_split_split_join = "\n".join([">SEQ_TO_SCORE"]+sequence_text_split) os.system("echo '"+sequence_text_split_split_join+"' > "+seq_to_align_location) expanded_MSA_location = model.MSA_folder + os.sep + "Sampled" + os.sep + "Expanded_" + model.MSA_filename.split(os.sep)[-1] clustalw_cline = ClustalOmegaCommandline(cmd=model.config.clustal_omega_location, profile1=sampled_MSA_location, profile2=seq_to_align_location, outfile=expanded_MSA_location, force=True) stdout, stderr = clustalw_cline() msa_data = process_msa_data(expanded_MSA_location) aligned_seqA, aligned_seqB = msa_data[">SEQ_TO_SCORE"], msa_data[">REFERENCE_SEQUENCE"] try: keep_column=[] for column_index_pairwise_alignment in range(len(aligned_seqA)): if aligned_seqA[column_index_pairwise_alignment]=="-" and aligned_seqB[column_index_pairwise_alignment]=="-": continue elif aligned_seqA[column_index_pairwise_alignment]=="-": keep_column.append(False) elif aligned_seqB[column_index_pairwise_alignment]=="-": MSA_log_prior=torch.cat((MSA_log_prior[:column_index_pairwise_alignment], torch.zeros(MSA_log_prior.shape[1]).view(1,-1).cuda(), MSA_log_prior[column_index_pairwise_alignment:]),dim=0) keep_column.append(True) #keep the zero column we just added else: keep_column.append(True) MSA_log_prior = MSA_log_prior[keep_column] MSA_end = MSA_start + len(MSA_log_prior) except: print("Error when processing the following alignment: {}".format(expanded_MSA_location)) return MSA_log_prior, MSA_start, MSA_end class MSA_processing: def __init__(self, MSA_location="", theta=0.2, use_weights=True, weights_location="./data/weights", preprocess_MSA=True, threshold_sequence_frac_gaps=0.5, threshold_focus_cols_frac_gaps=0.3, remove_sequences_with_indeterminate_AA_in_focus_cols=True ): """ This MSA_processing class is directly borrowed from the EVE codebase: https://github.com/OATML-Markslab/EVE Parameters: - msa_location: (path) Location of the MSA data. Constraints on input MSA format: - focus_sequence is the first one in the MSA data - first line is structured as follows: ">focus_seq_name/start_pos-end_pos" (e.g., >SPIKE_SARS2/310-550) - corespondding sequence data located on following line(s) - then all other sequences follow with ">name" on first line, corresponding data on subsequent lines - theta: (float) Sequence weighting hyperparameter. Generally: Prokaryotic and eukaryotic families = 0.2; Viruses = 0.01 - use_weights: (bool) If False, sets all sequence weights to 1. If True, checks weights_location -- if non empty uses that; otherwise compute weights from scratch and store them at weights_location - weights_location: (path) Location to load from/save to the sequence weights - preprocess_MSA: (bool) performs pre-processing of MSA to remove short fragments and positions that are not well covered. - threshold_sequence_frac_gaps: (float, between 0 and 1) Threshold value to define fragments - sequences with a fraction of gap characters above threshold_sequence_frac_gaps are removed - default is set to 0.5 (i.e., fragments with 50% or more gaps are removed) - threshold_focus_cols_frac_gaps: (float, between 0 and 1) Threshold value to define focus columns - positions with a fraction of gap characters above threshold_focus_cols_pct_gaps will be set to lower case (and not included in the focus_cols) - default is set to 0.3 (i.e., focus positions are the ones with 30% of gaps or less, i.e., 70% or more residue occupancy) - remove_sequences_with_indeterminate_AA_in_focus_cols: (bool) Remove all sequences that have indeterminate AA (e.g., B, J, X, Z) at focus positions of the wild type """ np.random.seed(2021) self.MSA_location = MSA_location self.weights_location = weights_location self.theta = theta self.alphabet = "ACDEFGHIKLMNPQRSTVWY" self.use_weights = use_weights self.preprocess_MSA = preprocess_MSA self.threshold_sequence_frac_gaps = threshold_sequence_frac_gaps self.threshold_focus_cols_frac_gaps = threshold_focus_cols_frac_gaps self.remove_sequences_with_indeterminate_AA_in_focus_cols = remove_sequences_with_indeterminate_AA_in_focus_cols self.gen_alignment() def gen_alignment(self, verbose=False): """ Read training alignment and store basics in class instance """ self.aa_dict = {} for i,aa in enumerate(self.alphabet): self.aa_dict[aa] = i self.seq_name_to_sequence = defaultdict(str) name = "" with open(self.MSA_location, "r") as msa_data: for i, line in enumerate(msa_data): line = line.rstrip() if line.startswith(">"): name = line if i==0: self.focus_seq_name = name else: self.seq_name_to_sequence[name] += line ## MSA pre-processing to remove inadequate columns and sequences if self.preprocess_MSA: msa_df = pd.DataFrame.from_dict(self.seq_name_to_sequence, orient='index', columns=['sequence']) # Data clean up msa_df.sequence = msa_df.sequence.apply(lambda x: x.replace(".","-")).apply(lambda x: ''.join([aa.upper() for aa in x])) # Remove columns that would be gaps in the wild type non_gap_wt_cols = [aa!='-' for aa in msa_df.sequence[self.focus_seq_name]] msa_df['sequence'] = msa_df['sequence'].apply(lambda x: ''.join([aa for aa,non_gap_ind in zip(x, non_gap_wt_cols) if non_gap_ind])) assert 0.0 <= self.threshold_sequence_frac_gaps <= 1.0,"Invalid fragment filtering parameter" assert 0.0 <= self.threshold_focus_cols_frac_gaps <= 1.0,"Invalid focus position filtering parameter" msa_array = np.array([list(seq) for seq in msa_df.sequence]) gaps_array = np.array(list(map(lambda seq: [aa=='-' for aa in seq], msa_array))) # Identify fragments with too many gaps seq_gaps_frac = gaps_array.mean(axis=1) seq_below_threshold = seq_gaps_frac <= self.threshold_sequence_frac_gaps if verbose: print("Proportion of sequences dropped due to fraction of gaps: "+str(round(float(1 - seq_below_threshold.sum()/seq_below_threshold.shape)*100,2))+"%") # Identify focus columns columns_gaps_frac = gaps_array[seq_below_threshold].mean(axis=0) index_cols_below_threshold = columns_gaps_frac <= self.threshold_focus_cols_frac_gaps if verbose: print("Proportion of non-focus columns removed: "+str(round(float(1 - index_cols_below_threshold.sum()/index_cols_below_threshold.shape)*100,2))+"%") # Lower case non focus cols and filter fragment sequences msa_df['sequence'] = msa_df['sequence'].apply(lambda x: ''.join([aa.upper() if upper_case_ind else aa.lower() for aa, upper_case_ind in zip(x, index_cols_below_threshold)])) msa_df = msa_df[seq_below_threshold] # Overwrite seq_name_to_sequence with clean version self.seq_name_to_sequence = defaultdict(str) for seq_idx in range(len(msa_df['sequence'])): self.seq_name_to_sequence[msa_df.index[seq_idx]] = msa_df.sequence[seq_idx] self.focus_seq = self.seq_name_to_sequence[self.focus_seq_name] self.focus_cols = [ix for ix, s in enumerate(self.focus_seq) if s == s.upper() and s!='-'] self.focus_seq_trimmed = [self.focus_seq[ix] for ix in self.focus_cols] self.seq_len = len(self.focus_cols) self.alphabet_size = len(self.alphabet) # Connect local sequence index with uniprot index (index shift inferred from 1st row of MSA) focus_loc = self.focus_seq_name.split("/")[-1] start,stop = focus_loc.split("-") self.focus_start_loc = int(start) self.focus_stop_loc = int(stop) self.uniprot_focus_col_to_wt_aa_dict \ = {idx_col+int(start):self.focus_seq[idx_col] for idx_col in self.focus_cols} self.uniprot_focus_col_to_focus_idx \ = {idx_col+int(start):idx_col for idx_col in self.focus_cols} # Move all letters to CAPS; keeps focus columns only self.raw_seq_name_to_sequence = self.seq_name_to_sequence.copy() for seq_name,sequence in self.seq_name_to_sequence.items(): sequence = sequence.replace(".","-") self.seq_name_to_sequence[seq_name] = [sequence[ix].upper() for ix in self.focus_cols] # Remove sequences that have indeterminate AA (e.g., B, J, X, Z) in the focus columns if self.remove_sequences_with_indeterminate_AA_in_focus_cols: alphabet_set = set(list(self.alphabet)) seq_names_to_remove = [] for seq_name,sequence in self.seq_name_to_sequence.items(): for letter in sequence: if letter not in alphabet_set and letter != "-": seq_names_to_remove.append(seq_name) continue seq_names_to_remove = list(set(seq_names_to_remove)) for seq_name in seq_names_to_remove: del self.seq_name_to_sequence[seq_name] # Encode the sequences self.one_hot_encoding = np.zeros((len(self.seq_name_to_sequence.keys()),len(self.focus_cols),len(self.alphabet))) if verbose: print("One-hot encoded sequences shape:" + str(self.one_hot_encoding.shape)) for i,seq_name in enumerate(self.seq_name_to_sequence.keys()): sequence = self.seq_name_to_sequence[seq_name] for j,letter in enumerate(sequence): if letter in self.aa_dict: k = self.aa_dict[letter] self.one_hot_encoding[i,j,k] = 1.0 if self.use_weights: try: self.weights = np.load(file=self.weights_location) if verbose: print("Loaded sequence weights from disk") except: if verbose: print ("Computing sequence weights") list_seq = self.one_hot_encoding list_seq = list_seq.reshape((list_seq.shape[0], list_seq.shape[1] * list_seq.shape[2])) def compute_weight(seq): number_non_empty_positions = np.dot(seq,seq) if number_non_empty_positions>0: denom = np.dot(list_seq,seq) / np.dot(seq,seq) denom = np.sum(denom > 1 - self.theta) return 1/denom else: return 0.0 #return 0 weight if sequence is fully empty self.weights = np.array(list(map(compute_weight,list_seq))) np.save(file=self.weights_location, arr=self.weights) else: # If not using weights, use an isotropic weight matrix if verbose: print("Not weighting sequence data") self.weights = np.ones(self.one_hot_encoding.shape[0]) self.Neff = np.sum(self.weights) self.num_sequences = self.one_hot_encoding.shape[0] self.seq_name_to_weight={} for i,seq_name in enumerate(self.seq_name_to_sequence.keys()): self.seq_name_to_weight[seq_name]=self.weights[i] if verbose: print ("Neff =",str(self.Neff)) print ("Data Shape =",self.one_hot_encoding.shape)