ndcg@k computation

sample_ndcg.py

import math
import numpy as np
def ndcg(tot_probs, tot_labels,k=1): # will calculate for ndcg@1 by default
    tot_ndcg=[]
    for i in range(len(tot_probs)):
        probs = tot_probs[i]
        labels = tot_labels[i]
        # sort labels for the ideal dcg
        sor_labels_=[i for i,j in enumerate(labels) if j==1 ]
        # rank the documents # print(sor_labels_)
        sor_scores=[i[0] for i in sorted(enumerate(probs), key=lambda x:x[1],reverse=True)]
        # keep the labels for the first k documents
        sor_labels=[i for i in sor_labels_ if i in sor_scores[:k]]
        # compute ndcg@k
        dcg = 0
        for n,r in enumerate(sor_scores[:min(k,len(labels))]):
            dcg+=((2**labels[r])-1)/(math.log(n+2,2))
        perf_dcg = 0
        for n,r in enumerate(sorted_labels_[:k]):
            perf_dcg+=((2**labels[r])-1)/(math.log(n+2,2))
        ndcg=0.0
        if perf_dcg!=0:
            ndcg=dcg/perf_dcg
        tot_ndcg.append(ndcg)
    return np.mean(np.array(tot_ndcg))