uqa_scoring/wand/
diagnostics.rs1use std::sync::Arc;
10
11use uqa_core::{DocId, Payload, PostingEntry, PostingList};
12
13use crate::error::invalid_input;
14use crate::scorer::Scorer;
15use crate::ScoringResult;
16
17#[derive(Debug, Default, Clone)]
20pub struct BoundTightnessAnalyzer {
21 pairs: Vec<(f64, f64)>,
22}
23
24impl BoundTightnessAnalyzer {
25 pub fn record(&mut self, upper_bound: f64, actual_max: f64) -> ScoringResult<()> {
26 if !upper_bound.is_finite() || upper_bound < 0.0 {
27 return Err(invalid_input(format!(
28 "upper bound must be finite and non-negative, got {upper_bound}"
29 )));
30 }
31 if !actual_max.is_finite() || actual_max < 0.0 {
32 return Err(invalid_input(format!(
33 "actual maximum must be finite and non-negative, got {actual_max}"
34 )));
35 }
36 if actual_max > upper_bound {
37 return Err(invalid_input(format!(
38 "actual maximum {actual_max} exceeds upper bound {upper_bound}"
39 )));
40 }
41 self.pairs.push((upper_bound, actual_max));
42 Ok(())
43 }
44
45 pub fn tightness_ratio(&self) -> f64 {
46 if self.pairs.is_empty() {
47 return 1.0;
48 }
49 let n = self.pairs.len() as f64;
50 let s: f64 = self
51 .pairs
52 .iter()
53 .map(|&(ub, am)| if ub > 0.0 { (am / ub).min(1.0) } else { 1.0 })
54 .sum();
55 s / n
56 }
57
58 pub fn slack(&self) -> f64 {
59 1.0 - self.tightness_ratio()
60 }
61
62 pub fn worst_bound_index(&self) -> usize {
63 self.pairs
64 .iter()
65 .enumerate()
66 .min_by(|(_, (ub_a, actual_a)), (_, (ub_b, actual_b))| {
67 let ratio_a = if *ub_a > 0.0 {
68 (*actual_a / *ub_a).min(1.0)
69 } else {
70 1.0
71 };
72 let ratio_b = if *ub_b > 0.0 {
73 (*actual_b / *ub_b).min(1.0)
74 } else {
75 1.0
76 };
77 ratio_a.total_cmp(&ratio_b)
78 })
79 .map_or(0, |(idx, _)| idx)
80 }
81
82 pub fn clear(&mut self) {
83 self.pairs.clear();
84 }
85}
86
87pub struct AdaptiveWANDScorer {
88 pub scorers: Vec<Arc<dyn Scorer>>,
89 pub k: usize,
90 pub posting_lists: Vec<PostingList>,
91 pub tightening_factor: f64,
92 pub analyzer: BoundTightnessAnalyzer,
93}
94
95impl AdaptiveWANDScorer {
96 pub fn new(
97 scorers: Vec<Arc<dyn Scorer>>,
98 k: usize,
99 posting_lists: Vec<PostingList>,
100 tightening_factor: f64,
101 ) -> ScoringResult<Self> {
102 validate_adaptive_inputs(&scorers, &posting_lists, tightening_factor)?;
103 Ok(Self {
104 scorers,
105 k,
106 posting_lists,
107 tightening_factor,
108 analyzer: BoundTightnessAnalyzer::default(),
109 })
110 }
111
112 pub fn compute_upper_bounds(&self) -> ScoringResult<Vec<f64>> {
113 validate_adaptive_inputs(&self.scorers, &self.posting_lists, self.tightening_factor)?;
114 self.scorers
115 .iter()
116 .zip(&self.posting_lists)
117 .map(|(scorer, pl)| {
118 let df = u64::try_from(pl.len())
119 .map_err(|_| invalid_input("posting-list length does not fit in u64"))?;
120 let bound = scorer.term_upper_bound(df) * self.tightening_factor;
121 if bound.is_finite() && bound >= 0.0 {
122 Ok(bound)
123 } else {
124 Err(invalid_input(format!(
125 "adaptive WAND bound must be finite and non-negative, got {bound}"
126 )))
127 }
128 })
129 .collect()
130 }
131
132 pub fn score_top_k(&mut self) -> ScoringResult<PostingList> {
133 validate_adaptive_inputs(&self.scorers, &self.posting_lists, self.tightening_factor)?;
134 self.analyzer.clear();
135 for (scorer, pl) in self.scorers.iter().zip(&self.posting_lists) {
136 let df = u64::try_from(pl.len())
137 .map_err(|_| invalid_input("posting-list length does not fit in u64"))?;
138 let upper = scorer.term_upper_bound(df);
139 let actual = pl
140 .iter()
141 .map(|entry| entry.payload.score)
142 .fold(0.0_f64, f64::max);
143 self.analyzer.record(upper, actual)?;
144 }
145
146 let mut scores: std::collections::BTreeMap<DocId, f64> = std::collections::BTreeMap::new();
147 for pl in &self.posting_lists {
148 for entry in pl {
149 let score = scores.entry(entry.doc_id).or_insert(0.0);
150 *score += entry.payload.score;
151 if !score.is_finite() || *score < 0.0 {
152 return Err(invalid_input(format!(
153 "adaptive WAND aggregate score must be finite and non-negative, got {score}"
154 )));
155 }
156 }
157 }
158 let mut entries: Vec<PostingEntry> = scores
159 .into_iter()
160 .map(|(doc_id, score)| PostingEntry::new(doc_id, Payload::with_score(score)))
161 .collect();
162 entries.sort_by(|a, b| {
163 b.payload
164 .score
165 .total_cmp(&a.payload.score)
166 .then_with(|| a.doc_id.cmp(&b.doc_id))
167 });
168 entries.truncate(self.k);
169 Ok(PostingList::from_unsorted(entries))
170 }
171}
172
173fn validate_adaptive_inputs(
174 scorers: &[Arc<dyn Scorer>],
175 posting_lists: &[PostingList],
176 tightening_factor: f64,
177) -> ScoringResult<()> {
178 if scorers.len() != posting_lists.len() {
179 return Err(invalid_input(format!(
180 "adaptive WAND requires one scorer per posting list, got {} scorers and {} lists",
181 scorers.len(),
182 posting_lists.len()
183 )));
184 }
185 if !tightening_factor.is_finite() || !(0.0..=1.0).contains(&tightening_factor) {
186 return Err(invalid_input(format!(
187 "adaptive WAND tightening factor must be finite and in [0, 1], got {tightening_factor}"
188 )));
189 }
190 for posting_list in posting_lists {
191 for entry in posting_list {
192 if !entry.payload.score.is_finite() || entry.payload.score < 0.0 {
193 return Err(invalid_input(format!(
194 "adaptive WAND input score must be finite and non-negative, got {} for document {}",
195 entry.payload.score, entry.doc_id
196 )));
197 }
198 }
199 }
200 Ok(())
201}