model/utils/inpainting_util.py

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, 'wt') 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)