1use std::cmp::Ordering;
30use std::collections::BTreeMap;
31use std::fmt;
32
33use rustev_contract::decimal::Decimal;
34use rustev_contract::ids::{CalibrationId, PlanId};
35
36#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct KindError(pub String);
40
41impl fmt::Display for KindError {
42 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43 f.write_str(&self.0)
44 }
45}
46
47impl std::error::Error for KindError {}
48
49fn err<T>(msg: impl Into<String>) -> Result<T, KindError> {
50 Err(KindError(msg.into()))
51}
52
53fn masses_in_order(options: &[String], map: &BTreeMap<String, f64>) -> Result<Vec<f64>, KindError> {
55 if let Some(extra) = map.keys().find(|k| !options.contains(k)) {
56 return err(format!("undeclared option {extra:?}"));
57 }
58 options
59 .iter()
60 .map(|o| match map.get(o) {
61 Some(v) => Ok(*v),
62 None => err(format!("missing option {o:?}")),
63 })
64 .collect()
65}
66
67fn check_distribution(masses: &[f64], tolerance: f64) -> Result<(), KindError> {
68 if masses.is_empty() {
69 return err("empty distribution");
70 }
71 for (i, m) in masses.iter().enumerate() {
72 if !m.is_finite() {
73 return err(format!("non-finite mass at option {i}"));
74 }
75 if *m < 0.0 {
76 return err(format!("negative mass at option {i}"));
77 }
78 }
79 let sum: f64 = masses.iter().fold(0.0, |acc, m| acc + m);
81 if (sum - 1.0).abs() > tolerance {
82 return err(format!(
83 "masses sum to {sum}, outside tolerance {tolerance}"
84 ));
85 }
86 Ok(())
87}
88
89fn argmax(values: &[f64]) -> usize {
91 let mut best = 0;
92 for (i, v) in values.iter().enumerate().skip(1) {
93 if *v > values[best] {
94 best = i;
95 }
96 }
97 best
98}
99
100pub(crate) fn softmax(z: &[f64]) -> Vec<f64> {
104 let max = z.iter().copied().fold(f64::NEG_INFINITY, f64::max);
105 let e: Vec<f64> = z
106 .iter()
107 .map(|v| {
108 if *v == f64::NEG_INFINITY {
109 0.0
110 } else {
111 crate::math::exp(v - max)
112 }
113 })
114 .collect();
115 let sum: f64 = e.iter().fold(0.0, |acc, v| acc + v);
116 e.into_iter().map(|v| v / sum).collect()
117}
118
119#[derive(Debug, Clone, PartialEq)]
122pub struct Distribution {
123 options: Vec<String>,
124 masses: Vec<f64>,
125}
126
127impl Distribution {
128 pub fn new(
131 options: &[String],
132 map: &BTreeMap<String, f64>,
133 tolerance: Decimal,
134 ) -> Result<Self, KindError> {
135 let masses = masses_in_order(options, map)?;
136 check_distribution(&masses, tolerance.to_f64_nearest())?;
137 Ok(Self {
138 options: options.to_vec(),
139 masses,
140 })
141 }
142
143 pub fn from_logits(
146 options: &[String],
147 map: &BTreeMap<String, f64>,
148 tolerance: Decimal,
149 ) -> Result<Self, KindError> {
150 let z = masses_in_order(options, map)?;
151 if let Some(i) = z.iter().position(|v| !v.is_finite()) {
152 return err(format!("non-finite logit at option {i}"));
153 }
154 let masses = softmax(&z);
155 check_distribution(&masses, tolerance.to_f64_nearest())?;
156 Ok(Self {
157 options: options.to_vec(),
158 masses,
159 })
160 }
161
162 pub fn options(&self) -> &[String] {
163 &self.options
164 }
165
166 pub fn masses(&self) -> &[f64] {
167 &self.masses
168 }
169
170 pub fn mass(&self, option: &str) -> Option<f64> {
171 self.options
172 .iter()
173 .position(|o| o == option)
174 .map(|i| self.masses[i])
175 }
176
177 pub fn top(&self) -> (&str, f64) {
179 let i = argmax(&self.masses);
180 (&self.options[i], self.masses[i])
181 }
182}
183
184#[derive(Debug, Clone, PartialEq)]
188pub struct CalibratedProbability {
189 calibration: CalibrationId,
190 options: Vec<String>,
191 masses: Vec<f64>,
192}
193
194impl CalibratedProbability {
195 pub(crate) fn from_applied(
197 calibration: CalibrationId,
198 options: Vec<String>,
199 masses: Vec<f64>,
200 tolerance: f64,
201 ) -> Result<Self, KindError> {
202 check_distribution(&masses, tolerance)?;
203 Ok(Self {
204 calibration,
205 options,
206 masses,
207 })
208 }
209
210 pub fn calibration(&self) -> &CalibrationId {
211 &self.calibration
212 }
213
214 pub fn options(&self) -> &[String] {
215 &self.options
216 }
217
218 pub fn masses(&self) -> &[f64] {
219 &self.masses
220 }
221
222 pub fn mass(&self, option: &str) -> Option<f64> {
223 self.options
224 .iter()
225 .position(|o| o == option)
226 .map(|i| self.masses[i])
227 }
228
229 pub fn top(&self) -> (&str, f64) {
230 let i = argmax(&self.masses);
231 (&self.options[i], self.masses[i])
232 }
233}
234
235#[derive(Debug, Clone, PartialEq, Eq)]
237pub struct SelectedLabel {
238 options: Vec<String>,
239 index: usize,
240}
241
242impl SelectedLabel {
243 pub fn new(options: &[String], label: &str) -> Result<Self, KindError> {
244 match options.iter().position(|o| o == label) {
245 Some(index) => Ok(Self {
246 options: options.to_vec(),
247 index,
248 }),
249 None => err(format!("undeclared label {label:?}")),
250 }
251 }
252
253 pub fn label(&self) -> &str {
254 &self.options[self.index]
255 }
256
257 pub fn index(&self) -> usize {
258 self.index
259 }
260}
261
262#[derive(Debug, Clone, PartialEq)]
265pub struct OrdinalLevel {
266 levels: Vec<String>,
267 repr: OrdinalRepr,
268}
269
270#[derive(Debug, Clone, PartialEq)]
271enum OrdinalRepr {
272 Distribution(Vec<f64>),
273 Single(usize),
274}
275
276impl OrdinalLevel {
277 pub fn from_distribution(d: Distribution) -> Self {
278 Self {
279 levels: d.options,
280 repr: OrdinalRepr::Distribution(d.masses),
281 }
282 }
283
284 pub fn from_label(l: SelectedLabel) -> Self {
285 Self {
286 levels: l.options,
287 repr: OrdinalRepr::Single(l.index),
288 }
289 }
290
291 pub fn levels(&self) -> &[String] {
292 &self.levels
293 }
294
295 pub fn has_distribution(&self) -> bool {
296 matches!(self.repr, OrdinalRepr::Distribution(_))
297 }
298
299 pub fn expectation(&self) -> f64 {
302 match &self.repr {
303 OrdinalRepr::Distribution(m) => m
304 .iter()
305 .enumerate()
306 .fold(0.0, |acc, (i, p)| acc + (i as f64) * p),
307 OrdinalRepr::Single(i) => *i as f64,
308 }
309 }
310
311 pub fn top(&self) -> &str {
313 match &self.repr {
314 OrdinalRepr::Distribution(m) => &self.levels[argmax(m)],
315 OrdinalRepr::Single(i) => &self.levels[*i],
316 }
317 }
318}
319
320#[derive(Debug, Clone, PartialEq, Eq, Hash)]
322pub struct ScoreScope {
323 pub plan: PlanId,
324 pub step: String,
325 pub instance: Vec<String>,
326}
327
328#[derive(Debug, Clone, PartialEq)]
330pub struct ModelScore {
331 scope: ScoreScope,
332 candidate: String,
333 score: f64,
334}
335
336#[derive(Debug, Clone, PartialEq, Eq)]
338pub struct ScopeMismatch;
339
340impl ModelScore {
341 pub fn new(scope: ScoreScope, candidate: String, score: f64) -> Result<Self, KindError> {
342 if !score.is_finite() {
343 return err(format!("non-finite score for {candidate:?}"));
344 }
345 Ok(Self {
346 scope,
347 candidate,
348 score,
349 })
350 }
351
352 pub fn candidate(&self) -> &str {
353 &self.candidate
354 }
355
356 pub fn scope(&self) -> &ScoreScope {
357 &self.scope
358 }
359
360 pub fn cmp_within(&self, other: &ModelScore) -> Result<Ordering, ScopeMismatch> {
362 if self.scope != other.scope {
363 return Err(ScopeMismatch);
364 }
365 Ok(self.score.total_cmp(&other.score))
366 }
367}
368
369#[derive(Debug, Clone, PartialEq)]
372pub struct RankPosition {
373 pub(crate) entries: Vec<(String, f64, BTreeMap<String, f64>)>,
374}
375
376impl RankPosition {
377 pub fn order(&self) -> impl Iterator<Item = &str> {
379 self.entries.iter().map(|(c, _, _)| c.as_str())
380 }
381
382 pub fn len(&self) -> usize {
383 self.entries.len()
384 }
385
386 pub fn is_empty(&self) -> bool {
387 self.entries.is_empty()
388 }
389}
390
391#[cfg(test)]
392mod tests {
393 use super::*;
394
395 fn opts(v: &[&str]) -> Vec<String> {
396 v.iter().map(|s| s.to_string()).collect()
397 }
398
399 fn map(v: &[(&str, f64)]) -> BTreeMap<String, f64> {
400 v.iter().map(|(k, m)| (k.to_string(), *m)).collect()
401 }
402
403 fn tol() -> Decimal {
404 Decimal::parse("0.000001").unwrap()
405 }
406
407 #[test]
408 fn a_valid_distribution_is_accepted_as_supplied() {
409 let d = Distribution::new(&opts(&["a", "b"]), &map(&[("a", 0.25), ("b", 0.75)]), tol())
410 .unwrap();
411 assert_eq!(d.masses(), &[0.25, 0.75]);
412 assert_eq!(d.top(), ("b", 0.75));
413 }
414
415 #[test]
416 fn invalid_distributions_are_refused() {
417 let o = opts(&["a", "b"]);
418 for (bad, why) in [
419 (map(&[("a", f64::NAN), ("b", 1.0)]), "NaN"),
420 (map(&[("a", f64::INFINITY), ("b", 0.0)]), "infinite"),
421 (map(&[("a", -0.1), ("b", 1.1)]), "negative"),
422 (map(&[("a", 1.0)]), "missing key"),
423 (map(&[("a", 0.5), ("b", 0.5), ("c", 0.0)]), "undeclared key"),
424 (map(&[("a", 0.5), ("b", 0.6)]), "sum above tolerance"),
425 ] {
426 assert!(Distribution::new(&o, &bad, tol()).is_err(), "{why}");
427 }
428 assert!(Distribution::new(&o, &map(&[("a", 0.5), ("b", 0.5000005)]), tol()).is_ok());
430 assert!(Distribution::new(&o, &map(&[("a", 0.5), ("b", 0.500002)]), tol()).is_err());
431 }
432
433 #[test]
434 fn logits_are_normalized_and_non_finite_logits_refused() {
435 let o = opts(&["a", "b"]);
436 let d = Distribution::from_logits(&o, &map(&[("a", 0.0), ("b", 0.0)]), tol()).unwrap();
437 assert_eq!(d.masses(), &[0.5, 0.5]);
438 let d = Distribution::from_logits(&o, &map(&[("a", 1000.0), ("b", 999.0)]), tol()).unwrap();
440 assert!(d.masses()[0] > d.masses()[1]);
441 assert!(
442 Distribution::from_logits(&o, &map(&[("a", f64::NAN), ("b", 0.0)]), tol()).is_err()
443 );
444 }
445
446 #[test]
447 fn ties_go_to_the_earlier_declared_option() {
448 let d =
449 Distribution::new(&opts(&["z", "a"]), &map(&[("z", 0.5), ("a", 0.5)]), tol()).unwrap();
450 assert_eq!(d.top().0, "z");
451 }
452
453 #[test]
454 fn a_label_must_be_declared() {
455 assert!(SelectedLabel::new(&opts(&["a"]), "a").is_ok());
456 assert!(SelectedLabel::new(&opts(&["a"]), "b").is_err());
457 }
458
459 #[test]
460 fn ordinal_expectation_is_an_index_summary() {
461 let d = Distribution::new(
462 &opts(&["calm", "frustrated", "very_angry"]),
463 &map(&[("calm", 0.2), ("frustrated", 0.3), ("very_angry", 0.5)]),
464 tol(),
465 )
466 .unwrap();
467 let o = OrdinalLevel::from_distribution(d);
468 assert!((o.expectation() - 1.3).abs() < 1e-12);
469 let l =
470 OrdinalLevel::from_label(SelectedLabel::new(&opts(&["low", "high"]), "high").unwrap());
471 assert_eq!(l.expectation(), 1.0);
472 assert!(!l.has_distribution());
473 }
474
475 #[test]
476 fn scores_compare_only_within_their_scope() {
477 let plan = PlanId::parse(&format!("sha256:{}", "0".repeat(64))).unwrap();
478 let s1 = ScoreScope {
479 plan: plan.clone(),
480 step: "r".into(),
481 instance: vec![],
482 };
483 let s2 = ScoreScope {
484 plan,
485 step: "other".into(),
486 instance: vec![],
487 };
488 let a = ModelScore::new(s1.clone(), "x".into(), 1.0).unwrap();
489 let b = ModelScore::new(s1, "y".into(), 2.0).unwrap();
490 let c = ModelScore::new(s2, "x".into(), 3.0).unwrap();
491 assert_eq!(a.cmp_within(&b), Ok(Ordering::Less));
492 assert_eq!(a.cmp_within(&c), Err(ScopeMismatch));
493 assert!(ModelScore::new(a.scope().clone(), "z".into(), f64::NAN).is_err());
494 }
495}