///|
pub fn precision_from_counts(
relevant_retrieved : Int,
retrieved : Int,
) -> Double {
to_ratio(relevant_retrieved, retrieved)
}
///|
pub fn recall_from_counts(relevant_retrieved : Int, relevant : Int) -> Double {
to_ratio(relevant_retrieved, relevant)
}
///|
pub fn f1_from_counts(
relevant_retrieved : Int,
retrieved : Int,
relevant : Int,
) -> Double {
let precision = precision_from_counts(relevant_retrieved, retrieved)
let recall = recall_from_counts(relevant_retrieved, relevant)
if precision + recall == 0.0 {
0.0
} else {
2.0 * precision * recall / (precision + recall)
}
}
///|
fn sorted_run_items(run : Array[RetrievedDoc]) -> Array[RetrievedDoc] {
let copied = run.copy()
copied.sort_by(fn(a, b) {
let by_score = b.score.compare(a.score)
if by_score == 0 {
a.doc_id.compare(b.doc_id)
} else {
by_score
}
})
copied
}
///|
fn top_relevances(
qrels : Array[JudgedDoc],
run : Array[RetrievedDoc],
cutoff : Int,
threshold : Int,
) -> (Array[Int], Int, Int) {
let relevance_map = build_relevance_map(qrels)
let mut relevant_total = 0
for _, relevance in relevance_map {
if relevance >= threshold {
relevant_total += 1
}
}
let ranked = sorted_run_items(run)
let limit = Int::min(Int::max(cutoff, 0), ranked.length())
let relevances : Array[Int] = []
for index in 0.. Double {
let (relevances, relevant_total, _) = top_relevances(
qrels, run, cutoff, threshold,
)
let mut hits = 0
for relevance in relevances {
if relevance >= threshold {
hits += 1
}
}
recall_from_counts(hits, relevant_total)
}
///|
pub fn precision_at(
qrels : Array[JudgedDoc],
run : Array[RetrievedDoc],
cutoff : Int,
threshold : Int,
) -> Double {
let (relevances, _, _) = top_relevances(qrels, run, cutoff, threshold)
let mut hits = 0
for relevance in relevances {
if relevance >= threshold {
hits += 1
}
}
precision_from_counts(hits, Int::max(cutoff, 0))
}
///|
pub fn f1_at(
qrels : Array[JudgedDoc],
run : Array[RetrievedDoc],
cutoff : Int,
threshold : Int,
) -> Double {
let (relevances, relevant_total, _) = top_relevances(
qrels, run, cutoff, threshold,
)
let mut hits = 0
for relevance in relevances {
if relevance >= threshold {
hits += 1
}
}
f1_from_counts(hits, Int::max(cutoff, 0), relevant_total)
}
///|
pub fn r_precision(
qrels : Array[JudgedDoc],
run : Array[RetrievedDoc],
threshold : Int,
) -> Double {
let relevance_map = build_relevance_map(qrels)
let mut relevant_total = 0
for _, relevance in relevance_map {
if relevance >= threshold {
relevant_total += 1
}
}
if relevant_total <= 0 {
return 0.0
}
precision_at(qrels, run, relevant_total, threshold)
}
///|
pub fn fallout_at(
qrels : Array[JudgedDoc],
run : Array[RetrievedDoc],
cutoff : Int,
threshold : Int,
) -> Double {
let relevance_map = build_relevance_map(qrels)
let mut non_relevant_total = 0
for _, relevance in relevance_map {
if relevance < threshold {
non_relevant_total += 1
}
}
if non_relevant_total <= 0 {
return 0.0
}
// The relevance array alone cannot distinguish an unjudged document from
// a judged non-relevant document; recompute from ranked IDs below.
let ranked = sorted_run_items(run)
let limit = Int::min(Int::max(cutoff, 0), ranked.length())
let mut false_positives = 0
for index in 0.. Map[String, Double] {
let values : Map[String, Double] = Map([])
values["precision@\{cutoff}"] = precision_at(qrels, run, cutoff, threshold)
values["recall@\{cutoff}"] = recall_at(qrels, run, cutoff, threshold)
values["f1@\{cutoff}"] = f1_at(qrels, run, cutoff, threshold)
values["r_precision@\{cutoff}"] = r_precision(qrels, run, threshold)
values["bpref@\{cutoff}"] = bpref_at(qrels, run, cutoff, threshold)
values["err@\{cutoff}"] = err_at(qrels, run, cutoff, threshold)
values["rbp@\{cutoff}"] = rbp_at(qrels, run, cutoff, threshold)
values["fallout@\{cutoff}"] = fallout_at(qrels, run, cutoff, threshold)
values["graded_precision@\{cutoff}"] = graded_precision_at(
qrels, run, cutoff, threshold, gain_scheme,
)
values
}