Spaces:
Runtime error
Runtime error
| import math | |
| import os | |
| import csv | |
| import random | |
| import torch | |
| from torch.utils import data | |
| import numpy as np | |
| from dateutil import parser | |
| import contigs | |
| from util import * | |
| from kinematics import * | |
| import pandas as pd | |
| import sys | |
| import torch.nn as nn | |
| from icecream import ic | |
| def write_pdb(filename, seq, atoms, Bfacts=None, prefix=None, chains=None): | |
| L = len(seq) | |
| ctr = 1 | |
| seq = seq.long() | |
| with open(filename, 'w+') as f: | |
| for i,s in enumerate(seq): | |
| if chains is None: | |
| chain='A' | |
| else: | |
| chain=chains[i] | |
| if (len(atoms.shape)==2): | |
| f.write ("%-6s%5s %4s %3s %s%4d %8.3f%8.3f%8.3f%6.2f%6.2f\n"%( | |
| "ATOM", ctr, " CA ", util.num2aa[s], | |
| chain, i+1, atoms[i,0], atoms[i,1], atoms[i,2], | |
| 1.0, Bfacts[i] ) ) | |
| ctr += 1 | |
| elif atoms.shape[1]==3: | |
| for j,atm_j in enumerate((" N "," CA "," C ")): | |
| f.write ("%-6s%5s %4s %3s %s%4d %8.3f%8.3f%8.3f%6.2f%6.2f\n"%( | |
| "ATOM", ctr, atm_j, num2aa[s], | |
| chain, i+1, atoms[i,j,0], atoms[i,j,1], atoms[i,j,2], | |
| 1.0, Bfacts[i] ) ) | |
| ctr += 1 | |
| else: | |
| atms = aa2long[s] | |
| for j,atm_j in enumerate(atms): | |
| if (atm_j is not None): | |
| f.write ("%-6s%5s %4s %3s %s%4d %8.3f%8.3f%8.3f%6.2f%6.2f\n"%( | |
| "ATOM", ctr, atm_j, num2aa[s], | |
| chain, i+1, atoms[i,j,0], atoms[i,j,1], atoms[i,j,2], | |
| 1.0, Bfacts[i] ) ) | |
| ctr += 1 | |
| def preprocess(xyz_t, t1d, DEVICE, masks_1d, ti_dev=None, ti_flip=None, ang_ref=None): | |
| B, _, L, _, _ = xyz_t.shape | |
| seq_tmp = t1d[...,:-1].argmax(dim=-1).reshape(-1,L).to(DEVICE, non_blocking=True) | |
| alpha, _, alpha_mask,_ = get_torsions(xyz_t.reshape(-1,L,27,3), seq_tmp, ti_dev, ti_flip, ang_ref) | |
| alpha_mask = torch.logical_and(alpha_mask, ~torch.isnan(alpha[...,0])) | |
| alpha[torch.isnan(alpha)] = 0.0 | |
| alpha = alpha.reshape(B,-1,L,10,2) | |
| alpha_mask = alpha_mask.reshape(B,-1,L,10,1) | |
| alpha_t = torch.cat((alpha, alpha_mask), dim=-1).reshape(B,-1,L,30) | |
| #t1d = torch.cat((t1d, chis.reshape(B,-1,L,30)), dim=-1) | |
| xyz_t = get_init_xyz(xyz_t) | |
| xyz_prev = xyz_t[:,0] | |
| state = t1d[:,0] | |
| alpha = alpha[:,0] | |
| t2d=xyz_to_t2d(xyz_t) | |
| return (t2d, alpha, alpha_mask, alpha_t, t1d, xyz_t, xyz_prev, state) | |
| def TemplFeaturizeFixbb(seq, conf_1d=None): | |
| """ | |
| Template 1D featurizer for fixed BB examples : | |
| Parameters: | |
| seq (torch.tensor, required): Integer sequence | |
| conf_1d (torch.tensor, optional): Precalcualted confidence tensor | |
| """ | |
| L = seq.shape[-1] | |
| t1d = torch.nn.functional.one_hot(seq, num_classes=21) # one hot sequence | |
| if conf_1d is None: | |
| conf = torch.ones_like(seq)[...,None] | |
| else: | |
| conf = conf_1d[:,None] | |
| t1d = torch.cat((t1d, conf), dim=-1) | |
| return t1d | |
| def MSAFeaturize_fixbb(msa, params): | |
| ''' | |
| Input: full msa information | |
| Output: Single sequence, with some percentage of amino acids mutated (but no resides 'masked') | |
| This is modified from autofold2, to remove mutations of the single sequence | |
| ''' | |
| N, L = msa.shape | |
| # raw MSA profile | |
| raw_profile = torch.nn.functional.one_hot(msa, num_classes=22) | |
| raw_profile = raw_profile.float().mean(dim=0) | |
| b_seq = list() | |
| b_msa_clust = list() | |
| b_msa_seed = list() | |
| b_msa_extra = list() | |
| b_mask_pos = list() | |
| for i_cycle in range(params['MAXCYCLE']): | |
| assert torch.max(msa) < 22 | |
| msa_onehot = torch.nn.functional.one_hot(msa[:1],num_classes=22) | |
| msa_fakeprofile_onehot = torch.nn.functional.one_hot(msa[:1],num_classes=26) #add the extra two indel planes, which will be set to zero | |
| msa_full_onehot = torch.cat((msa_onehot, msa_fakeprofile_onehot), dim=-1) | |
| #make fake msa_extra | |
| msa_extra_onehot = torch.nn.functional.one_hot(msa[:1],num_classes=25) | |
| #make fake msa_clust and mask_pos | |
| msa_clust = msa[:1] | |
| mask_pos = torch.full_like(msa_clust, 1).bool() | |
| b_seq.append(msa[0].clone()) | |
| b_msa_seed.append(msa_full_onehot[:1].clone()) #masked single sequence onehot (nb no mask so just single sequence onehot) | |
| b_msa_extra.append(msa_extra_onehot[:1].clone()) #masked single sequence onehot (nb no mask so just single sequence onehot) | |
| b_msa_clust.append(msa_clust[:1].clone()) #unmasked original single sequence | |
| b_mask_pos.append(mask_pos[:1].clone()) #mask positions in single sequence (all zeros) | |
| b_seq = torch.stack(b_seq) | |
| b_msa_clust = torch.stack(b_msa_clust) | |
| b_msa_seed = torch.stack(b_msa_seed) | |
| b_msa_extra = torch.stack(b_msa_extra) | |
| b_mask_pos = torch.stack(b_mask_pos) | |
| return b_seq, b_msa_clust, b_msa_seed, b_msa_extra, b_mask_pos | |
| def MSAFeaturize(msa, params): | |
| ''' | |
| Input: full msa information | |
| Output: Single sequence, with some percentage of amino acids mutated (but no resides 'masked') | |
| This is modified from autofold2, to remove mutations of the single sequence | |
| ''' | |
| N, L = msa.shape | |
| # raw MSA profile | |
| raw_profile = torch.nn.functional.one_hot(msa, num_classes=22) | |
| raw_profile = raw_profile.float().mean(dim=0) | |
| b_seq = list() | |
| b_msa_clust = list() | |
| b_msa_seed = list() | |
| b_msa_extra = list() | |
| b_mask_pos = list() | |
| for i_cycle in range(params['MAXCYCLE']): | |
| assert torch.max(msa) < 22 | |
| msa_onehot = torch.nn.functional.one_hot(msa,num_classes=22) | |
| msa_fakeprofile_onehot = torch.nn.functional.one_hot(msa,num_classes=26) #add the extra two indel planes, which will be set to zero | |
| msa_full_onehot = torch.cat((msa_onehot, msa_fakeprofile_onehot), dim=-1) | |
| #make fake msa_extra | |
| msa_extra_onehot = torch.nn.functional.one_hot(msa,num_classes=25) | |
| #make fake msa_clust and mask_pos | |
| msa_clust = msa | |
| mask_pos = torch.full_like(msa_clust, 1).bool() | |
| b_seq.append(msa[0].clone()) | |
| b_msa_seed.append(msa_full_onehot.clone()) #masked single sequence onehot (nb no mask so just single sequence onehot) | |
| b_msa_extra.append(msa_extra_onehot.clone()) #masked single sequence onehot (nb no mask so just single sequence onehot) | |
| b_msa_clust.append(msa_clust.clone()) #unmasked original single sequence | |
| b_mask_pos.append(mask_pos.clone()) #mask positions in single sequence (all zeros) | |
| b_seq = torch.stack(b_seq) | |
| b_msa_clust = torch.stack(b_msa_clust) | |
| b_msa_seed = torch.stack(b_msa_seed) | |
| b_msa_extra = torch.stack(b_msa_extra) | |
| b_mask_pos = torch.stack(b_mask_pos) | |
| return b_seq, b_msa_clust, b_msa_seed, b_msa_extra, b_mask_pos | |
| def mask_inputs(seq, msa_masked, msa_full, xyz_t, t1d, input_seq_mask=None, input_str_mask=None, input_t1dconf_mask=None, loss_seq_mask=None, loss_str_mask=None): | |
| """ | |
| Parameters: | |
| seq (torch.tensor, required): (B,I,L) integer sequence | |
| msa_masked (torch.tensor, required): (B,I,N_short,L,46) | |
| msa_full (torch,.tensor, required): (B,I,N_long,L,23) | |
| xyz_t (torch,tensor): (B,T,L,14,3) template crds BEFORE they go into get_init_xyz | |
| t1d (torch.tensor, required): (B,I,L,22) this is the t1d before tacking on the chi angles | |
| str_mask_1D (torch.tensor, required): Shape (L) rank 1 tensor where structure is masked at False positions | |
| seq_mask_1D (torch.tensor, required): Shape (L) rank 1 tensor where seq is masked at False positions | |
| """ | |
| ########### | |
| B,_,_ = seq.shape | |
| assert B == 1, 'batch sizes > 1 not supported' | |
| seq_mask = input_seq_mask[0] | |
| seq[:,:,~seq_mask] = 21 # mask token categorical value | |
| ### msa_masked ### | |
| ################## | |
| msa_masked[:,:,:,~seq_mask,:20] = 0 | |
| msa_masked[:,:,:,~seq_mask,20] = 0 | |
| msa_masked[:,:,:,~seq_mask,21] = 1 # set to the unkown char | |
| # index 44/45 is insertion/deletion | |
| # index 43 is the unknown token | |
| # index 42 is the masked token | |
| msa_masked[:,:,:,~seq_mask,22:42] = 0 | |
| msa_masked[:,:,:,~seq_mask,43] = 1 | |
| msa_masked[:,:,:,~seq_mask,42] = 0 | |
| # insertion/deletion stuff | |
| msa_masked[:,:,:,~seq_mask,44:] = 0 | |
| ### msa_full ### | |
| ################ | |
| msa_full[:,:,:,~seq_mask,:20] = 0 | |
| msa_full[:,:,:,~seq_mask,21] = 1 | |
| msa_full[:,:,:,~seq_mask,20] = 0 | |
| msa_full[:,:,:,~seq_mask,-1] = 0 #NOTE: double check this is insertions/deletions and 0 makes sense | |
| ### t1d ### | |
| ########### | |
| # NOTE: Not adjusting t1d last dim (confidence) from sequence mask | |
| t1d[:,:,~seq_mask,:20] = 0 | |
| t1d[:,:,~seq_mask,20] = 1 # unknown | |
| t1d[:,:,:,21] *= input_t1dconf_mask | |
| #JG added in here to make sure everything fits | |
| print('expanding t1d to 24 dims') | |
| t1d = torch.cat((t1d, torch.zeros((t1d.shape[0],t1d.shape[1],t1d.shape[2],2)).float()), -1).to(seq.device) | |
| xyz_t[:,:,~seq_mask,3:,:] = float('nan') | |
| # Structure masking | |
| str_mask = input_str_mask[0] | |
| xyz_t[:,:,~str_mask,:,:] = float('nan') | |
| return seq, msa_masked, msa_full, xyz_t, t1d | |
| ########################################################### | |
| #Functions for randomly translating/rotation input residues | |
| ########################################################### | |
| def get_translated_coords(args): | |
| ''' | |
| Parses args.res_translate | |
| ''' | |
| #get positions to translate | |
| res_translate = [] | |
| for res in args.res_translate.split(":"): | |
| temp_str = [] | |
| for i in res.split(','): | |
| temp_str.append(i) | |
| if temp_str[-1][0].isalpha() is True: | |
| temp_str.append(2.0) #set default distance | |
| for i in temp_str[:-1]: | |
| if '-' in i: | |
| start = int(i.split('-')[0][1:]) | |
| while start <= int(i.split('-')[1]): | |
| res_translate.append((i.split('-')[0][0] + str(start),float(temp_str[-1]))) | |
| start += 1 | |
| else: | |
| res_translate.append((i, float(temp_str[-1]))) | |
| start = 0 | |
| output = [] | |
| for i in res_translate: | |
| temp = (i[0], i[1], start) | |
| output.append(temp) | |
| start += 1 | |
| return output | |
| def get_tied_translated_coords(args, untied_translate=None): | |
| ''' | |
| Parses args.tie_translate | |
| ''' | |
| #pdb_idx = list(parsed_pdb['idx']) | |
| #xyz = parsed_pdb['xyz'] | |
| #get positions to translate | |
| res_translate = [] | |
| block = 0 | |
| for res in args.tie_translate.split(":"): | |
| temp_str = [] | |
| for i in res.split(','): | |
| temp_str.append(i) | |
| if temp_str[-1][0].isalpha() is True: | |
| temp_str.append(2.0) #set default distance | |
| for i in temp_str[:-1]: | |
| if '-' in i: | |
| start = int(i.split('-')[0][1:]) | |
| while start <= int(i.split('-')[1]): | |
| res_translate.append((i.split('-')[0][0] + str(start),float(temp_str[-1]), block)) | |
| start += 1 | |
| else: | |
| res_translate.append((i, float(temp_str[-1]), block)) | |
| block += 1 | |
| #sanity check | |
| if untied_translate != None: | |
| checker = [i[0] for i in res_translate] | |
| untied_check = [i[0] for i in untied_translate] | |
| for i in checker: | |
| if i in untied_check: | |
| print(f'WARNING: residue {i} is specified both in --res_translate and --tie_translate. Residue {i} will be ignored in --res_translate, and instead only moved in a tied block (--tie_translate)') | |
| final_output = res_translate | |
| for i in untied_translate: | |
| if i[0] not in checker: | |
| final_output.append((i[0],i[1],i[2] + block + 1)) | |
| else: | |
| final_output = res_translate | |
| return final_output | |
| def translate_coords(parsed_pdb, res_translate): | |
| ''' | |
| Takes parsed list in format [(chain_residue,distance,tieing_block)] and randomly translates residues accordingly. | |
| ''' | |
| pdb_idx = parsed_pdb['pdb_idx'] | |
| xyz = np.copy(parsed_pdb['xyz']) | |
| translated_coord_dict = {} | |
| #get number of blocks | |
| temp = [int(i[2]) for i in res_translate] | |
| blocks = np.max(temp) | |
| for block in range(blocks + 1): | |
| init_dist = 1.01 | |
| while init_dist > 1: #gives equal probability to any direction (as keeps going until init_dist is within unit circle) | |
| x = random.uniform(-1,1) | |
| y = random.uniform(-1,1) | |
| z = random.uniform(-1,1) | |
| init_dist = np.sqrt(x**2 + y**2 + z**2) | |
| x=x/init_dist | |
| y=y/init_dist | |
| z=z/init_dist | |
| translate_dist = random.uniform(0,1) #now choose distance (as proportion of maximum) that coordinates will be translated | |
| for res in res_translate: | |
| if res[2] == block: | |
| res_idx = pdb_idx.index((res[0][0],int(res[0][1:]))) | |
| original_coords = np.copy(xyz[res_idx,:,:]) | |
| for i in range(14): | |
| if parsed_pdb['mask'][res_idx, i]: | |
| xyz[res_idx,i,0] += np.float32(x * translate_dist * float(res[1])) | |
| xyz[res_idx,i,1] += np.float32(y * translate_dist * float(res[1])) | |
| xyz[res_idx,i,2] += np.float32(z * translate_dist * float(res[1])) | |
| translated_coords = xyz[res_idx,:,:] | |
| translated_coord_dict[res[0]] = (original_coords.tolist(), translated_coords.tolist()) | |
| return xyz[:,:,:], translated_coord_dict | |
| def parse_block_rotate(args): | |
| block_translate = [] | |
| block = 0 | |
| for res in args.block_rotate.split(":"): | |
| temp_str = [] | |
| for i in res.split(','): | |
| temp_str.append(i) | |
| if temp_str[-1][0].isalpha() is True: | |
| temp_str.append(10) #set default angle to 10 degrees | |
| for i in temp_str[:-1]: | |
| if '-' in i: | |
| start = int(i.split('-')[0][1:]) | |
| while start <= int(i.split('-')[1]): | |
| block_translate.append((i.split('-')[0][0] + str(start),float(temp_str[-1]), block)) | |
| start += 1 | |
| else: | |
| block_translate.append((i, float(temp_str[-1]), block)) | |
| block += 1 | |
| return block_translate | |
| def rotate_block(xyz, block_rotate,pdb_index): | |
| rotated_coord_dict = {} | |
| #get number of blocks | |
| temp = [int(i[2]) for i in block_rotate] | |
| blocks = np.max(temp) | |
| for block in range(blocks + 1): | |
| idxs = [pdb_index.index((i[0][0],int(i[0][1:]))) for i in block_rotate if i[2] == block] | |
| angle = [i[1] for i in block_rotate if i[2] == block][0] | |
| block_xyz = xyz[idxs,:,:] | |
| com = [float(torch.mean(block_xyz[:,:,i])) for i in range(3)] | |
| origin_xyz = np.copy(block_xyz) | |
| for i in range(np.shape(origin_xyz)[0]): | |
| for j in range(14): | |
| origin_xyz[i,j] = origin_xyz[i,j] - com | |
| rotated_xyz = rigid_rotate(origin_xyz,angle,angle,angle) | |
| recovered_xyz = np.copy(rotated_xyz) | |
| for i in range(np.shape(origin_xyz)[0]): | |
| for j in range(14): | |
| recovered_xyz[i,j] = rotated_xyz[i,j] + com | |
| recovered_xyz=torch.tensor(recovered_xyz) | |
| rotated_coord_dict[f'rotated_block_{block}_original'] = block_xyz | |
| rotated_coord_dict[f'rotated_block_{block}_rotated'] = recovered_xyz | |
| xyz_out = torch.clone(xyz) | |
| for i in range(len(idxs)): | |
| xyz_out[idxs[i]] = recovered_xyz[i] | |
| return xyz_out,rotated_coord_dict | |
| def rigid_rotate(xyz,a=180,b=180,c=180): | |
| #TODO fix this to make it truly uniform | |
| a=(a/180)*math.pi | |
| b=(b/180)*math.pi | |
| c=(c/180)*math.pi | |
| alpha = random.uniform(-a, a) | |
| beta = random.uniform(-b, b) | |
| gamma = random.uniform(-c, c) | |
| rotated = [] | |
| for i in range(np.shape(xyz)[0]): | |
| for j in range(14): | |
| try: | |
| x = xyz[i,j,0] | |
| y = xyz[i,j,1] | |
| z = xyz[i,j,2] | |
| x2 = x*math.cos(alpha) - y*math.sin(alpha) | |
| y2 = x*math.sin(alpha) + y*math.cos(alpha) | |
| x3 = x2*math.cos(beta) - z*math.sin(beta) | |
| z2 = x2*math.sin(beta) + z*math.cos(beta) | |
| y3 = y2*math.cos(gamma) - z2*math.sin(gamma) | |
| z3 = y2*math.sin(gamma) + z2*math.cos(gamma) | |
| rotated.append([x3,y3,z3]) | |
| except: | |
| rotated.append([float('nan'),float('nan'),float('nan')]) | |
| rotated=np.array(rotated) | |
| rotated=np.reshape(rotated, [np.shape(xyz)[0],14,3]) | |
| return rotated | |
| ######## from old pred_util.py | |
| def find_contigs(mask): | |
| """ | |
| Find contiguous regions in a mask that are True with no False in between | |
| Parameters: | |
| mask (torch.tensor or np.array, required): 1D boolean array | |
| Returns: | |
| contigs (list): List of tuples, each tuple containing the beginning and the | |
| """ | |
| assert len(mask.shape) == 1 # 1D tensor of bools | |
| contigs = [] | |
| found_contig = False | |
| for i,b in enumerate(mask): | |
| if b and not found_contig: # found the beginning of a contig | |
| contig = [i] | |
| found_contig = True | |
| elif b and found_contig: # currently have contig, continuing it | |
| pass | |
| elif not b and found_contig: # found the end, record previous index as end, reset indicator | |
| contig.append(i) | |
| found_contig = False | |
| contigs.append(tuple(contig)) | |
| else: # currently don't have a contig, and didn't find one | |
| pass | |
| # fence post bug - check if the very last entry was True and we didn't get to finish | |
| if b: | |
| contig.append(i+1) | |
| found_contig = False | |
| contigs.append(tuple(contig)) | |
| return contigs | |
| def reindex_chains(pdb_idx): | |
| """ | |
| Given a list of (chain, index) tuples, and the indices where chains break, create a reordered indexing | |
| Parameters: | |
| pdb_idx (list, required): List of tuples (chainID, index) | |
| breaks (list, required): List of indices where chains begin | |
| """ | |
| new_breaks, new_idx = [],[] | |
| current_chain = None | |
| chain_and_idx_to_torch = {} | |
| for i,T in enumerate(pdb_idx): | |
| chain, idx = T | |
| if chain != current_chain: | |
| new_breaks.append(i) | |
| current_chain = chain | |
| # create new space for chain id listings | |
| chain_and_idx_to_torch[chain] = {} | |
| # map original pdb (chain, idx) pair to index in tensor | |
| chain_and_idx_to_torch[chain][idx] = i | |
| # append tensor index to list | |
| new_idx.append(i) | |
| new_idx = np.array(new_idx) | |
| # now we have ordered list and know where the chainbreaks are in the new order | |
| num_additions = 0 | |
| for i in new_breaks[1:]: # skip the first trivial one | |
| new_idx[np.where(new_idx==(i+ num_additions*500))[0][0]:] += 500 | |
| num_additions += 1 | |
| return new_idx, chain_and_idx_to_torch,new_breaks[1:] | |
| class ObjectView(object): | |
| ''' | |
| Easy wrapper to access dictionary values with "dot" notiation instead | |
| ''' | |
| def __init__(self, d): | |
| self.__dict__ = d | |
| def split_templates(xyz_t, t1d, multi_templates,mappings,multi_tmpl_conf=None): | |
| templates = multi_templates.split(":") | |
| if multi_tmpl_conf is not None: | |
| multi_tmpl_conf = [float(i) for i in multi_tmpl_conf.split(",")] | |
| assert len(templates) == len(multi_tmpl_conf), "Number of templates must equal number of confidences specified in --multi_tmpl_conf flag" | |
| for idx, template in enumerate(templates): | |
| parts = template.split(",") | |
| template_mask = torch.zeros(xyz_t.shape[2]).bool() | |
| for part in parts: | |
| start = int(part.split("-")[0][1:]) | |
| end = int(part.split("-")[1]) + 1 | |
| chain = part[0] | |
| for i in range(start, end): | |
| try: | |
| ref_pos = mappings['complex_con_ref_pdb_idx'].index((chain, i)) | |
| hal_pos_0 = mappings['complex_con_hal_idx0'][ref_pos] | |
| except: | |
| ref_pos = mappings['con_ref_pdb_idx'].index((chain, i)) | |
| hal_pos_0 = mappings['con_hal_idx0'][ref_pos] | |
| template_mask[hal_pos_0] = True | |
| xyz_t_temp = torch.clone(xyz_t) | |
| xyz_t_temp[:,:,~template_mask,:,:] = float('nan') | |
| t1d_temp = torch.clone(t1d) | |
| t1d_temp[:,:,~template_mask,:20] =0 | |
| t1d_temp[:,:,~template_mask,20] = 1 | |
| if multi_tmpl_conf is not None: | |
| t1d_temp[:,:,template_mask,21] = multi_tmpl_conf[idx] | |
| if idx != 0: | |
| xyz_t_out = torch.cat((xyz_t_out, xyz_t_temp),dim=1) | |
| t1d_out = torch.cat((t1d_out, t1d_temp),dim=1) | |
| else: | |
| xyz_t_out = xyz_t_temp | |
| t1d_out = t1d_temp | |
| return xyz_t_out, t1d_out | |
| class ContigMap(): | |
| ''' | |
| New class for doing mapping. | |
| Supports multichain or multiple crops from a single receptor chain. | |
| Also supports indexing jump (+200) or not, based on contig input. | |
| Default chain outputs are inpainted chains as A (and B, C etc if multiple chains), and all fragments of receptor chain on the next one (generally B) | |
| Output chains can be specified. Sequence must be the same number of elements as in contig string | |
| ''' | |
| def __init__(self, parsed_pdb, contigs=None, inpaint_seq=None, inpaint_str=None, length=None, ref_idx=None, hal_idx=None, idx_rf=None, inpaint_seq_tensor=None, inpaint_str_tensor=None, topo=False): | |
| #sanity checks | |
| if contigs is None and ref_idx is None: | |
| sys.exit("Must either specify a contig string or precise mapping") | |
| if idx_rf is not None or hal_idx is not None or ref_idx is not None: | |
| if idx_rf is None or hal_idx is None or ref_idx is None: | |
| sys.exit("If you're specifying specific contig mappings, the reference and output positions must be specified, AND the indexing for RoseTTAFold (idx_rf)") | |
| self.chain_order='ABCDEFGHIJKLMNOPQRSTUVWXYZ' | |
| if length is not None: | |
| if '-' not in length: | |
| self.length = [int(length),int(length)+1] | |
| else: | |
| self.length = [int(length.split("-")[0]),int(length.split("-")[1])+1] | |
| else: | |
| self.length = None | |
| self.ref_idx = ref_idx | |
| self.hal_idx=hal_idx | |
| self.idx_rf=idx_rf | |
| self.inpaint_seq = ','.join(inpaint_seq).split(",") if inpaint_seq is not None else None | |
| self.inpaint_str = ','.join(inpaint_str).split(",") if inpaint_str is not None else None | |
| self.inpaint_seq_tensor=inpaint_seq_tensor | |
| self.inpaint_str_tensor=inpaint_str_tensor | |
| self.parsed_pdb = parsed_pdb | |
| self.topo=topo | |
| if ref_idx is None: | |
| #using default contig generation, which outputs in rosetta-like format | |
| self.contigs=contigs | |
| self.sampled_mask,self.contig_length,self.n_inpaint_chains = self.get_sampled_mask() | |
| self.receptor_chain = self.chain_order[self.n_inpaint_chains] | |
| self.receptor, self.receptor_hal, self.receptor_rf, self.inpaint, self.inpaint_hal, self.inpaint_rf= self.expand_sampled_mask() | |
| self.ref = self.inpaint + self.receptor | |
| self.hal = self.inpaint_hal + self.receptor_hal | |
| self.rf = self.inpaint_rf + self.receptor_rf | |
| else: | |
| #specifying precise mappings | |
| self.ref=ref_idx | |
| self.hal=hal_idx | |
| self.rf = rf_idx | |
| self.mask_1d = [False if i == ('_','_') else True for i in self.ref] | |
| #take care of sequence and structure masking | |
| if self.inpaint_seq_tensor is None: | |
| if self.inpaint_seq is not None: | |
| self.inpaint_seq = self.get_inpaint_seq_str(self.inpaint_seq) | |
| else: | |
| self.inpaint_seq = np.array([True if i != ('_','_') else False for i in self.ref]) | |
| else: | |
| self.inpaint_seq = self.inpaint_seq_tensor | |
| if self.inpaint_str_tensor is None: | |
| if self.inpaint_str is not None: | |
| self.inpaint_str = self.get_inpaint_seq_str(self.inpaint_str) | |
| else: | |
| self.inpaint_str = np.array([True if i != ('_','_') else False for i in self.ref]) | |
| else: | |
| self.inpaint_str = self.inpaint_str_tensor | |
| #get 0-indexed input/output (for trb file) | |
| self.ref_idx0,self.hal_idx0, self.ref_idx0_inpaint, self.hal_idx0_inpaint, self.ref_idx0_receptor, self.hal_idx0_receptor=self.get_idx0() | |
| def get_sampled_mask(self): | |
| ''' | |
| Function to get a sampled mask from a contig. | |
| ''' | |
| length_compatible=False | |
| count = 0 | |
| while length_compatible is False: | |
| inpaint_chains=0 | |
| contig_list = self.contigs | |
| sampled_mask = [] | |
| sampled_mask_length = 0 | |
| #allow receptor chain to be last in contig string | |
| if all([i[0].isalpha() for i in contig_list[-1].split(",")]): | |
| contig_list[-1] = f'{contig_list[-1]},0' | |
| for con in contig_list: | |
| if ((all([i[0].isalpha() for i in con.split(",")[:-1]]) and con.split(",")[-1] == '0')) or self.topo is True: | |
| #receptor chain | |
| sampled_mask.append(con) | |
| else: | |
| inpaint_chains += 1 | |
| #chain to be inpainted. These are the only chains that count towards the length of the contig | |
| subcons = con.split(",") | |
| subcon_out = [] | |
| for subcon in subcons: | |
| if subcon[0].isalpha(): | |
| subcon_out.append(subcon) | |
| if '-' in subcon: | |
| sampled_mask_length += (int(subcon.split("-")[1])-int(subcon.split("-")[0][1:])+1) | |
| else: | |
| sampled_mask_length += 1 | |
| else: | |
| if '-' in subcon: | |
| length_inpaint=random.randint(int(subcon.split("-")[0]),int(subcon.split("-")[1])) | |
| subcon_out.append(f'{length_inpaint}-{length_inpaint}') | |
| sampled_mask_length += length_inpaint | |
| elif subcon == '0': | |
| subcon_out.append('0') | |
| else: | |
| length_inpaint=int(subcon) | |
| subcon_out.append(f'{length_inpaint}-{length_inpaint}') | |
| sampled_mask_length += int(subcon) | |
| sampled_mask.append(','.join(subcon_out)) | |
| #check length is compatible | |
| if self.length is not None: | |
| if sampled_mask_length >= self.length[0] and sampled_mask_length < self.length[1]: | |
| length_compatible = True | |
| else: | |
| length_compatible = True | |
| count+=1 | |
| if count == 100000: #contig string incompatible with this length | |
| sys.exit("Contig string incompatible with --length range") | |
| return sampled_mask, sampled_mask_length, inpaint_chains | |
| def expand_sampled_mask(self): | |
| chain_order='ABCDEFGHIJKLMNOPQRSTUVWXYZ' | |
| receptor = [] | |
| inpaint = [] | |
| receptor_hal = [] | |
| inpaint_hal = [] | |
| receptor_idx = 1 | |
| inpaint_idx = 1 | |
| inpaint_chain_idx=-1 | |
| receptor_chain_break=[] | |
| inpaint_chain_break = [] | |
| for con in self.sampled_mask: | |
| if (all([i[0].isalpha() for i in con.split(",")[:-1]]) and con.split(",")[-1] == '0') or self.topo is True: | |
| #receptor chain | |
| subcons = con.split(",")[:-1] | |
| assert all([i[0] == subcons[0][0] for i in subcons]), "If specifying fragmented receptor in a single block of the contig string, they MUST derive from the same chain" | |
| assert all(int(subcons[i].split("-")[0][1:]) < int(subcons[i+1].split("-")[0][1:]) for i in range(len(subcons)-1)), "If specifying multiple fragments from the same chain, pdb indices must be in ascending order!" | |
| for idx, subcon in enumerate(subcons): | |
| ref_to_add = [(subcon[0], i) for i in np.arange(int(subcon.split("-")[0][1:]),int(subcon.split("-")[1])+1)] | |
| receptor.extend(ref_to_add) | |
| receptor_hal.extend([(self.receptor_chain,i) for i in np.arange(receptor_idx, receptor_idx+len(ref_to_add))]) | |
| receptor_idx += len(ref_to_add) | |
| if idx != len(subcons)-1: | |
| idx_jump = int(subcons[idx+1].split("-")[0][1:]) - int(subcon.split("-")[1]) -1 | |
| receptor_chain_break.append((receptor_idx-1,idx_jump)) #actual chain break in pdb chain | |
| else: | |
| receptor_chain_break.append((receptor_idx-1,200)) #200 aa chain break | |
| else: | |
| inpaint_chain_idx += 1 | |
| for subcon in con.split(","): | |
| if subcon[0].isalpha(): | |
| ref_to_add=[(subcon[0], i) for i in np.arange(int(subcon.split("-")[0][1:]),int(subcon.split("-")[1])+1)] | |
| inpaint.extend(ref_to_add) | |
| inpaint_hal.extend([(chain_order[inpaint_chain_idx], i) for i in np.arange(inpaint_idx,inpaint_idx+len(ref_to_add))]) | |
| inpaint_idx += len(ref_to_add) | |
| else: | |
| inpaint.extend([('_','_')] * int(subcon.split("-")[0])) | |
| inpaint_hal.extend([(chain_order[inpaint_chain_idx], i) for i in np.arange(inpaint_idx,inpaint_idx+int(subcon.split("-")[0]))]) | |
| inpaint_idx += int(subcon.split("-")[0]) | |
| inpaint_chain_break.append((inpaint_idx-1,200)) | |
| if self.topo is True or inpaint_hal == []: | |
| receptor_hal = [(i[0], i[1]) for i in receptor_hal] | |
| else: | |
| receptor_hal = [(i[0], i[1] + inpaint_hal[-1][1]) for i in receptor_hal] #rosetta-like numbering | |
| #get rf indexes, with chain breaks | |
| inpaint_rf = np.arange(0,len(inpaint)) | |
| receptor_rf = np.arange(len(inpaint)+200,len(inpaint)+len(receptor)+200) | |
| for ch_break in inpaint_chain_break[:-1]: | |
| receptor_rf[:] += 200 | |
| inpaint_rf[ch_break[0]:] += ch_break[1] | |
| for ch_break in receptor_chain_break[:-1]: | |
| receptor_rf[ch_break[0]:] += ch_break[1] | |
| return receptor, receptor_hal, receptor_rf.tolist(), inpaint, inpaint_hal, inpaint_rf.tolist() | |
| def get_inpaint_seq_str(self, inpaint_s): | |
| ''' | |
| function to generate inpaint_str or inpaint_seq masks specific to this contig | |
| ''' | |
| s_mask = np.copy(self.mask_1d) | |
| inpaint_s_list = [] | |
| for i in inpaint_s: | |
| if '-' in i: | |
| inpaint_s_list.extend([(i[0],p) for p in range(int(i.split("-")[0][1:]), int(i.split("-")[1])+1)]) | |
| else: | |
| inpaint_s_list.append((i[0],int(i[1:]))) | |
| for res in inpaint_s_list: | |
| if res in self.ref: | |
| s_mask[self.ref.index(res)] = False #mask this residue | |
| return np.array(s_mask) | |
| def get_idx0(self): | |
| ref_idx0=[] | |
| hal_idx0=[] | |
| ref_idx0_inpaint=[] | |
| hal_idx0_inpaint=[] | |
| ref_idx0_receptor=[] | |
| hal_idx0_receptor=[] | |
| for idx, val in enumerate(self.ref): | |
| if val != ('_','_'): | |
| assert val in self.parsed_pdb['pdb_idx'],f"{val} is not in pdb file!" | |
| hal_idx0.append(idx) | |
| ref_idx0.append(self.parsed_pdb['pdb_idx'].index(val)) | |
| for idx, val in enumerate(self.inpaint): | |
| if val != ('_','_'): | |
| hal_idx0_inpaint.append(idx) | |
| ref_idx0_inpaint.append(self.parsed_pdb['pdb_idx'].index(val)) | |
| for idx, val in enumerate(self.receptor): | |
| if val != ('_','_'): | |
| hal_idx0_receptor.append(idx) | |
| ref_idx0_receptor.append(self.parsed_pdb['pdb_idx'].index(val)) | |
| return ref_idx0, hal_idx0, ref_idx0_inpaint, hal_idx0_inpaint, ref_idx0_receptor, hal_idx0_receptor | |
| def get_mappings(rm): | |
| mappings = {} | |
| mappings['con_ref_pdb_idx'] = [i for i in rm.inpaint if i != ('_','_')] | |
| mappings['con_hal_pdb_idx'] = [rm.inpaint_hal[i] for i in range(len(rm.inpaint_hal)) if rm.inpaint[i] != ("_","_")] | |
| mappings['con_ref_idx0'] = rm.ref_idx0_inpaint | |
| mappings['con_hal_idx0'] = rm.hal_idx0_inpaint | |
| if rm.inpaint != rm.ref: | |
| mappings['complex_con_ref_pdb_idx'] = [i for i in rm.ref if i != ("_","_")] | |
| mappings['complex_con_hal_pdb_idx'] = [rm.hal[i] for i in range(len(rm.hal)) if rm.ref[i] != ("_","_")] | |
| mappings['receptor_con_ref_pdb_idx'] = [i for i in rm.receptor if i != ("_","_")] | |
| mappings['receptor_con_hal_pdb_idx'] = [rm.receptor_hal[i] for i in range(len(rm.receptor_hal)) if rm.receptor[i] != ("_","_")] | |
| mappings['complex_con_ref_idx0'] = rm.ref_idx0 | |
| mappings['complex_con_hal_idx0'] = rm.hal_idx0 | |
| mappings['receptor_con_ref_idx0'] = rm.ref_idx0_receptor | |
| mappings['receptor_con_hal_idx0'] = rm.hal_idx0_receptor | |
| mappings['inpaint_str'] = rm.inpaint_str | |
| mappings['inpaint_seq'] = rm.inpaint_seq | |
| mappings['sampled_mask'] = rm.sampled_mask | |
| mappings['mask_1d'] = rm.mask_1d | |
| return mappings | |
| def lddt_unbin(pred_lddt): | |
| nbin = pred_lddt.shape[1] | |
| bin_step = 1.0 / nbin | |
| lddt_bins = torch.linspace(bin_step, 1.0, nbin, dtype=pred_lddt.dtype, device=pred_lddt.device) | |
| pred_lddt = nn.Softmax(dim=1)(pred_lddt) | |
| return torch.sum(lddt_bins[None,:,None]*pred_lddt, dim=1) | |