File size: 5,522 Bytes
4523bda | 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 | import gc
import time
import ray
import numpy as np
from openea.modules.finding.similarity import sim
from openea.modules.utils.util import task_divide
def calculate_rank_(idx, sim_mat, top_k, accurate, total_num):
assert 1 in top_k
mr = 0
mrr = 0
hits = [0] * len(top_k)
hits1_rest = set()
for i in range(len(idx)):
gold = idx[i]
if accurate:
rank = (-sim_mat[i, :]).argsort()
else:
rank = np.argpartition(-sim_mat[i, :], np.array(top_k) - 1)
hits1_rest.add((gold, rank[0]))
assert gold in rank
rank_index = np.where(rank == gold)[0][0]
mr += (rank_index + 1)
mrr += 1 / (rank_index + 1)
for j in range(len(top_k)):
if rank_index < top_k[j]:
hits[j] += 1
mr /= total_num
mrr /= total_num
return mr, mrr, hits, hits1_rest
@ray.remote(num_cpus=1)
def calculate_rank(idx, sim_mat, top_k, accurate, total_num):
assert 1 in top_k
mr = 0
mrr = 0
hits = [0] * len(top_k)
hits1_rest = set()
for i in range(len(idx)):
gold = idx[i]
if accurate:
rank = (-sim_mat[i, :]).argsort()
else:
rank = np.argpartition(-sim_mat[i, :], np.array(top_k) - 1)
hits1_rest.add((gold, rank[0]))
assert gold in rank
rank_index = np.where(rank == gold)[0][0]
mr += (rank_index + 1)
mrr += 1 / (rank_index + 1)
for j in range(len(top_k)):
if rank_index < top_k[j]:
hits[j] += 1
mr /= total_num
mrr /= total_num
return mr, mrr, hits, hits1_rest
def greedy_alignment(embed1, embed2, top_k, nums_threads, metric, normalize, csls_k, accurate, is_print):
t = time.time()
sim_mat = sim(embed1, embed2, metric=metric, normalize=normalize, csls_k=csls_k)
num = sim_mat.shape[0]
if nums_threads > 1:
hits = [0] * len(top_k)
mr, mrr = 0, 0
alignment_rest = set()
rests = list()
search_tasks = task_divide(np.array(range(num)), nums_threads)
for task in search_tasks:
mat = sim_mat[task, :]
res = calculate_rank.remote(task, mat, top_k, accurate, num)
rests.append(res)
for rest in ray.get(rests):
sub_mr, sub_mrr, sub_hits, sub_hits1_rest = rest
mr += sub_mr
mrr += sub_mrr
hits += np.array(sub_hits)
alignment_rest |= sub_hits1_rest
else:
mr, mrr, hits, alignment_rest = calculate_rank_(list(range(num)), sim_mat, top_k, accurate, num)
assert len(alignment_rest) == num
hits = np.array(hits) / num
for i in range(len(hits)):
hits[i] = round(hits[i], 3)
cost = time.time() - t
if is_print:
if accurate:
if csls_k > 0:
print("accurate results with csls: csls={}, hits@{} = {}, mr = {:.3f}, mrr = {:.6f}, time = {:.3f} s ".
format(csls_k, top_k, hits, mr, mrr, cost))
else:
print("accurate results: hits@{} = {}, mr = {:.3f}, mrr = {:.6f}, time = {:.3f} s ".
format(top_k, hits, mr, mrr, cost))
else:
if csls_k > 0:
print("quick results with csls: csls={}, hits@{} = {}, time = {:.3f} s ".format(csls_k, top_k, hits,
cost))
else:
print("quick results: hits@{} = {}, time = {:.3f} s ".format(top_k, hits, cost))
sim_list = []
for i, j in alignment_rest:
sim_list.append(sim_mat[i, j])
del embed1
del embed2
del sim_mat
gc.collect()
return alignment_rest, hits, mr, mrr, sim_list
def valid(embeds1, embeds2, mapping, top_k, threads_num, metric='inner', normalize=False, csls_k=0, accurate=False,
is_print=True):
if mapping is None:
_, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(embeds1, embeds2, top_k, threads_num,
metric, normalize, csls_k, accurate, is_print)
else:
test_embeds1_mapped = np.matmul(embeds1, mapping)
_, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(test_embeds1_mapped, embeds2, top_k, threads_num,
metric, normalize, csls_k, accurate, is_print)
return hits1_12, mrr_12, sim_list
def test(embeds1, embeds2, mapping, top_k, threads_num, metric='inner', normalize=False, csls_k=0, accurate=True,
is_print=True):
if mapping is None:
alignment_rest_12, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(embeds1, embeds2, top_k, threads_num,
metric, normalize, csls_k, accurate,
is_print)
else:
test_embeds1_mapped = np.matmul(embeds1, mapping)
alignment_rest_12, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(test_embeds1_mapped, embeds2, top_k,
threads_num,
metric, normalize, csls_k, accurate,
is_print)
return alignment_rest_12, hits1_12, mrr_12, sim_list
|