Spaces:

merle
/

PROTEIN_GENERATOR

Runtime error

Jacob Gershon

new b

59a9ccf about 1 year ago

No virus

41.9 kB

	# 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[2i : 2i+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