def test_gap_mask2(self):
s = ('222222222222222222.11112222222222222222222222222'
'222222222222222222222222222222222222222222222222'
'22222222...::::::..:2:22::2:::::::..11.111...::.'
'::::::::::.::::......:::::::::::222:.::::::::.11'
'.:::::::::.:22.::::::::::::2:::::::::::::::1::..'
'.::::::::::::::::::::::22:2:2::::::::::1::::::::'
'::::22222::::::::::1::::::.')
# N, M = 197, 283
gap_mask(s)
def test_gap_mask(self):
s = ":11::22:"
res = gap_mask(s)
exp = np.array([[1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 0, 1]])
npt.assert_equal(res, exp)
s = ":11:.:22:"
res = gap_mask(s)
exp = np.array([[1, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1]])
npt.assert_equal(res, exp)
def __getitem__(self, i):
""" Gets alignment pair.
Parameters
----------
i : int
Index of item
Returns
-------
gene : torch.Tensor
Encoded representation of protein of interest
pos : torch.Tensor
Encoded representation of protein that aligns with `gene`.
states : torch.Tensor
Alignment string
alignment_matrix : torch.Tensor
Ground truth alignment matrix
path_matrix : torch.Tensor
Pairwise path distances, where the smallest distance
to the path is computed for every element in the matrix.
"""
gene = self.pairs.iloc[i]['chain1']
pos = self.pairs.iloc[i]['chain2']
st = self.pairs.iloc[i]['alignment']
states = list(map(tmstate_f, st))
if self.clip_ends:
gene, pos, states, st = clip_boundaries(gene, pos, states, st)
if self.pad_ends:
states = [m] + states + [m]
states = torch.Tensor(states).long()
gene = self.tokenizer(str.encode(gene))
pos = self.tokenizer(str.encode(pos))
gene = torch.Tensor(gene).long()
pos = torch.Tensor(pos).long()
alignment_matrix = torch.from_numpy(states2matrix(states))
path_matrix = torch.empty(*alignment_matrix.shape)
g_mask = torch.ones(*alignment_matrix.shape)
if self.construct_paths:
pi = states2edges(states)
path_matrix = torch.from_numpy(path_distance_matrix(pi))
path_matrix = reshape(path_matrix, len(gene), len(pos))
if self.mask_gaps:
g_mask = torch.from_numpy(gap_mask(st)).bool()
alignment_matrix = reshape(alignment_matrix, len(gene), len(pos))
g_mask = reshape(g_mask, len(gene), len(pos))
if not self.return_names:
return gene, pos, states, alignment_matrix, path_matrix, g_mask
else:
gene_name = self.pairs.iloc[i]['chain1_name']
pos_name = self.pairs.iloc[i]['chain2_name']
return (gene, pos, states, alignment_matrix, path_matrix, g_mask,
gene_name, pos_name)