metrics.py

import numpy as np
import unittest


def ndcg(score_list):
    """
    measure the ndcg score of order list
    https://en.wikipedia.org/wiki/Discounted_cumulative_gain
    parameter:
        score_list: np.array, shape=(sample_num,1)

    e.g. predict rank score list :
    >>> scores =  [3, 2, 3, 0, 1, 2]
    >>> ndcg_score = ndcg(scores)
    """

    def dcg(score_list):
        n = len(score_list)
        cost = .0
        for i in range(n):
            cost += float(np.power(2, score_list[i])) / np.log((i + 1) + 1)
        return cost

    dcg_cost = dcg(score_list)
    score_ranking = sorted(score_list, reverse=True)
    ideal_cost = dcg(score_ranking)
    return dcg_cost / ideal_cost


class TestNDCG(unittest.TestCase):
    def test_array(self):
        a = [3, 2, 3, 0, 1, 2]
        value = ndcg(a)
        self.assertAlmostEqual(0.9583, value, places=3)


if __name__ == '__main__':
    unittest.main()