MultiKD-DTA/datahelper.py at master · NethanHu/MultiKD-DTA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import numpy as np
import json, pickle
from collections import OrderedDict
import esm
import torch
from rdkit import Chem


def generate_protein_pretraining_representation(dataset, prots):
    data_dict = {}
    prots_tuple = [(str(i), prots[i][:1022]) for i in range(len(prots))]
    model, alphabet = esm.pretrained.esm2_t30_150M_UR50D()
    batch_converter = alphabet.get_batch_converter()
    i = 0
    batch = 1

    while (batch * i) < len(prots):
        print('Converting protein batch: {}, batch * i = {}, length of protein = {}'.format(i, batch * i, len(prots)))
        if (i + batch) < len(prots):
            pt = prots_tuple[batch * i:batch * (i + 1)]
        else:
            pt = prots_tuple[batch * i:]

        batch_labels, batch_strs, batch_tokens = batch_converter(pt)
        with torch.no_grad():
            results = model(batch_tokens, repr_layers=[30], return_contacts=True)
        token_representations = results["representations"][30].numpy()
        data_dict[i] = token_representations
        i += 1
    np.savez(dataset + '.npz', dict=data_dict)


datasets = ['davis', 'kiba']
for dataset in datasets:
    fpath = 'data/' + dataset + '/'
    train_fold = json.load(open(fpath + "folds/train_fold_setting1.txt"))
    valid_fold = json.load(open(fpath + "folds/test_fold_setting1.txt"))
    folds = train_fold + [valid_fold]
    valid_ids = [5, 4, 3, 2, 1]
    valid_folds = [folds[vid] for vid in valid_ids]
    train_folds = []
    for i in range(5):
        temp = []
        for j in range(6):
            if j != valid_ids[i]:
                temp += folds[j]
        train_folds.append(temp)

    ligands = json.load(open(fpath + "ligands_can.txt"), object_pairs_hook=OrderedDict)
    proteins = json.load(open(fpath + "proteins.txt"), object_pairs_hook=OrderedDict)
    affinity = pickle.load(open(fpath + "Y", "rb"), encoding='latin1')
    drugs = []
    prots = []
    for d in ligands.keys():
        lg = Chem.MolToSmiles(Chem.MolFromSmiles(ligands[d]), isomericSmiles=True)
        drugs.append(lg)
    for t in proteins.keys():
        prots.append(proteins[t])
    if dataset == 'davis':
        affinity = [-np.log10(y / 1e9) for y in affinity]

    # protein pretraing presentation
    generate_protein_pretraining_representation(dataset, prots)

    affinity = np.asarray(affinity)
    opts = ['train', 'test']
    for i in range(5):
        train_fold = train_folds[i]
        valid_fold = valid_folds[i]
        for opt in opts:
            rows, cols = np.where(np.isnan(affinity) == False)
            if opt == 'train':
                rows, cols = rows[train_fold], cols[train_fold]
            elif opt == 'test':
                rows, cols = rows[valid_fold], cols[valid_fold]

            if i == 0:
                # generating standard data
                print('generating standard data')
                with open('data/' + dataset + '_' + opt + '.csv', 'w') as f:
                    f.write('compound_iso_smiles,target_sequence,affinity,protein_id\n')
                    for pair_ind in range(len(rows)):
                        ls = []
                        ls += [drugs[rows[pair_ind]]]
                        ls += [prots[cols[pair_ind]]]
                        ls += [affinity[rows[pair_ind], cols[pair_ind]]]
                        ls += [cols[pair_ind]]
                        f.write(','.join(map(str, ls)) + '\n')


            # 5-fold validation data
            print('generating 5-fold validation data')
            with open('data/' + dataset + '/' + dataset + '_' + opt + '_fold' + str(i) + '.csv', 'w') as f:
                f.write('compound_iso_smiles,target_sequence,affinity,protein_id\n')
                for pair_ind in range(len(rows)):
                    ls = []
                    ls += [drugs[rows[pair_ind]]]
                    ls += [prots[cols[pair_ind]]]
                    ls += [affinity[rows[pair_ind], cols[pair_ind]]]
                    ls += [cols[pair_ind]]
                    f.write(','.join(map(str, ls)) + '\n')