Skip to main content

rustev_core/
kinds.rs

1//! Value kinds (spec 002, 3.5). Each is its own type with private fields and
2//! fallible construction; there is no `From` or `Into` between any two, so a
3//! score never becomes a probability and a label never acquires mass.
4//!
5//! ```compile_fail
6//! # use rustev_core::kinds::*;
7//! fn launder(label: SelectedLabel) -> Distribution { label.into() }
8//! ```
9//!
10//! ```compile_fail
11//! # use rustev_core::kinds::*;
12//! fn launder(d: Distribution) -> CalibratedProbability { d.into() }
13//! ```
14//!
15//! ```compile_fail
16//! # use rustev_core::kinds::*;
17//! fn launder(s: ModelScore) -> Distribution { Distribution::from(s) }
18//! ```
19//!
20//! ```compile_fail
21//! # use rustev_core::kinds::*;
22//! # use rustev_contract::ids::CalibrationId;
23//! // A calibrated probability cannot be built without applying an artifact.
24//! fn forge(id: CalibrationId) -> CalibratedProbability {
25//!     CalibratedProbability { calibration: id, options: vec![], masses: vec![] }
26//! }
27//! ```
28
29use std::cmp::Ordering;
30use std::collections::BTreeMap;
31use std::fmt;
32
33use rustev_contract::decimal::Decimal;
34use rustev_contract::ids::{CalibrationId, PlanId};
35
36/// Why a supplied value is not a valid value of its kind. Becomes
37/// `invalid_backend_output` at the backend boundary.
38#[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
53/// Masses keyed exactly by `options`, in declared order.
54fn 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    // Summed in declared option order (spec 002, 3.5.2).
80    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
89/// Index of the largest value; ties go to the earliest (declared order).
90fn 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
100/// Softmax with max subtraction, in declared order. Uses `exp`, so it is
101/// outside exact computation (spec 002, 3.4.6), but it computes the same
102/// bits on every target (spec 047).
103pub(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/// Non-negative mass over declared options, summing to 1 within tolerance.
120/// Uncalibrated: it is not a calibrated probability.
121#[derive(Debug, Clone, PartialEq)]
122pub struct Distribution {
123    options: Vec<String>,
124    masses: Vec<f64>,
125}
126
127impl Distribution {
128    /// Validate a supplied distribution. Stored as supplied, never
129    /// renormalized.
130    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    /// Normalize logits by softmax (design 8.2: the core normalizes logits,
144    /// it never invents mass for a label backend).
145    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    /// The top option; ties go to the earlier declared option.
178    pub fn top(&self) -> (&str, f64) {
179        let i = argmax(&self.masses);
180        (&self.options[i], self.masses[i])
181    }
182}
183
184/// A distribution passed through a calibration artifact bound to this step.
185/// It records which transformation was applied; it does not establish that
186/// the result is calibrated on current data (spec 002, 3.6.2).
187#[derive(Debug, Clone, PartialEq)]
188pub struct CalibratedProbability {
189    calibration: CalibrationId,
190    options: Vec<String>,
191    masses: Vec<f64>,
192}
193
194impl CalibratedProbability {
195    /// Only [`crate::calibrate::apply`] calls this, after checking the binding.
196    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/// One option, with no mass and no uncertainty measure.
236#[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/// A level on a declared rubric, low to high, as a distribution over levels
263/// or a single level. Its expectation is an index summary only.
264#[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    /// `sum(index * mass)` in level order, or the single level's index. An
300    /// index summary, not an interval quantity.
301    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    /// The most likely level (ties to the lower level), or the single level.
312    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/// Where a model score may be compared: one plan, step and request.
321#[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/// An unnormalized score, comparable only within its scope.
329#[derive(Debug, Clone, PartialEq)]
330pub struct ModelScore {
331    scope: ScoreScope,
332    candidate: String,
333    score: f64,
334}
335
336/// Two scores from different scopes were compared.
337#[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    /// Order two scores of the same scope; refused across scopes.
361    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/// An order over candidates with the scores and components that produced it;
370/// not a probability of being best.
371#[derive(Debug, Clone, PartialEq)]
372pub struct RankPosition {
373    pub(crate) entries: Vec<(String, f64, BTreeMap<String, f64>)>,
374}
375
376impl RankPosition {
377    /// Candidates in rank order.
378    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        // Negative control at the tolerance edge: 1 + 5e-7 passes, 1 + 2e-6 fails.
429        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        // Large logits do not overflow thanks to max subtraction.
439        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}