| 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} |
