-
Notifications
You must be signed in to change notification settings - Fork 1
/
RunningNDCG.py
53 lines (46 loc) · 2.08 KB
/
RunningNDCG.py
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
import numpy as np
import math
from pyltr.util.sort import get_sorted_y_positions
from pyltr.util.group import get_groups
class RunningNDCG(object):
def __init__(self, max_k):
self.max_k = max_k
self.discounts = np.array([1/math.log(i+2, 2) for i in range(max_k)])
def init(self, targets):
self.targets = np.array([math.pow(2, t)-1 for t in targets])
self.ideal_dcg = np.dot(self.discounts, np.array(sorted(self.targets, key=lambda r:-r)[0:self.max_k]))
if abs(self.ideal_dcg) < 1e-10:
self.ideal_dcg = 1
self.cur_dcg = np.dot(self.discounts, self.targets[0:self.max_k])
return self.cur_dcg / self.ideal_dcg
def swap_delta(self, i, j):
# Remove i&j contribution to DCG.
discount_i = 0 if i>=self.max_k else self.discounts[i]
discount_j = 0 if j>=self.max_k else self.discounts[j]
swapped_dcg_delta = (discount_j - discount_i) * (self.targets[i] - self.targets[j])
return swapped_dcg_delta / self.ideal_dcg
def calc(self, targets):
targets = np.array([math.pow(2, t)-1 for t in targets])
k = min(self.max_k, len(targets))
cur_dcg = np.dot(self.discounts[0:k], targets[0:k])
ideal_dcg = np.dot(self.discounts[0:k], np.array(sorted(targets, key=lambda r:-r)[0:k]))
if abs(ideal_dcg) < 1e-10:
ideal_dcg = 1
return cur_dcg /ideal_dcg
def calc_mean(self, qids, y, y_pred):
count = 0
sum = 0.
for qid, a, b in get_groups(qids):
currentRanking = list(get_sorted_y_positions(-y[a:b], y_pred[a:b]))
targetsForCurrentRanking = y[a:b][currentRanking]
sum += self.calc(targetsForCurrentRanking)
count += 1
return sum/count
def calc_mean_random(self, qids, y):
count = 0
sum = 0.
for qid, a, b in get_groups(qids):
targetsForRandomRanking = np.random.permutation(y[a:b])
sum += self.calc(targetsForRandomRanking)
count += 1
return sum/count