# utility functions for dealing with contigs during hallucination
import numpy as np
import random, copy, torch, geometry, os, sys
from kinematics import xyz_to_t2d
def parse_range_string(el):
''' Splits string with integer or integer range into start and end ints. '''
if '-' in el:
s,e = el.split('-')
s,e = int(s), int(e)
else:
s,e = int(el), int(el)
return s,e
def ranges_to_indexes(range_string):
'''Converts a string containig comma-separated numeric ranges to a list of integers'''
idx = []
for x in range_string.split(','):
start, end = parse_range_string(x)
idx.extend(np.arange(start, end+1))
return np.array(idx)
def parse_contigs(contig_input, pdb_id):
'''
Input: contig start/end by pdb chain and residue number as in the pdb file
ex - B12-17
Output: corresponding start/end indices of the "features" numpy array (idx0)
'''
contigs = []
for con in contig_input.split(','):
pdb_ch = con[0]
pdb_s, pdb_e = parse_range_string(con[1:])
np_s = pdb_id.index((pdb_ch, pdb_s))
np_e = pdb_id.index((pdb_ch, pdb_e))
contigs.append([np_s, np_e])
return contigs
def mk_feat_hal_and_mappings(hal_2_ref_idx0, pdb_out):
#####################################
# rearrange ref features according to hal_2_ref_idx0
#####################################
#1. find corresponding idx0 in hal and ref
hal_idx0 = []
ref_idx0 = []
for hal, ref in enumerate(hal_2_ref_idx0):
if ref is not None:
hal_idx0.append(hal)
ref_idx0.append(ref)
hal_idx0 = np.array(hal_idx0, dtype=int)
ref_idx0 = np.array(ref_idx0, dtype=int)
#2. rearrange the 6D features
hal_len = len(hal_2_ref_idx0)
if 'feat' in pdb_out:
d_feat = pdb_out['feat'].shape[3:]
feat_hal = np.zeros((1, hal_len, hal_len) + d_feat)
feat_ref = pdb_out['feat'] # (B,L,L,...)
feat_hal[:, hal_idx0[:,None], hal_idx0[None,:]] = feat_ref[:, ref_idx0[:,None], ref_idx0[None,:]]
else:
feat_hal = None
#3. make the 1d binary mask, for backwards compatibility
hal_2_ref_idx0 = np.array(hal_2_ref_idx0, dtype=np.float32) # convert None to NaN
mask_1d = (~np.isnan(hal_2_ref_idx0)).astype(float)
mask_1d = mask_1d[None]
#####################################
# mappings between hal and ref
#####################################
mappings = {
'con_hal_idx0': hal_idx0.tolist(),
'con_ref_idx0': ref_idx0.tolist(),
'con_hal_pdb_idx': [('A',i+1) for i in hal_idx0],
'con_ref_pdb_idx': [pdb_out['pdb_idx'][i] for i in ref_idx0],
'mask_1d': mask_1d,
}
return feat_hal, mappings
def scatter_feats(template_mask, feat_1d_ref=None, feat_2d_ref=None, pdb_idx=None):
'''
Scatters 1D and/or 2D reference features according to mappings in hal_2_ref_idx0
Inputs
----------
hal_2_ref_idx0: (list; length=L_hal)
List mapping hal_idx0 positions to ref_idx0 positions.
"None" used for indices that do not map to ref.
ex: [None, None, 3, 4, 5, None, None, None, 34, 35, 36]
feat_1d_ref: (np.array; (batch, L_ref, ...))
1D refence features to scatter
feat_1d_ref: (np.array; (batch, L_ref, L_ref, ...))
pdb_idx: (list)
List of pdb chain and residue numbers, in the order that pdb features were read/parsed.
Outputs
----------
feat_1d_hal: (np.array, (batch, L_hal, ...))
Scattered 1d reference features. "None" mappings are 0.
feat_2d_hal: (np.array, (batch, L_hal, L_hal, ...))
Scattered 2d reference features. "None" mappings are 0.
mappings: (dict)
Keeps track of corresponding possitions in ref and hal proteins.
'''
hal_2_ref_idx0, _ = contigs.sample_mask(template_mask, pdb_idx)
out = {}
# Find corresponding idx0 in hal and ref
hal_idx0 = []
ref_idx0 = []
hal_len = len(hal_2_ref_idx0)
for hal, ref in enumerate(hal_2_ref_idx0):
if ref is not None:
hal_idx0.append(hal)
ref_idx0.append(ref)
hal_idx0 = np.array(hal_idx0, dtype=int)
ref_idx0 = np.array(ref_idx0, dtype=int)
# Make the 1d binary mask, for backwards compatibility
hal_2_ref_idx0 = np.array(hal_2_ref_idx0, dtype=np.float32) # convert None to NaN
mask_1d = (~np.isnan(hal_2_ref_idx0)).astype(float)
mask_1d = mask_1d[None]
# scatter 2D features
if feat_2d_ref is not None:
B = feat_2d_ref.shape[0]
d_feat = feat_2d_ref.shape[3:]
feat_2d_hal = np.zeros((B, hal_len, hal_len)+d_feat)
feat_2d_hal[:, hal_idx0[:,None], hal_idx0[None,:]] = feat_2d_ref[:, ref_idx0[:,None], ref_idx0[None,:]]
out['feat_2d_hal'] = feat_2d_hal
# scatter 1D features
if feat_1d_ref is not None:
B = feat_1d_ref.shape[0]
d_feat = feat_1d_ref.shape[2:]
feat_1d_hal = np.zeros((B, hal_len)+d_feat)
feat_1d_hal[:, hal_idx0] = feat_1d_ref[:, ref_idx0]
out['feat_1d_hal'] = feat_1d_hal
# Mappings between hal and ref
mappings = {
'con_hal_idx0': hal_idx0.tolist(),
'con_ref_idx0': ref_idx0.tolist(),
'mask_1d': mask_1d,
}
if pdb_idx is not None:
mappings.update({
'con_hal_pdb_idx': [('A',i+1) for i in hal_idx0],
'con_ref_pdb_idx': [pdb_idx[i] for i in ref_idx0],
})
out['mappings'] = mappings
return out
def scatter_contigs(contigs, pdb_out, L_range, keep_order=False, min_gap=0):
'''
Randomly places contigs in a protein within the length range.
Inputs
Contig: A continuous range of residues from the pdb.
Inclusive of the begining and end
Must start with the chain number. Comma separated
ex: B6-11,A12-19
pdb_out: dictionary from the prep_input function
L_range: String range of possible lengths.
ex: 90-110
ex: 70
keep_order: keep contigs in the provided order or randomly permute
min_gap: minimum number of amino acids separating contigs
Outputs
feat_hal: target pdb features to hallucinate
mappings: dictionary of ways to convert from the hallucinated protein
to the reference protein
'''
ref_pdb_2_idx0 = {pdb_idx:i for i, pdb_idx in enumerate(pdb_out['pdb_idx'])}
#####################################
# make a map from hal_idx0 to ref_idx0. Has None for gap regions
#####################################
#1. Permute contig order
contigs = contigs.split(',')
if not keep_order:
random.shuffle(contigs)
#2. convert to ref_idx0
contigs_ref_idx0 = []
for con in contigs:
chain = con[0]
s, e = parse_range_string(con[1:])
contigs_ref_idx0.append( [ref_pdb_2_idx0[(chain, i)] for i in range(s, e+1)] )
#3. Add minimum gap size
for i in range(len(contigs_ref_idx0) - 1):
contigs_ref_idx0[i] += [None] * min_gap
#4. Sample protein length
L_low, L_high = parse_range_string(L_range)
L_hal = np.random.randint(L_low, L_high+1)
L_con = 0
for con in contigs_ref_idx0:
L_con += len(con)
L_gaps = L_hal - L_con
if L_gaps <= 1:
print("Error: The protein isn't long enough to incorporate all the contigs."
"Consider reduce the min_gap or increasing L_range")
return
#5. Randomly insert contigs into gaps
hal_2_ref_idx0 = np.array([None] * L_gaps, dtype=float) # inserting contigs into this
n_contigs = len(contigs_ref_idx0)
insertion_idxs = np.random.randint(L_gaps + 1, size=n_contigs)
insertion_idxs.sort()
for idx, con in zip(insertion_idxs[::-1], contigs_ref_idx0[::-1]):
hal_2_ref_idx0 = np.insert(hal_2_ref_idx0, idx, con)
#6. Convert mask to feat_hal and mappings
hal_2_ref_idx0 = [int(el) if ~np.isnan(el) else None for el in hal_2_ref_idx0] # convert nan to None
feat_hal, mappings = mk_feat_hal_and_mappings(hal_2_ref_idx0, pdb_out)
#7. Generate str of the sampled mask
contig_positive = np.array(hal_2_ref_idx0) != None
boundaries = np.where(np.diff(contig_positive))[0]
start_idx0 = np.concatenate([np.array([0]), boundaries+1])
end_idx0 = np.concatenate([boundaries, np.array([contig_positive.shape[0]])-1])
lengths = end_idx0 - start_idx0 + 1
is_contig = contig_positive[start_idx0]
sampled_mask = []
con_counter = 0
for i, is_con in enumerate(is_contig):
if is_con:
sampled_mask.append(contigs[con_counter])
con_counter += 1
else:
len_gap = lengths[i]
sampled_mask.append(f'{len_gap}-{len_gap}')
sampled_mask = ','.join(sampled_mask)
mappings['sampled_mask'] = sampled_mask
return feat_hal, mappings
def get_receptor_contig(ref_pdb_idx):
rec_pdb_idx = [idx for idx in ref_pdb_idx if idx[0]=='R']
return SampledMask.contract(rec_pdb_idx)
def mk_con_to_set(mask, set_id=None, args=None, ref_pdb_idx=None):
'''
Maps a mask or list of contigs to a set_id. If no set_id is provided, it treats
everything as set 0.
Input
-----------
mask (str): Mask or list of contigs. Ex: 3,B6-11,12,A12-19,9 or Ex: B6-11,A12-19
ref_pdb_idx (List(ch, res)): pdb idxs of the reference pdb. Ex: [(A, 2), (A, 3), ...]
args: Arguments object. Must have args.receptor
set_id (list): List of integers. Length must match contigs in mask. Ex: [0,1]
Output
-----------
con_to_set (dict): Maps str of contig to integer
'''
# Extract contigs
cons = [l for l in mask.split(',') if l[0].isalpha()]
# Assign all contigs to set 0 if set_id is not passed
if set_id is None:
set_id = [0] * len(cons)
con_to_set = dict(zip(cons, set_id))
# Assign receptor to set 0
if args.receptor:
receptor_contig = get_receptor_contig(ref_pdb_idx)
con_to_set.update({receptor_contig: 0})
return con_to_set
def parse_range(_range):
if '-' in _range:
s, e = _range.split('-')
else:
s, e = _range, _range
return int(s), int(e)
def parse_contig(contig):
'''
Return the chain, start and end residue in a contig or gap str.
Ex:
'A4-8' --> 'A', 4, 8
'A5' --> 'A', 5, 5
'4-8' --> None, 4, 8
'A' --> 'A', None, None
'''
# is contig
if contig[0].isalpha():
ch = contig[0]
if len(contig) > 1:
s, e = parse_range(contig[1:])
else:
s, e = None, None
# is gap
else:
ch = None
s, e = parse_range(contig)
return ch, s, e
def mask_as_list(sampled_mask):
'''
Make a length L_hal list, with each position pointing to a ref_pdb_idx (or None)
'''
mask_list = []
for l in sampled_mask.split(','):
ch, s, e = parse_contig(l)
# contig
if ch is not None:
mask_list += [(ch, idx) for idx in range(s, e+1)]
# gap
else:
mask_list += [None for _ in range(s, e+1)]
return mask_list
def mask_subset(sampled_mask, subset):
'''
Returns a 1D boolean array of where a subset of the contig is in the hallucinated protein
Input
---------
subset (str): Some chain and residue subset of the contigs. Ex: A10-15
Can also just pass chain. All contig residues from that chain are selected. Ex: R
Ouput
---------
m_1d (np.array): Boolean array where subset appears in the hallucinated protein
'''
mask_list = mask_as_list(sampled_mask)
m_1d = []
ch_subset, s, e = parse_contig(subset)
assert ch_subset.isalpha(), '"Subset" must include a chain reference'
if (s is None) or (e is None):
s = -np.inf
e = np.inf
for l in mask_list:
if l is None:
continue
ch, idx = l
if (ch == ch_subset) and (idx >= s) and (idx <= e):
m_1d.append(True)
else:
m_1d.append(False)
return np.array(m_1d)
def mk_cce_and_hal_mask_2d(sampled_mask, con_to_set=None):
'''
Makes masks for ij pixels where the cce and hallucination loss should be applied.
Inputs
---------------
sampled_mask (str): String of where contigs should be applied. Ex: 3,B6-11,12,A12-19,9
cce_cutoff (float): Apply cce loss to cb-cb distances less than this value. Angstroms.
con_to_set (dict): Dictionary mapping the string of a contig (ex: 'B6-11') to an integer.
L_rec (int): Length of the receptor, if hallucinating in the context of the receptor.
Outputs
---------------
mask_cce (np.array, (L_hal, L_hal)): Boolean array. True where cce loss should be applied.
mask_hal (np.array, (L_hal, L_hal)): Boolean array. True where hallucination loss should be applied.
'''
if con_to_set is None:
con_to_set = mk_con_to_set(sampled_mask)
# Length of hallucinated protein
L_hal, L_max = mask_len(sampled_mask)
assert L_hal == L_max, 'A sampled mask must have gaps of a single length.'
# Map each contig to a 1D boolean mask
m_con = dict()
start_idx = 0
for l in sampled_mask.split(','):
if l[0].isalpha():
s, e = parse_range_string(l[1:])
L_con = e - s + 1
m = np.zeros(L_hal, dtype=bool)
m[start_idx:start_idx+L_con] = True
m_con[l] = m
start_idx += L_con
else:
L_gap, _ = parse_range_string(l)
start_idx += L_gap
# Combine contigs masks from each set to make 2D mask
mask_cce = np.zeros((L_hal, L_hal), dtype=bool)
for set_id in set(con_to_set.values()):
# gather all masks from contigs in the same set
masks = [m_con[k] for k,v in con_to_set.items() if v == set_id]
mask_1D = np.any(masks, axis=0)
update = mask_1D[:,None] * mask_1D[None,:]
mask_cce = np.any([mask_cce, update], axis=0)
# Make mask_hal
mask_hal = ~mask_cce
# Don't apply ANY losses on diagonal
mask_cce[np.arange(L_hal), np.arange(L_hal)] = False
mask_hal[np.arange(L_hal), np.arange(L_hal)] = False
# Don't apply ANY losses to receptor
m_1d_rec = mask_subset(sampled_mask, 'R')
m_2d_rec = m_1d_rec[:, None] * m_1d_rec[None, :]
mask_cce *= ~m_2d_rec
mask_hal *= ~m_2d_rec
return mask_cce, mask_hal
def apply_mask(mask, pdb_out):
'''
Uniformly samples gap lengths, then gathers the ref features
into the target hal features
Inputs
--------------
mask: specify the order and ranges of contigs and gaps
Contig - A continuous range of residues from the pdb.
Inclusive of the begining and end
Must start with the chain number
ex: B6-11
Gap - a gap length or a range of gaps lengths the
model is free to hallucinate
Gap ranges are inclusive of the end
ex: 9-21
ex - '3,B6-11,9-21,A36-42,20-30,A12-24,3-6'
pdb_out: dictionary from the prep_input function
Outputs
-------------
feat_hal: features from pdb_out scattered according to the sampled mask
mappings: dict keeping track of corresponding positions in the ref and hal features
'''
ref_pdb_2_idx0 = {pdb_idx:i for i, pdb_idx in enumerate(pdb_out['pdb_idx'])}
#1. make a map from hal_idx0 to ref_idx0. Has None for gap regions
hal_2_ref_idx0 = []
sampled_mask = []
for el in mask.split(','):
if el[0].isalpha(): # el is a contig
sampled_mask.append(el)
chain = el[0]
s,e = parse_range_string(el[1:])
for i in range(s, e+1):
idx0 = ref_pdb_2_idx0[(chain, i)]
hal_2_ref_idx0.append(idx0)
else: # el is a gap
# sample gap length
s,e = parse_range_string(el)
gap_len = np.random.randint(s, e+1)
hal_2_ref_idx0 += [None]*gap_len
sampled_mask.append(f'{gap_len}-{gap_len}')
#2. Convert mask to feat_hal and mappings
feat_hal, mappings = mk_feat_hal_and_mappings(hal_2_ref_idx0, pdb_out)
#3. Record the mask that was sampled
mappings['sampled_mask'] = ','.join(sampled_mask)
return feat_hal, mappings
def sample_mask(mask, pdb_idx):
'''
Uniformly samples gap lengths, then gathers the ref features
into the target hal features
Inputs
--------------
mask: specify the order and ranges of contigs and gaps
Contig - A continuous range of residues from the pdb.
Inclusive of the begining and end
Must start with the chain number
ex: B6-11
Gap - a gap length or a range of gaps lengths the
model is free to hallucinate
Gap ranges are inclusive of the end
ex: 9-21
ex - '3,B6-11,9-21,A36-42,20-30,A12-24,3-6'
Outputs
-------------
hal_2_ref_idx0: (list; length=L_hal)
List mapping hal_idx0 positions to ref_idx0 positions.
"None" used for indices that do not map to ref.
ex: [None, None, 3, 4, 5, None, None, None, 34, 35, 36]
sampled_mask: (str)
string of the sampled mask, so the transformations can be reapplied
ex - '3-3,B6-11,9-9,A36-42,20-20,A12-24,5-5'
'''
ref_pdb_2_idx0 = {pdb_i:i for i, pdb_i in enumerate(pdb_idx)}
#1. make a map from hal_idx0 to ref_idx0. Has None for gap regions
hal_2_ref_idx0 = []
sampled_mask = []
for el in mask.split(','):
if el[0].isalpha(): # el is a contig
sampled_mask.append(el)
chain = el[0]
s,e = parse_range_string(el[1:])
for i in range(s, e+1):
idx0 = ref_pdb_2_idx0[(chain, i)]
hal_2_ref_idx0.append(idx0)
else: # el is a gap
# sample gap length
s,e = parse_range_string(el)
gap_len = np.random.randint(s, e+1)
hal_2_ref_idx0 += [None]*gap_len
sampled_mask.append(f'{gap_len}-{gap_len}')
return hal_2_ref_idx0, sampled_mask
class GapResampler():
def __init__(self, use_bkg=True):
'''
'''
self.counts_passed = {} # dictionary for tallying counts of gap lengths for designs passing some threshold
self.counts_bkg = {}
self.use_bkg = use_bkg
def clean_mask(self, mask):
'''
Makes mask into a cononical form.
Ensures masks always alternate gap, contig and that
masks begin and end with a gap (even of length 0)
Input
-----------
masks: list of masks (str). Mask format: comma separted list
of alternating gap_length (int or int-int), contig.
Ex - 9,A12-19,15,B45-52 OR 9-9,A12-19,15-15,B45-52
Output
-----------
A canonicalized mask. Ex: N,9,A12-19,15,B45-52,0,C
'''
mask = mask.split(',')
mask_out = []
was_contig = True
was_gap = False
for i, el in enumerate(mask):
is_contig = el[0].isalpha()
is_gap = not is_contig
is_last = i == len(mask) - 1
# accepting gaps as either x-x or just x
if is_gap:
if '-' in el:
x1, x2 = el.split('-')
if x1 != x2:
print(f"Error: Gap must not be a range: {mask}")
return None
gap = x1
else:
gap = el
if is_contig:
contig = el
# gap -> contig: just append new contig
if (was_gap and is_contig):
mask_out.append(contig)
# contig -> gap: just append gap
elif (was_contig and is_gap):
mask_out.append(gap)
# contig -> contig: insert gap of 0, then add contig
elif (was_contig and is_contig):
mask_out.append('0')
mask_out.append(contig)
# gap -> gap: add them
elif (was_gap and is_gap):
combined_len = int(mask_out[-1]) + int(gap)
mask_out[-1] = str(combined_len)
# ensure last mask element is a gap
if (is_last and is_contig):
mask_out.append('0')
# update what previous element was
was_contig = el[0].isalpha()
was_gap = ~is_contig
# add 'N' and 'C' contigs
mask_out.insert(0, 'N')
mask_out.append('C')
return ','.join(mask_out)
def add_mask(self, mask, counting_dict):
'''
Adds counts of gap lengths to counting_dict
Inputs
-----------
masks: list of masks (str). Mask format: comma separted list
of alternating gap_length (int or int-int), contig.
Ex - 9,A12-19,15,B45-52 OR 9-9,A12-19,15-15,B45-52
'''
mask = self.clean_mask(mask)
mask = mask.split(',')
n_gaps = len(mask) // 2
# count occurances of contig,gap,contig triples
for i in range(n_gaps):
con1, gap, con2 = mask[2*i : 2*i+3]
# count gap length
if con1 in counting_dict:
if (gap, con2) in counting_dict[con1]:
counting_dict[con1][(gap, con2)] += 1
else:
counting_dict[con1][(gap, con2)] = 1
else:
counting_dict[con1] = {(gap, con2): 1}
def add_mask_pass(self, mask):
'''
Add a mask that passed to self.counts_passed
'''
self.add_mask(mask, self.counts_passed)
def add_mask_bkg(self, mask):
'''
Add a mask that passed to self.counts_bkg
'''
self.add_mask(mask, self.counts_bkg)
def get_enrichment(self):
'''
Calculate the ratio of counts_passed / count_bkg
Also notes all contigs
'''
if self.use_bkg is False:
print('Please pass in background masks and set self.use_bkg=True')
return
self.counts_enrich = copy.copy(self.counts_passed)
self.con_all = set()
for con1 in self.counts_enrich.keys():
self.con_all |= set([con1])
for gap, con2 in self.counts_enrich[con1].keys():
self.con_all |= set([con2])
bkg = self.counts_bkg[con1][(gap, con2)]
cnt = self.counts_passed[con1][(gap, con2)]
self.counts_enrich[con1][(gap, con2)] = cnt / bkg
def sample_mask(self):
'''
Sample a mask
'''
searching = True
while searching:
n_gaps = len(self.con_all) - 1
mask = ['N']
if self.use_bkg:
counts = self.counts_enrich
else:
counts = self.counts_passed
for i in range(n_gaps):
con_last = mask[-1]
# only allow jump to C as last option
if i == n_gaps - 1:
con_used = set(mask[::2])
else:
con_used = set(mask[::2]+['C'])
con_free = self.con_all - con_used
# get available "jumps" (con -> gap, con) you can make
jumps_all = counts[con_last]
jumps_free = {k:v for k,v in jumps_all.items() if k[1] in con_free}
if len(jumps_free) == 0:
print('No available jumps to continue the mask. Sampling again...')
else:
# normalize counts and sample move
mvs, cnt = zip(*jumps_free.items())
cnt = np.array(cnt)
prob = cnt / cnt.sum()
idx = np.random.choice(len(prob), p=prob)
mv = mvs[idx]
# add to the mask
mask.append(mv[0])
mask.append(mv[1])
# check that mask has the right number of elements
if len(mask) == 2*n_gaps + 1:
searching = False
else:
searching = True
return ','.join(mask[1:-1])
def gaps_as_ranges(self, mask):
'''
Convert gaps of a single int to ranges, for
backwards compatibility reasons
'''
mask_out = []
for el in mask.split(','):
if el[0].isalpha():
mask_out.append(el)
else:
mask_out.append(f'{el}-{el}')
return ','.join(mask_out)
def recover_mask(trb):
'''
Recover the string of the sampled mask given the trb file
'''
L_hal = trb['mask_contig'].shape[0]
mask = []
for idx0 in range(L_hal):
# what is the current idx
if idx0 in trb['con_hal_idx0']:
is_con = True
is_gap = False
else:
is_con = False
is_gap = True
# dealing with the first entry
if idx0 == 0:
if is_gap:
L_gap = 1
elif is_con:
ch, idx = trb['con_ref_pdb_idx'][ trb['con_hal_idx0'].tolist().index(idx0) ]
con_start = f'{ch}{idx}'
# take action based on what happend last time
else:
if (was_gap) and (is_gap):
L_gap +=1
#elif (was_con) and (is_con):
# continue
elif (was_gap) and (is_con):
# end gap
mask.append(f'{L_gap}-{L_gap}')
# start con
ch, idx = trb['con_ref_pdb_idx'][ trb['con_hal_idx0'].tolist().index(idx0) ]
con_start = f'{ch}{idx}'
elif (was_con) and (is_gap):
# end con
ch, idx = trb['con_ref_pdb_idx'][ trb['con_hal_idx0'].tolist().index(idx0) ]
mask.append(f'{con_start}-{idx}')
# start gap
L_gap = 1
# dealing with last entry
if idx0 == L_hal-1:
if is_gap:
mask.append(f'{L_gap}-{L_gap}')
elif is_con: # (edge case not handled: con starts and ends on last idx)
ch, idx = trb['con_ref_pdb_idx'][ trb['con_hal_idx0'].tolist().index(idx0-1) ]
mask.append(f'{con_start}-{idx}')
# update what last position was
was_con = copy.copy(is_con)
was_gap = copy.copy(is_gap)
return ','.join(mask)
def mask_len(mask):
'''
Calculate the min and max possible length that can
be sampled given a mask
'''
L_min = 0
L_max = 0
for el in mask.split(','):
if el[0].isalpha(): # is con
con_s, con_e = el[1:].split('-')
con_s, con_e = int(con_s), int(con_e)
L_con = con_e - con_s + 1
L_min += L_con
L_max += L_con
else: # is gap
if '-' in el:
gap_min, gap_max = el.split('-')
gap_min, gap_max = int(gap_min), int(gap_max)
L_min += gap_min
L_max += gap_max
else:
L_min += int(el)
L_max += int(el)
return L_min, L_max
class SampledMask():
def __init__(self, mask_str, ref_pdb_idx, con_to_set=None):
self.str = mask_str
self.L_hal = len(self)
self.L_ref = len(ref_pdb_idx)
#################
# con indices in hal and ref
#################
self.ref_pdb_idx = ref_pdb_idx
self.hal_pdb_idx = [('A', i) for i in range(1, len(self)+1)]
hal_idx0 = 0
con_ref_pdb_idx = []
con_hal_pdb_idx = []
con_ref_idx0 = []
con_hal_idx0 = []
for l in mask_str.split(','):
ch, s, e = SampledMask.parse_contig(l)
# contig
if ch:
for res in range(s, e+1):
con_ref_pdb_idx.append((ch, res))
con_hal_pdb_idx.append(('A', hal_idx0+1))
con_ref_idx0.append(self.ref_pdb_idx.index((ch, res)))
con_hal_idx0.append(hal_idx0)
hal_idx0 += 1
# gap
else:
for _ in range(s):
hal_idx0 += 1
self.con_mappings = {
'ref_pdb_idx': con_ref_pdb_idx,
'hal_pdb_idx': con_hal_pdb_idx,
'ref_idx0': con_ref_idx0,
'hal_idx0': con_hal_idx0,
}
#################
# con_to_set mapping
#################
if con_to_set:
self.con_to_set = con_to_set
else:
contigs = self.get_contigs()
self.con_to_set = dict(zip(contigs, len(contigs)*[0]))
# set_to_con mapping
set_to_con = {}
for k, v in self.con_to_set.items():
set_to_con[v] = set_to_con.get(v, []) + [k] # invert a dictionary with non-unique values
self.set_to_con = set_to_con
def __len__(self,):
_, L_max = self.mask_len(self.str)
return L_max
def map(self, sel, src, dst):
'''
Convert the contig selection in one indexing scheme to another.
Will return None if selection is not in a contig.
Input
----------
sel (str): selection of a contig range or idx0 range. Can take multiple comma separated values of same type. Ex: A5-10,B2-8 or 3-8,14-21
src (str): <'ref', 'hal'>
dst (str): <'ref_pdb_idx', 'hal_pdb_idx', 'ref_idx0', 'hal_idx0>
'''
out = []
for con in sel.split(','):
ch, s, e = SampledMask.parse_contig(con)
# selection type is pdb_idx
if ch:
src_long = f'{src}_pdb_idx'
mapping = dict(zip(self.con_mappings[src_long], self.con_mappings[dst]))
out += [mapping.get((ch, res)) for res in range(s, e+1)]
# selection type is idx0
else:
src_long = f'{src}_idx0'
mapping = dict(zip(self.con_mappings[src_long], self.con_mappings[dst]))
out += [mapping.get(i) for i in range(s, e+1)]
return out
@staticmethod
def expand(mask_str):
'''
Ex: '2,A3-5,3' --> [None, None, (A,3), (A,4), (A,5), None, None, None]
'''
expanded = []
for l in mask_str.split(','):
ch, s, e = SampledMask.parse_contig(l)
# contig
if ch:
expanded += [(ch, res) for res in range(s, e+1)]
# gap
else:
expanded += [None for _ in range(s)]
return expanded
@staticmethod
def contract(pdb_idx):
'''
Inverse of expand
Ex: [None, None, (A,3), (A,4), (A,5), None, None, None] --> '2,A3-5,3'
'''
contracted = []
l_prev = (None, -200)
first_el_written = False
for l_curr in pdb_idx:
if l_curr is None:
l_curr = (None, -100)
# extend gap
if l_curr == l_prev:
L_gap += 1
# extend con
elif l_curr == (l_prev[0], l_prev[1]+1):
con_e = l_curr[1]
# new gap
elif (l_curr != l_prev) and (l_curr[0] is None):
# write prev con
if 'con_ch' in locals():
contracted.append(f'{con_ch}{con_s}-{con_e}')
L_gap = 1
# new con
elif (l_curr != l_prev) and isinstance(l_curr[0], str):
# write prev con
if isinstance(l_prev[0], str) and ('con_ch' in locals()):
contracted.append(f'{con_ch}{con_s}-{con_e}')
# write prev gap
elif 'L_gap' in locals():
contracted.append(str(L_gap))
con_ch = l_curr[0]
con_s = l_curr[1]
con_e = l_curr[1]
# update l_prev
l_prev = l_curr
# write last element
if isinstance(l_prev[0], str) and ('con_ch' in locals()):
contracted.append(f'{con_ch}{con_s}-{con_e}')
elif 'L_gap' in locals():
contracted.append(str(L_gap))
return ','.join(contracted)
def subset(self, sub):
'''
Make a mask_str that is a subset of the original mask_str
Ex: self.mask_str = '2,A5-20,4', sub='A5-10' --> '2,A5-10,14'
'''
# map from hal_idx0 to ref_pdb_idx
hal_idx0 = self.map(sub, 'ref', 'hal_idx0')
ref_pdb_idx = SampledMask.expand(sub)
mapping = dict(zip(hal_idx0, ref_pdb_idx))
expanded = [mapping.get(idx0) for idx0 in range(len(self))]
return self.contract(expanded)
def mask_len(self, mask):
'''
Technically, can take both sampled and unsampled mask
'''
L_min = 0
L_max = 0
for l in self.str.split(','):
ch, s, e = SampledMask.parse_contig(l)
# contig
if ch:
L_min += e - s + 1
L_max += e - s + 1
# gap
else:
L_min += s
L_max += e
return L_min, L_max
def get_contigs(self, include_receptor=True):
'''
Get a list of all contigs in the mask
'''
[con for con in self.str.split(',') if SampledMask.parse_contig(con)[0]]
contigs = []
for con in self.str.split(','):
ch = SampledMask.parse_contig(con)[0]
if ch == 'R' and include_receptor == False:
continue
if ch:
contigs.append(con)
return contigs
def get_gaps(self,):
'''
Get a list of all gaps in the mask
'''
return [con for con in self.str.split(',') if SampledMask.parse_contig(con)[0] is None]
@staticmethod
def parse_range(_range):
if '-' in _range:
s, e = _range.split('-')
else:
s, e = _range, _range
return int(s), int(e)
@staticmethod
def parse_contig(contig):
'''
Return the chain, start and end residue in a contig or gap str.
Ex:
'A4-8' --> 'A', 4, 8
'A5' --> 'A', 5, 5
'4-8' --> None, 4, 8
'A' --> 'A', None, None
'''
# is contig
if contig[0].isalpha():
ch = contig[0]
if len(contig) > 1:
s, e = SampledMask.parse_range(contig[1:])
else:
s, e = None, None
# is gap
else:
ch = None
s, e = SampledMask.parse_range(contig)
return ch, s, e
def remove_diag(self, m_2d):
'''
Set the diagonal of a 2D boolean array to False
'''
L = m_2d.shape[0]
m_2d[np.arange(L), np.arange(L)] = False
return m_2d
def get_receptor_contig(self,):
'''
Returns None if there is no chain R in the mask_str
'''
receptor_contig = [l for l in self.get_contigs() if 'R' in l]
if len(receptor_contig) == 0:
receptor_contig = None
else:
receptor_contig = ','.join(receptor_contig)
return receptor_contig
def remove_receptor(self, m_2d):
'''
Remove intra-receptor contacts (chain R) from a mask
'''
receptor_contig = self.get_receptor_contig()
if receptor_contig: # has chain R
m_1d = np.zeros(self.L_hal, dtype=bool)
idx = np.array(self.map(receptor_contig, 'ref', 'hal_idx0'))
m_1d[idx] = True
update = m_1d[:, None] * m_1d[None, :]
m_2d = m_2d * ~update
return m_2d
def get_mask_con(self, include_receptor=False):
# Make a 2D boolean mask for each contig set
L = self.L_hal
mask_con = np.zeros([L, L], dtype=bool)
for set_id, contigs in self.set_to_con.items():
m_1d = np.zeros(L, dtype=bool)
for con in contigs:
idx = self.map(con, 'ref', 'hal_idx0')
idx = [l for l in idx if l != None]
idx = np.array(idx, dtype=int)
m_1d[idx] = True
update = m_1d[:, None] * m_1d[None, :]
mask_con = np.any([mask_con, update], axis=0)
# clean up
mask_con = self.remove_diag(mask_con)
if not include_receptor:
mask_con = self.remove_receptor(mask_con)
return mask_con
def get_mask_hal(self,):
mask_hal = ~self.get_mask_con()
mask_hal = self.remove_diag(mask_hal)
mask_hal = self.remove_receptor(mask_hal)
return mask_hal
def get_mask_cce(self, pdb, cce_cutoff=20., include_receptor=False):
'''
Remove ij pixels where contig distances are greater than cce_cutoff.
'''
# start with mask_con
mask_con = self.get_mask_con(include_receptor=include_receptor)
# get ref dists
xyz_ref = torch.tensor(pdb['xyz'][:,:3,:]).float()
c6d_ref = geometry.xyz_to_c6d(xyz_ref[None].permute(0,2,1,3),{'DMAX':20.0}).numpy()
dist = c6d_ref[0,:,:,0] # (L_ref, L_ref)
# scatter
dist_scattered = self.scatter_2d(dist)
# apply cce cuttoff
update = dist_scattered < cce_cutoff
mask_cce = np.all([mask_con, update], axis=0)
return mask_cce
def scatter_2d(self, ref_feat_2d):
'''
Inputs
---------
ref_feat_2d (np.array; (L_ref, L_ref, ...)): Features to be scattered. The first two leading dimensions must be equal to L_ref.
'''
assert ref_feat_2d.shape[:2] == (self.L_ref, self.L_ref), 'ERROR: feat_2d must have leading dimensions of (L_ref, L_ref)'
trailing_dims = ref_feat_2d.shape[2:]
dtype = ref_feat_2d.dtype
hal_feat_2d = np.zeros((self.L_hal, self.L_hal)+trailing_dims, dtype=dtype)
con_hal_idx0 = np.array(self.con_mappings['hal_idx0'])
ref_hal_idx0 = np.array(self.con_mappings['ref_idx0'])
hal_feat_2d[con_hal_idx0[:, None], con_hal_idx0[None, :]] = ref_feat_2d[ref_hal_idx0[:, None], ref_hal_idx0[None, :]]
return hal_feat_2d
def scatter_1d(self, ref_feat_1d):
'''
Inputs
---------
ref_feat_1d (np.array; (L_ref, ...)): Features to be scattered. The first leading dimension must be equal to L_ref.
'''
assert ref_feat_1d.shape[0] == self.L_ref, 'ERROR: feat_1d must have leading dimensions of (L_ref,)'
trailing_dims = ref_feat_1d.shape[1:]
dtype = ref_feat_1d.dtype
hal_feat_1d = np.zeros((self.L_hal,)+trailing_dims, dtype=dtype)
con_hal_idx0 = np.array(self.con_mappings['hal_idx0'])
ref_hal_idx0 = np.array(self.con_mappings['ref_idx0'])
hal_feat_1d[con_hal_idx0] = ref_feat_1d[ref_hal_idx0]
return hal_feat_1d
def idx_for_template(self, gap=200):
'''
Essentially return hal_idx0, except have a large jump for chain B,
to simulate a chain break. If B contains internal jumps in residue
numbering, these are preserved.
'''
is_rec = self.m1d_receptor()
resi_rec = np.array([idx[1] for idx in SampledMask.expand(self.str)
if idx is not None and idx[0]=='R'])
L_binder = sum(~is_rec)
if len(resi_rec)>0:
if is_rec[0]:
# receptor first
idx_tmpl = np.arange(resi_rec[-1]+gap+1, resi_rec[-1]+gap+1+L_binder)
idx_tmpl = np.concatenate([resi_rec, idx_tmpl])
else:
# receptor second
idx_tmpl = np.arange(L_binder)
if resi_rec[0] <= idx_tmpl[-1]+gap:
resi_rec += idx_tmpl[-1] - resi_rec[0] + gap + 1
idx_tmpl = np.concatenate([idx_tmpl, resi_rec])
else:
#when no receptor
idx_tmpl = np.arange(L_binder)
return idx_tmpl
def m1d_receptor(self,):
'''
Get a boolean array, True if the position corresponds to the receptor
'''
m1d = [(l is not None) and (l[0] == 'R') for l in SampledMask.expand(self.str)]
return np.array(m1d)
def erode(self, N_term=True, C_term=True):
'''
Reduce non-receptor contigs by 1 residue from the N and/or C terminus.
'''
x = SampledMask.expand(self.str)
if N_term:
for i, l in enumerate(x):
if (l is not None) and (l[0] != 'R'):
x[i] = None
break
if C_term:
x = x[::-1]
for i, l in enumerate(x):
if (l is not None) and (l[0] != 'R'):
x[i] = None
break
x = x[::-1]
self.str = self.contract(x)
return
def len_contigs(self, include_receptor=False):
con_str = ','.join(self.get_contigs(include_receptor))
return len(SampledMask.expand(con_str))
def make_template_features(pdb, args, device, hal_2_ref_idx0=None, sm_loss=None):
'''
Inputs
----------
sm_loss: Instance of a contig.SampledMask object used for making the loss masks.
'''
PARAMS = {
"DMIN" : 2.0,
"DMAX" : 20.0,
"DBINS" : 36,
"ABINS" : 36,
}
if args.use_template:
B,T = 1,1 # batch, templates
# spoof reference features
xyz_t = torch.tensor(pdb['xyz'][:, :3][None, None]) # (batch,templ,nres,3,3)
t0d = torch.ones((1,1,3)) # (batch, templ, 3)
t2d_ref = xyz_to_t2d(xyz_t=xyz_t, t0d=t0d, params=PARAMS) # (B,T,L,L,...)
L_ref = t2d_ref.shape[2]
#t1d_ref = torch.ones(size=(B,T,L_ref,3), dtype=torch.float32, device=device)
a = 2 * torch.ones([B,T,L_ref], dtype=torch.float32, device=device)
b = 0 * torch.ones([B,T,L_ref], dtype=torch.float32, device=device)
c = 1 * torch.ones([B,T,L_ref], dtype=torch.float32, device=device)
t1d_ref = torch.stack([a,b,c], axis=-1)
# Get the mask_str for scattering template features
#1. Template mask = sampled mask
if (args.use_template.lower() == 't') or (args.use_template.lower() == 'true'):
sm_tmpl = sm_loss
#2. Template mask is a subset of the sampled mask
else:
subset_contigs = args.use_template
if args.receptor:
receptor_contig = sm_loss.get_receptor_contig()
subset_contigs = ','.join([subset_contigs, receptor_contig])
mask_str_tmpl = sm_loss.subset(subset_contigs)
sm_tmpl = SampledMask(mask_str=mask_str_tmpl, ref_pdb_idx=pdb['pdb_idx'])
# scatter template features
# make leading dims (L,(L),...)
t1d_ref = t1d_ref.permute(2,3,0,1) # (L, ..., B, T)
t2d_ref = t2d_ref.permute(2,3,4,0,1) # (L, L, ..., B, T)
t1d_tmpl = sm_tmpl.scatter_1d(t1d_ref.cpu().numpy())
t2d_tmpl = sm_tmpl.scatter_2d(t2d_ref.cpu().numpy())
# update t2d_tmpl with mask_con (could update with mask_cce instead?)
mask_con = sm_tmpl.get_mask_con(include_receptor=True)
t2d_tmpl = (t2d_tmpl.T * mask_con.T).T # trick to broadcast arrays if leading dimensions match
t1d_tmpl = torch.tensor(t1d_tmpl, device=device)
t2d_tmpl = torch.tensor(t2d_tmpl, device=device)
# Permute B and T dims back to front
t1d_tmpl = t1d_tmpl.permute(2,3,0,1)
t2d_tmpl = t2d_tmpl.permute(3,4,0,1,2)
# Make last 3 idx of last dim all 1 to mimick Ivan's template feature
t2d_tmpl[..., -3:] = 1.
idx = torch.tensor(sm_tmpl.idx_for_template(gap=200), device=device)[None]
net_kwargs = {
'idx': idx,
't1d': t1d_tmpl,
't2d': t2d_tmpl
}
elif args.template_pdbs is not None:
B,T = 1, len(args.template_pdbs) # batch, templates
# get xyz features of all templates
xyz_t = [torch.tensor(parse_pdb(f_pdb)['xyz'][:, :3]) for f_pdb in args.template_pdbs]
xyz_t = torch.stack(xyz_t, axis=0)[None] # (batch, template, nres, 3, 3)
t0d = torch.ones(B,T,3)
t2d_tmpl = xyz_to_t2d(xyz_t=xyz_t, t0d=t0d, params=PARAMS).to(device) # (B,T,L,L,...)
L_tmpl = t2d_tmpl.shape[2]
t1d_tmpl = torch.ones(size=(B,T,L_tmpl,3), dtype=torch.float32, device=device)
# spoof pdb idx
idx_tmpl = torch.range(0, L_tmpl-1, dtype=torch.long, device=device)[None]
# Net() kwargs
net_kwargs = {
'idx': idx_tmpl,
't1d': t1d_tmpl,
't2d': t2d_tmpl
}
else:
net_kwargs = {}
return net_kwargs