|
|
import torch
|
|
|
|
|
|
def gather(x, indices):
|
|
|
indices = indices.view(-1, indices.shape[-1]).tolist()
|
|
|
out = torch.cat([x[i] for i in indices])
|
|
|
|
|
|
return out
|
|
|
|
|
|
def gather_nd(x, indices):
|
|
|
newshape = indices.shape[:-1] + x.shape[indices.shape[-1]:]
|
|
|
indices = indices.view(-1, indices.shape[-1]).tolist()
|
|
|
out = torch.cat([x[tuple(i)] for i in indices])
|
|
|
|
|
|
return out.reshape(newshape)
|
|
|
|
|
|
def gen_node_indices(size_list):
|
|
|
'''generate node index for extraction of nodes of each graph from batched data'''
|
|
|
node_num = []
|
|
|
node_range = []
|
|
|
size_list = [int(i) for i in size_list]
|
|
|
for i, n in enumerate(size_list):
|
|
|
node_num.extend([i]*n)
|
|
|
node_range.extend(list(range(n)))
|
|
|
|
|
|
node_num = torch.tensor(node_num)
|
|
|
node_range = torch.tensor(node_range)
|
|
|
indices = torch.stack([node_num, node_range], axis=1)
|
|
|
return indices, node_num, node_range
|
|
|
|
|
|
def segment_max(x, size_list):
|
|
|
size_list = [int(i) for i in size_list]
|
|
|
return torch.stack([torch.max(v, 0).values for v in torch.split(x, size_list)])
|
|
|
|
|
|
def segment_sum(x, size_list):
|
|
|
size_list = [int(i) for i in size_list]
|
|
|
return torch.stack([torch.sum(v, 0) for v in torch.split(x, size_list)])
|
|
|
|
|
|
def segment_softmax(gate, size_list):
|
|
|
segmax = segment_max(gate, size_list)
|
|
|
|
|
|
segmax_expand = torch.cat([segmax[i].repeat(n,1) for i,n in enumerate(size_list)], dim=0)
|
|
|
subtract = gate - segmax_expand
|
|
|
exp = torch.exp(subtract)
|
|
|
segsum = segment_sum(exp, size_list)
|
|
|
|
|
|
segsum_expand = torch.cat([segsum[i].repeat(n,1) for i,n in enumerate(size_list)], dim=0)
|
|
|
attention = exp / (segsum_expand + 1e-16)
|
|
|
|
|
|
return attention
|
|
|
|
|
|
def pad_V(V, max_n):
|
|
|
N, C = V.shape
|
|
|
if max_n > N:
|
|
|
zeros = torch.zeros(max_n-N, C)
|
|
|
V = torch.cat([V, zeros], dim=0)
|
|
|
return V
|
|
|
|
|
|
def pad_A(A, max_n):
|
|
|
N, L, _ = A.shape
|
|
|
if max_n > N:
|
|
|
zeros = torch.zeros(N, L, max_n-N)
|
|
|
A = torch.cat([A, zeros], dim=-1)
|
|
|
zeros = torch.zeros(max_n-N, L, max_n)
|
|
|
A = torch.cat([A, zeros], dim=0)
|
|
|
|
|
|
return A
|
|
|
|
|
|
def pad_prot(P, max_n):
|
|
|
N, = P.shape
|
|
|
if max_n > N:
|
|
|
zeros = torch.zeros(max_n-N)
|
|
|
P = torch.cat([P, zeros], dim=0)
|
|
|
|
|
|
return P.type(torch.IntTensor)
|
|
|
|
|
|
def create_batch(input, pad=False, device=torch.device('cpu')):
|
|
|
vl = []
|
|
|
al = []
|
|
|
gsl = []
|
|
|
msl = []
|
|
|
ssl = []
|
|
|
lbl = []
|
|
|
idxs = []
|
|
|
smis = []
|
|
|
|
|
|
for d in input:
|
|
|
vl.append(d['V'])
|
|
|
al.append(d['A'])
|
|
|
gsl.append(d['G'])
|
|
|
msl.append(d['mol_size'])
|
|
|
ssl.append(d['subgraph_size'])
|
|
|
lbl.append(d['label'])
|
|
|
idxs.append(d['index'])
|
|
|
smis.append(d['smiles'])
|
|
|
|
|
|
if gsl[0] is not None:
|
|
|
gsl = torch.stack(gsl, dim=0).to(device)
|
|
|
|
|
|
if pad:
|
|
|
max_n = max(map(lambda x:x.shape[0], vl))
|
|
|
vl1 = []
|
|
|
for v in vl:
|
|
|
vl1.append(pad_V(v, max_n))
|
|
|
al1 = []
|
|
|
for a in al:
|
|
|
al1.append(pad_A(a, max_n))
|
|
|
|
|
|
return {'V': torch.stack(vl1, dim=0).to(device),
|
|
|
'A': torch.stack(al1, dim=0).to(device),
|
|
|
'G': gsl,
|
|
|
'mol_size': torch.cat(msl, dim=0).to(device),
|
|
|
'subgraph_size': torch.stack(ssl, dim=0).to(device),
|
|
|
'label': torch.stack(lbl, dim=0).to(device),
|
|
|
'index': idxs,
|
|
|
'smiles': smis}
|
|
|
|
|
|
return {'V': torch.stack(vl, dim=0).to(device),
|
|
|
'A': torch.stack(al, dim=0).to(device),
|
|
|
'G': gsl,
|
|
|
'mol_size': torch.cat(msl, dim=0).to(device),
|
|
|
'subgraph_size': torch.stack(ssl, dim=0).to(device),
|
|
|
'label': torch.stack(lbl, dim=0).to(device),
|
|
|
'index': idxs,
|
|
|
'smiles': smis}
|
|
|
|
|
|
def create_mol_protein_batch(input, pad=False, device=torch.device('cpu'), pr=True):
|
|
|
vl = []
|
|
|
al = []
|
|
|
gsl = []
|
|
|
msl = []
|
|
|
ssl = []
|
|
|
prot = []
|
|
|
seq = []
|
|
|
lbl = []
|
|
|
idxs = []
|
|
|
smis = []
|
|
|
fpl = []
|
|
|
|
|
|
for d in input:
|
|
|
vl.append(d['V'])
|
|
|
al.append(d['A'])
|
|
|
gsl.append(d['G'])
|
|
|
msl.append(d['mol_size'])
|
|
|
ssl.append(d['subgraph_size'])
|
|
|
prot.append(d['protein_seq'])
|
|
|
seq.append(d['protein'])
|
|
|
lbl.append(d['label'])
|
|
|
idxs.append(d['index'])
|
|
|
smis.append(d['smiles'])
|
|
|
if 'fp' in d:
|
|
|
fpl.append(d['fp'])
|
|
|
|
|
|
if gsl[0] is not None:
|
|
|
if pad:
|
|
|
gsl = torch.stack(gsl, dim=0).to(device)
|
|
|
else:
|
|
|
gsl = [torch.unsqueeze(g, 0) for g in gsl]
|
|
|
|
|
|
if pad:
|
|
|
max_n = max(map(lambda x:x.shape[0], vl))
|
|
|
vl1 = []
|
|
|
if pr:
|
|
|
print('\tPadding V to max_n:', max_n)
|
|
|
for v in vl:
|
|
|
vl1.append(pad_V(v, max_n))
|
|
|
|
|
|
al1 = []
|
|
|
if pr:
|
|
|
print('\tPadding A to max_n:', max_n)
|
|
|
for a in al:
|
|
|
al1.append(pad_A(a, max_n))
|
|
|
|
|
|
max_prot = max(map(lambda x:x.shape[0], prot))
|
|
|
prot1 = []
|
|
|
if pr:
|
|
|
print('\tPadding protein_seq to max_n:', max_prot)
|
|
|
for p in prot:
|
|
|
prot1.append(pad_prot(p, max_prot))
|
|
|
|
|
|
fpt = None
|
|
|
if fpl:
|
|
|
fpt = torch.stack(fpl, dim=0).to(device)
|
|
|
|
|
|
return {'V': torch.stack(vl1, dim=0).to(device),
|
|
|
'A': torch.stack(al1, dim=0).to(device),
|
|
|
'G': gsl,
|
|
|
'fp': fpt,
|
|
|
'mol_size': torch.cat(msl, dim=0).to(device),
|
|
|
'subgraph_size': torch.stack(ssl, dim=0).to(device),
|
|
|
'protein_seq': torch.stack(prot1, dim=0).to(device),
|
|
|
'label': torch.stack(lbl, dim=0).view(-1).to(device),
|
|
|
'index': idxs,
|
|
|
'smiles': smis,
|
|
|
'protein': seq}
|
|
|
|
|
|
return {'V': [torch.unsqueeze(v, 0) for v in vl],
|
|
|
'A': [torch.unsqueeze(a, 0) for a in al],
|
|
|
'G': gsl,
|
|
|
'fp': fpt,
|
|
|
'mol_size': torch.cat(msl, dim=0).to(device),
|
|
|
'subgraph_size': [torch.unsqueeze(s, 0) for s in ssl],
|
|
|
'protein_seq': [torch.unsqueeze(p, 0) for p in prot],
|
|
|
'label': torch.stack(lbl, dim=0).view(-1).to(device),
|
|
|
'index': idxs,
|
|
|
'smiles': smis,
|
|
|
'protein': seq}
|
|
|
|
|
|
def create_mol_protein_fp_batch(input, pad=False, device=torch.device('cpu'), pr=True):
|
|
|
fp = []
|
|
|
prot = []
|
|
|
lbl = []
|
|
|
idxs = []
|
|
|
smis = []
|
|
|
|
|
|
for d in input:
|
|
|
fp.append(d['fp'])
|
|
|
prot.append(d['protein_seq'])
|
|
|
lbl.append(d['label'])
|
|
|
idxs.append(d['index'])
|
|
|
smis.append(d['smiles'])
|
|
|
|
|
|
if pad:
|
|
|
max_prot = max(map(lambda x:x.shape[0], prot))
|
|
|
prot1 = []
|
|
|
if pr:
|
|
|
print('\tPadding protein_seq to max_n:', max_prot)
|
|
|
for p in prot:
|
|
|
prot1.append(pad_prot(p, max_prot))
|
|
|
|
|
|
return {'fp': torch.stack(fp, dim=0).to(device),
|
|
|
'protein_seq': torch.stack(prot1, dim=0).to(device),
|
|
|
'label': torch.stack(lbl, dim=0).view(-1).to(device),
|
|
|
'index': idxs,
|
|
|
'smiles': smis}
|
|
|
|
|
|
return {'fp': [torch.unsqueeze(f, 0) for f in fp],
|
|
|
'protein_seq': [torch.unsqueeze(p, 0) for p in prot],
|
|
|
'label': torch.stack(lbl, dim=0).view(-1).to(device),
|
|
|
'index': idxs,
|
|
|
'smiles': smis} |
