PFresGO/predict.py at main · BioColLab/PFresGO · 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
import csv
import json
import pickle
import h5py
import argparse
import numpy as np
import tensorflow as tf


from pfresgo.PFresGO import PFresGO
from pfresgo.PFresGO_util import seq2onehot,load_GO_annot
from preprocessing.Seq2TFRecord import read_fasta

if __name__ == "__main__":

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-hsize', '--hidden_size', type=int, default=128, nargs='+', help="Hidden dimension of decoder layers.")
    parser.add_argument('-hlayer', '--num_hidden_layers', type=int, default=2, nargs='+', help="Number of hidden layers.")
    parser.add_argument('-numhead', '--num_heads', type=int, default=8, nargs='+', help="Number of attention heads.")
    parser.add_argument('-drop', '--dropout', type=float, default=0.3, help="Dropout rate.")
    parser.add_argument('-lr', type=float, default=0.0001, help="Initial learning rate.")
    parser.add_argument('-dff', type=int, default=1024, help="Dimension of ffn layer of decoder.")
    parser.add_argument('-e', '--epochs', type=int, default=100, help="Number of epochs to train.")
    parser.add_argument('-bs', '--batch_size', type=int, default=32, help="Batch size.")
    parser.add_argument('-pd', '--pad_len', type=int, default=1000, help="Padd length (max len of protein sequences in train set).")
    parser.add_argument('-ont', '--ontology', type=str, default='mf', choices=['mf', 'bp', 'cc'], help="Ontology.")
    parser.add_argument('--model_name', type=str, default='MF_PFresGO', help="Name of the model.")
    parser.add_argument('--autoencoder_name', type=str, default='./trained_model/Autoencoder_128.h5', help="Name of trained autoencoder model.")
    parser.add_argument('--train_tfrecord_fn', type=str,  default='./Datasets/TFRecords_sequences/PDB_GO_train', help='Train tfrecords.')
    parser.add_argument('--valid_tfrecord_fn', type=str, default='./Datasets/TFRecords_sequences/PDB_GO_valid',help='valid tfrecords.')
    parser.add_argument('--annot_fn', type=str, default='./Datasets/nrPDB-GO_2019.06.18_annot.tsv', help="File (*tsv) with GO term annotations.")
    parser.add_argument('--test_list', type=str, default='./Datasets/nrPDB-GO_2019.06.18_test.txt', help="File with test PDB chains.")
    parser.add_argument('--fasta_seq', type=str, default='./Datasets/nrPDB-GO_2019.06.18_sequences.fasta', help="File with fasta chains.")
    parser.add_argument('--label_embedding_dict', type=str, default='./Datasets/label-embedding-128.npy', help="File with go term embedding generated by Anc2vec.")
    parser.add_argument('--res_embeddings', type=str, default='./Datasets/per_residue_embeddings.h5',  help="File with sequences embeddings generated by ProtT5.")

    args = parser.parse_args()

    train_tfrecord_fn = args.train_tfrecord_fn + '*'
    valid_tfrecord_fn = args.valid_tfrecord_fn + '*'

    # load annotations
    prot2annot, goterms, gonames, counts = load_GO_annot(args.annot_fn)
    goterms = goterms[args.ontology]
    gonames = gonames[args.ontology]
    output_dim = len(goterms)

    go_emb = []
    noexist_go_emb = []
    go_id_dict = np.load(args.label_embedding_dict, allow_pickle=True)
    es = go_id_dict.item()
    for i in range(len(goterms)):
        if es.get(goterms[i]) is not None:
            go_embedding = es[goterms[i]]
            go_emb.append(go_embedding)
        else:
            noexist_go_emb.append(goterms[i])

    model = PFresGO(output_dim=output_dim, n_channels=26, lr=args.lr, drop=args.dropout, num_hidden_layers=args.num_hidden_layers, train=False, dff=args.dff,
                    model_name_prefix=args.model_name, label_embedding=go_emb, hidden_size=args.hidden_size, num_heads=args.num_heads, autoencoder_name=args.autoencoder_name)

    model.load_weights('./trained_model/'+args.model_name + '_best_train_model.h5')

    Y_pred = []
    Y_true = []
    proteins = []
    sequence = read_fasta(args.fasta_seq)

    with open(args.test_list, mode='r') as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')
        for row in csv_reader:
            prot = row[0]
            print("this protein is", prot)
            embed_index = str(prot) + ' ' + 'nrPDB'
            with h5py.File(args.res_embeddings, "r") as f:
                res_embed = f[embed_index][:]

            S = seq2onehot(str(sequence[prot]))
            padding = np.ones(len(S), np.float32)

            S = S.reshape(1, *S.shape)
            padding = padding.reshape(1, *padding.shape)
            res_embed = tf.expand_dims(res_embed, axis=0)

            # results
            proteins.append(prot)
            Y_pred.append(model.predict([S, padding, res_embed]).reshape(1, output_dim))
            Y_true.append(prot2annot[prot][args.ontology].reshape(1, output_dim))

    pickle.dump({'proteins': np.asarray(proteins),
                 'Y_pred': np.concatenate(Y_pred, axis=0),
                 'Y_true': np.concatenate(Y_true, axis=0),
                 'ontology': args.ontology,
                 'goterms': goterms,
                 'gonames': gonames},
                open(args.model_name + '_results.pckl', 'wb'))

    print("finish predicting")