| from dataclasses import dataclass |
| import numpy as np |
| import torch |
|
|
| from extra_utils import paired_msa_numbering, unpaired_msa_numbering, create_alignment |
|
|
|
|
| class AbAlignment: |
|
|
| def __init__(self, device = 'cpu', ncpu = 1): |
| |
| self.device = device |
| self.ncpu = ncpu |
| |
| def number_sequences(self, seqs, chain = 'H', fragmented = False): |
| if chain == 'HL': |
| numbered_seqs, seqs, number_alignment = paired_msa_numbering(seqs, fragmented = fragmented, n_jobs = self.ncpu) |
| else: |
| assert chain == 'HL', 'Currently "Align==True" only works for paired sequences. \nPlease use paired sequences or Align=False.' |
| numbered_seqs, seqs, number_alignment = unpaired_msa_numbering( |
| seqs, chain = chain, fragmented = fragmented, n_jobs = self.ncpu |
| ) |
| |
| return numbered_seqs, seqs, number_alignment |
| |
| def align_encodings(self, encodings, numbered_seqs, seqs, number_alignment): |
| |
| aligned_list = [ |
| create_alignment( |
| res_embed, numbered_seq, seq, number_alignment |
| ) for res_embed, numbered_seq, seq in zip(encodings, numbered_seqs, seqs) |
| ] |
| aligned_encodings = np.concatenate([aligned_list], axis=0) |
| return aligned_encodings |
| |
| |
| def reformat_subsets( |
| self, |
| subset_list, |
| mode = 'seqcoding', |
| align = False, |
| numbered_seqs = None, |
| seqs = None, |
| number_alignment = None, |
| ): |
| |
| if mode in [ |
| 'seqcoding', |
| 'restore', |
| 'pseudo_log_likelihood', |
| 'confidence' |
| ]: |
| return np.concatenate(subset_list) |
| elif align: |
| subset_list = [ |
| self.align_encodings( |
| subset, |
| numbered_seqs[num*len(subset):(num+1)*len(subset)], |
| seqs[num*len(subset):(num+1)*len(subset)], |
| number_alignment |
| ) for num, subset in enumerate(subset_list) |
| ] |
| |
| subset = np.concatenate(subset_list) |
| |
| return aligned_results( |
| aligned_seqs = [''.join(alist) for alist in subset[:,:,-1]], |
| aligned_embeds = subset[:,:,:-1].astype(float), |
| number_alignment=number_alignment.apply(lambda x: '{}{}'.format(*x[0]), axis=1).values |
| ) |
| |
| elif not align: |
| return sum(subset_list, []) |
| else: |
| return np.concatenate(subset_list) |
| |
|
|
| @dataclass |
| class aligned_results(): |
| """ |
| Dataclass used to store output. |
| """ |
| |
| aligned_seqs: None |
| aligned_embeds: None |
| number_alignment: None |
