elinor 0.4.0

Evaluation Library in Information Retrieval
Documentation 
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

use std::collections::BTreeMap;

use crate::PredScore;
use crate::Relevance;
use crate::TrueScore;

#[derive(Clone, Copy, Debug)]
pub enum DcgWeighting {
    Jarvelin,
    Burges,
}

fn weighted_score(rel: TrueScore, weighting: DcgWeighting) -> f64 {
    match weighting {
        DcgWeighting::Jarvelin => rel as f64,
        DcgWeighting::Burges => 2.0_f64.powi(rel as i32) - 1.0,
    }
}

/// Computes the DCG at k.
pub fn compute_dcg<K>(
    trues: &BTreeMap<K, TrueScore>,
    sorted_preds: &[Relevance<K, PredScore>],
    k: usize,
    weighting: DcgWeighting,
) -> f64
where
    K: Eq + Ord,
{
    let k = if k == 0 { sorted_preds.len() } else { k };
    let mut dcg = 0.0;
    for (i, pred) in sorted_preds.iter().take(k).enumerate() {
        if let Some(&rel) = trues.get(&pred.doc_id) {
            dcg += weighted_score(rel, weighting) / (i as f64 + 2.0).log2();
        }
    }
    dcg
}

/// Computes the NDCG at k.
pub fn compute_ndcg<K>(
    trues: &BTreeMap<K, TrueScore>,
    sorted_trues: &[Relevance<K, TrueScore>],
    sorted_preds: &[Relevance<K, PredScore>],
    k: usize,
    weighting: DcgWeighting,
) -> f64
where
    K: Eq + Ord + Clone,
{
    let sorted_trues = sorted_trues
        .iter()
        .map(|r| Relevance {
            doc_id: r.doc_id.clone(),
            score: PredScore::from(r.score),
        })
        .collect::<Vec<_>>();
    let dcg = compute_dcg(trues, sorted_preds, k, weighting);
    let idcg = compute_dcg(trues, &sorted_trues, k, weighting);
    if idcg == 0.0 {
        1.0
    } else {
        dcg / idcg
    }
}