///|
pub fn threshold_sweep(
qrels : Array[JudgedDoc],
run : Array[RetrievedDoc],
thresholds~ : Array[Int],
cutoffs~ : Array[Int],
) -> Array[ThresholdEvaluation] {
let result : Array[ThresholdEvaluation] = []
let normalized_cutoffs = normalize_cutoffs(cutoffs)
let normalized_thresholds = thresholds.map(fn(value) { Int::max(value, 0) })
let unique_thresholds : Array[Int] = []
for threshold in normalized_thresholds {
if !unique_thresholds.contains(threshold) {
unique_thresholds.push(threshold)
}
}
unique_thresholds.sort()
for threshold in unique_thresholds {
for cutoff in normalized_cutoffs {
result.push({
threshold,
cutoff,
precision: precision_at(qrels, run, cutoff, threshold),
recall: recall_at(qrels, run, cutoff, threshold),
f1: f1_at(qrels, run, cutoff, threshold),
ndcg: ndcg_at(qrels, run, cutoff, threshold, GainScheme::linear()),
})
}
}
result
}
///|
pub fn score_bucket_counts(
run : Array[RetrievedDoc],
bucket_count : Int,
) -> Array[Int] {
let count = Int::max(bucket_count, 0)
if count == 0 {
return []
}
let buckets : Array[Int] = []
for _ in 0.. String {
let lines : Array[String] = ["threshold,cutoff,precision,recall,f1,ndcg"]
for row in rows {
lines.push(
"\{row.threshold},\{row.cutoff},\{row.precision},\{row.recall},\{row.f1},\{row.ndcg}",
)
}
lines.join("\n")
}
///|
pub fn best_threshold(rows : Array[ThresholdEvaluation], cutoff : Int) -> Int {
let mut found = false
let mut best = 0
let mut best_score = -1.0
for row in rows {
if row.cutoff == cutoff && (!found || row.f1 > best_score) {
found = true
best = row.threshold
best_score = row.f1
}
}
best
}
///|
pub fn threshold_frontier(
rows : Array[ThresholdEvaluation],
) -> Array[ThresholdEvaluation] {
let frontier : Array[ThresholdEvaluation] = []
for row in rows {
let mut dominated = false
for other in rows {
if other.cutoff == row.cutoff &&
other.threshold != row.threshold &&
other.f1 >= row.f1 &&
other.ndcg >= row.ndcg &&
(other.f1 > row.f1 || other.ndcg > row.ndcg) {
dominated = true
}
}
if !dominated {
frontier.push(row)
}
}
frontier
}