ranking metrics that account for imperfect recall of first pass ranker

ranking_metrics.py

def calc_ndcg(a, b, num_rel):
    ranking_len = len(a)
    len_diff = int(num_rel - a.sum())
    a = np.r_[a, np.ones(len_diff)]
    b = np.r_[b, (b.min() - 1) * np.ones(len_diff)]
    if (a.mean() < 1) and (a.mean() > 0): return ndcg_score([a], [b], k=ranking_len)
    else: return np.nan
    
def calc_auc(a, b, num_rel, population_size):
    assert a.sum() <= num_rel
    assert len(a) <= population_size
    assert num_rel <= population_size
    if (a.mean() < 1) and (a.mean() > 0):
        top_k_auc = roc_auc_score(a, b)
    else:
        top_k_auc = 1
    ranking_len = len(a)
    num_rel_top_k = a.sum()
    num_rel_rest = num_rel - num_rel_top_k
    rest_size = population_size - ranking_len
    num_irrel_top_k = ranking_len - num_rel_top_k
    num_irrel_rest = rest_size - num_rel_rest
    num_irrel = num_irrel_top_k + num_irrel_rest
    top_k_bottom_rest_auc = num_rel_top_k / ranking_len * (rest_size / num_irrel_rest)
    bottom_rest_top_k_auc = 0
    bottom_rest_auc = 1/2
    prob_top_k = num_rel_top_k / num_rel * num_irrel_top_k / num_irrel
    prob_top_k_bottom_rest = num_rel_top_k / num_rel * num_irrel_rest / num_irrel
    prob_bottom_rest_top_k = num_irrel_top_k / num_irrel * num_rel_rest / num_rel
    prob_bottom_rest = num_rel_rest / num_rel * num_irrel_rest / num_irrel
    return (top_k_auc * prob_top_k +
            bottom_rest_top_k_auc * prob_bottom_rest_top_k +
            np.nan_to_num(top_k_bottom_rest_auc) * prob_top_k_bottom_rest +
            bottom_rest_auc * prob_bottom_rest)