File size: 7,244 Bytes
94391f2 | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | import sys
import os
import json
import copy
import numpy as np
import scipy.stats as stats
import math
from rdkit.ML.Scoring.Scoring import CalcBEDROC, CalcAUC, CalcEnrichment
def cal_metrics(y_score, y_true):
# concate res_single and labels
scores = np.expand_dims(y_score, axis=1)
y_true = np.expand_dims(y_true, axis=1)
scores = np.concatenate((scores, y_true), axis=1)
# inverse sort scores based on first column
scores = scores[scores[:, 0].argsort()[::-1]]
bedroc = CalcBEDROC(scores, 1, 80.5)
count = 0
# sort y_score, return index
index = np.argsort(y_score)[::-1]
for i in range(int(len(index) * 0.005)):
if y_true[index[i]] == 1:
count += 1
auc = CalcAUC(scores, 1)
ef_list = CalcEnrichment(scores, 1, [0.005, 0.01, 0.02, 0.05])
return {
"BEDROC": bedroc,
"AUROC": auc,
"EF0.5": ef_list[0],
"EF1": ef_list[1],
"EF5": ef_list[3]
}
def print_avg_metric(metric_dict, name):
metric_lst = list(metric_dict.values())
ret_metric = copy.deepcopy(metric_lst[0])
for m in metric_lst[1:]:
for k in m:
ret_metric[k] += m[k]
for k in ret_metric:
ret_metric[k] = ret_metric[k] / len(metric_lst)
print(name, ret_metric)
def read_zeroshot_res(res_dir):
targets = sorted(list(os.listdir(res_dir)))
res_dict = {}
for target in targets:
real_dg = np.load(f"{res_dir}/{target}/saved_labels.npy")
if os.path.exists(f"{res_dir}/{target}/saved_preds.npy"):
pred_dg = np.load(f"{res_dir}/{target}/saved_preds.npy")
else:
mol_reps = np.load(f"{res_dir}/{target}/saved_mols_embed.npy")
pocket_reps = np.load(f"{res_dir}/{target}/saved_target_embed.npy")
res = pocket_reps @ mol_reps.T
pred_dg = res.max(axis=0)
res_dict[target] = {
"pred": pred_dg,
"exp": real_dg
}
return res_dict
def get_ensemble_res(res_list, begin=0, end=-1):
if end == -1:
end = len(res_list)
ret = copy.deepcopy(res_list[begin])
for res in res_list[begin+1:end]:
for k in ret.keys():
ret[k]["pred"] = np.array(ret[k]["pred"]) + np.array(res[k]["pred"])
for k in ret.keys():
ret[k]["pred"] = np.array(ret[k]["pred"]) / (end-begin)
return ret
def avg_metric(metric_lst_all):
ret_metric_dict = {}
for metric_lst in metric_lst_all:
ret_metric = copy.deepcopy(metric_lst[0])
for m in metric_lst[1:]:
for k in ["pearsonr", "spearmanr", "r2"]:
ret_metric[k] += m[k]
for k in ["spearmanr", "pearsonr", "r2"]:
ret_metric[k] = ret_metric[k] / len(metric_lst)
ret_metric_dict[ret_metric["target"]] = ret_metric
return ret_metric_dict
def get_metric(res):
metric_dict = {}
for k in sorted(list(res.keys())):
pred = res[k]["pred"]
exp = res[k]["exp"]
spearmanr = stats.spearmanr(exp, pred).statistic
pearsonr = stats.pearsonr(exp, pred).statistic
if math.isnan(pearsonr):
pearsonr = 0
if math.isnan(spearmanr):
spearmanr = 0
metric_dict[k] = {
"pearsonr":pearsonr,
"spearmanr":spearmanr,
"r2":max(pearsonr, 0)**2,
"target":k
}
return metric_dict
if __name__ == '__main__':
mode = sys.argv[1]
if mode == "zeroshot":
test_sets = sys.argv[2:]
for test_set in test_sets:
if test_set in ["DUDE", "PCBA", "DEKOIS"]:
metrics = {}
target_id_list = sorted(list(os.listdir(f"./result/pocket_ranking/{test_set}")))
for target_id in target_id_list:
lig_act = np.load(f"./result/pocket_ranking/{test_set}/{target_id}/saved_labels.npy")
score_1 = np.load(f"./result/pocket_ranking/{test_set}/{target_id}/GNN_res_epoch9.npy")
score_2 = np.load(f"./result/protein_ranking/{test_set}/{target_id}/GNN_res_epoch9.npy")
score = (score_1 + score_2)/2
metrics[target_id] = cal_metrics(score, lig_act)
json.dump(metrics, open(f"./result/pocket_ranking/{test_set}_metrics.json", "w"))
print_avg_metric(metrics, "Ours")
elif test_set in ["FEP"]:
target_id_list = sorted(list(os.listdir(f"./result/pocket_ranking/{test_set}")))
res_all_pocket, res_all_protein = [], []
for repeat in range(1, 6):
res_pocket = read_zeroshot_res(f"./result/pocket_ranking/{test_set}/repeat_{repeat}")
res_protein = read_zeroshot_res(f"./result/protein_ranking/{test_set}/repeat_{repeat}")
res_all_pocket.append(res_pocket)
res_all_protein.append(res_protein)
res_all_fusion = get_ensemble_res(res_all_pocket + res_all_protein)
metrics = get_metric(res_all_fusion)
json.dump(metrics, open(f"./result/pocket_ranking/{test_set}_metrics.json", "w"))
print_avg_metric(metrics, "Ours")
elif mode == "fewshot":
test_set = sys.argv[2]
support_num = sys.argv[3]
begin = 15
end = 20
metric_fusion_all = []
for seed in range(1, 11):
res_repeat_pocket = []
res_repeat_seq = []
if test_set in ["TIME", "OOD"]:
res_file_pocket = f"./result/pocket_ranking/{test_set}/random_{seed}_sup{support_num}.jsonl"
res_file_seq = f"./result/pocket_ranking/{test_set}/random_{seed}_sup{support_num}.jsonl"
if not os.path.exists(res_file_pocket):
continue
res_repeat_pocket = [json.loads(line) for line in open(res_file_pocket)][1:]
res_repeat_seq = [json.loads(line) for line in open(res_file_seq)][1:]
elif test_set in ["FEP_fewshot"]:
for repeat in range(1, 6):
res_file_pocket = f"./result/pocket_ranking/{test_set}/repeat_{repeat}/random_{seed}_sup{support_num}.jsonl"
res_file_seq = f"./result/pocket_ranking/{test_set}/repeat_{repeat}/random_{seed}_sup{support_num}.jsonl"
if not os.path.exists(res_file_pocket):
continue
res_pocket = [json.loads(line) for line in open(res_file_pocket)][1:]
res_seq = [json.loads(line) for line in open(res_file_seq)][1:]
res_pocket = get_ensemble_res(res_pocket, begin, end)
res_seq = get_ensemble_res(res_seq, begin, end)
res_repeat_pocket.append(res_pocket)
res_repeat_seq.append(res_seq)
res_repeat_fusion = get_ensemble_res(res_repeat_pocket + res_repeat_seq)
metric_fusion_all.append(get_metric(res_repeat_fusion))
metric_fusion_all = avg_metric(list(map(list, zip(*metric_fusion_all))))
json.dump(metric_fusion_all, open(f"./result/pocket_ranking/{test_set}_metrics.json", "w"))
print_avg_metric(metric_fusion_all, "Ours")
|