rustev_core/
calibrate.rs

1//! Applying a calibration artifact (spec 002, 3.6).
2//!
3//! The only method is `temperature/1`: `softmax(z / T)` over logits, or over
4//! the natural logarithms of a distribution's masses (a zero mass stays
5//! zero). Applying it establishes the transformation and its binding, not
6//! empirical calibration. Fitting is out of scope.
7
8use std::collections::BTreeMap;
9
10use rustev_contract::calibration::{CalibrationArtifact, CalibrationMethod, CalibrationParameters};
11use rustev_contract::decimal::Decimal;
12use rustev_contract::ids::{ArtifactId, CalibrationId};
13
14use crate::kinds::{CalibratedProbability, KindError, softmax};
15
16/// What the step expects the artifact to be bound to.
17#[derive(Debug, Clone)]
18pub struct BindingCheck<'a> {
19    pub artifact: &'a ArtifactId,
20    pub task: &'a str,
21    pub question: &'a str,
22    pub options: &'a [String],
23}
24
25/// The supplied output being calibrated.
26#[derive(Debug, Clone, Copy)]
27pub enum CalibrationInput<'a> {
28    Logits(&'a BTreeMap<String, f64>),
29    Distribution(&'a BTreeMap<String, f64>),
30}
31
32/// Why no calibrated value was produced.
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub enum CalibrationError {
35    /// The artifact is bound to something other than this step.
36    Binding(String),
37    /// The temperature is outside `(0, 1000]`.
38    Parameters(String),
39    /// The input is not a valid logit vector or distribution.
40    Input(KindError),
41}
42
43impl std::fmt::Display for CalibrationError {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        match self {
46            CalibrationError::Binding(s) => write!(f, "calibration binding: {s}"),
47            CalibrationError::Parameters(s) => write!(f, "calibration parameters: {s}"),
48            CalibrationError::Input(e) => write!(f, "calibration input: {e}"),
49        }
50    }
51}
52
53impl std::error::Error for CalibrationError {}
54
55/// Check an artifact's binding against a step, without applying it. Used by
56/// the compiler (spec 002, 3.9.5 category 10).
57pub fn check_binding(
58    artifact: &CalibrationArtifact,
59    check: &BindingCheck<'_>,
60) -> Result<(), CalibrationError> {
61    check_static(artifact, check.task, check.question, check.options)?;
62    if &artifact.binding.artifact != check.artifact {
63        return Err(CalibrationError::Binding(format!(
64            "artifact {} is not the bound {}",
65            artifact.binding.artifact, check.artifact
66        )));
67    }
68    Ok(())
69}
70
71/// The part of the binding known before a backend is bound: task, question,
72/// options, method and parameters (compile phase c).
73pub fn check_static(
74    artifact: &CalibrationArtifact,
75    task: &str,
76    question: &str,
77    options: &[String],
78) -> Result<(), CalibrationError> {
79    let check = BindingCheck {
80        artifact: &artifact.binding.artifact,
81        task,
82        question,
83        options,
84    };
85    let b = &artifact.binding;
86    if b.task != check.task {
87        return Err(CalibrationError::Binding(format!(
88            "task {:?} is not {:?}",
89            b.task, check.task
90        )));
91    }
92    if b.question != check.question {
93        return Err(CalibrationError::Binding(format!(
94            "question {:?} is not {:?}",
95            b.question, check.question
96        )));
97    }
98    if artifact.options != check.options {
99        return Err(CalibrationError::Binding(
100            "options differ from the step's declared options".into(),
101        ));
102    }
103    match b.method {
104        CalibrationMethod::Temperature1 => {}
105    }
106    temperature(artifact).map(|_| ())
107}
108
109fn temperature(artifact: &CalibrationArtifact) -> Result<f64, CalibrationError> {
110    let CalibrationParameters::Temperature { temperature } = &artifact.parameters;
111    let max = Decimal::from_i64(1000);
112    if *temperature <= Decimal::ZERO || *temperature > max {
113        return Err(CalibrationError::Parameters(format!(
114            "temperature {temperature} outside (0, 1000]"
115        )));
116    }
117    Ok(temperature.to_f64_nearest())
118}
119
120/// Apply `artifact` (identified by `id`) to `input`, after checking that it
121/// is bound to this step. The only constructor of [`CalibratedProbability`].
122pub fn apply(
123    artifact: &CalibrationArtifact,
124    id: &CalibrationId,
125    check: &BindingCheck<'_>,
126    input: CalibrationInput<'_>,
127    tolerance: Decimal,
128) -> Result<CalibratedProbability, CalibrationError> {
129    check_binding(artifact, check)?;
130    let t = temperature(artifact)?;
131    let (map, is_logits) = match input {
132        CalibrationInput::Logits(m) => (m, true),
133        CalibrationInput::Distribution(m) => (m, false),
134    };
135    if let Some(extra) = map.keys().find(|k| !check.options.contains(k)) {
136        return Err(CalibrationError::Input(KindError(format!(
137            "undeclared option {extra:?}"
138        ))));
139    }
140    let mut z = Vec::with_capacity(check.options.len());
141    for o in check.options {
142        let v = *map
143            .get(o)
144            .ok_or_else(|| CalibrationError::Input(KindError(format!("missing option {o:?}"))))?;
145        if !v.is_finite() || (!is_logits && v < 0.0) {
146            return Err(CalibrationError::Input(KindError(format!(
147                "invalid value for option {o:?}"
148            ))));
149        }
150        z.push(if is_logits {
151            v
152        } else if v == 0.0 {
153            f64::NEG_INFINITY
154        } else {
155            crate::math::ln(v)
156        });
157    }
158    // Shift by the maximum before dividing, so a large logit and a small
159    // temperature cannot overflow: softmax is shift-invariant.
160    let max = z.iter().copied().fold(f64::NEG_INFINITY, f64::max);
161    let z: Vec<f64> = z
162        .iter()
163        .map(|v| {
164            if *v == f64::NEG_INFINITY {
165                *v
166            } else {
167                (v - max) / t
168            }
169        })
170        .collect();
171    if !is_logits {
172        // The input must itself be a distribution before it is recalibrated.
173        let sum: f64 = check.options.iter().fold(0.0, |acc, o| acc + map[o]);
174        if (sum - 1.0).abs() > tolerance.to_f64_nearest() {
175            return Err(CalibrationError::Input(KindError(format!(
176                "masses sum to {sum}"
177            ))));
178        }
179    }
180    if z.iter().all(|v| *v == f64::NEG_INFINITY) {
181        return Err(CalibrationError::Input(KindError(
182            "no option has mass".into(),
183        )));
184    }
185    let masses = softmax(&z);
186    CalibratedProbability::from_applied(
187        id.clone(),
188        check.options.to_vec(),
189        masses,
190        tolerance.to_f64_nearest(),
191    )
192    .map_err(CalibrationError::Input)
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use rustev_contract::calibration::CalibrationBinding;
199    use rustev_contract::ids::DatasetId;
200    use rustev_contract::{Identified, schema};
201
202    fn id(c: char) -> String {
203        format!("sha256:{}", c.to_string().repeat(64))
204    }
205
206    fn artifact(t: &str) -> CalibrationArtifact {
207        CalibrationArtifact {
208            schema: schema::CALIBRATION.into(),
209            binding: CalibrationBinding {
210                artifact: ArtifactId::parse(&id('a')).unwrap(),
211                task: "support.topic".into(),
212                question: "topic".into(),
213                dataset: DatasetId::parse(&id('d')).unwrap(),
214                method: CalibrationMethod::Temperature1,
215            },
216            options: vec!["a".into(), "b".into()],
217            parameters: CalibrationParameters::Temperature {
218                temperature: Decimal::parse(t).unwrap(),
219            },
220        }
221    }
222
223    fn tol() -> Decimal {
224        Decimal::parse("0.000001").unwrap()
225    }
226
227    #[test]
228    fn temperature_one_on_logits_is_softmax_and_higher_temperature_flattens() {
229        let art = artifact("1");
230        let cid = art.id().unwrap();
231        let aid = ArtifactId::parse(&id('a')).unwrap();
232        let opts = vec!["a".to_string(), "b".to_string()];
233        let check = BindingCheck {
234            artifact: &aid,
235            task: "support.topic",
236            question: "topic",
237            options: &opts,
238        };
239        let logits: BTreeMap<String, f64> = [("a".to_string(), 2.0), ("b".to_string(), 0.0)].into();
240        let p1 = apply(&art, &cid, &check, CalibrationInput::Logits(&logits), tol()).unwrap();
241        let hot = artifact("2");
242        let p2 = apply(
243            &hot,
244            &hot.id().unwrap(),
245            &check,
246            CalibrationInput::Logits(&logits),
247            tol(),
248        )
249        .unwrap();
250        assert!(p1.masses()[0] > p2.masses()[0]);
251        assert!(p2.masses()[0] > 0.5);
252        assert_eq!(p1.calibration(), &cid);
253    }
254
255    #[test]
256    fn a_distribution_input_keeps_zero_mass_at_zero() {
257        let art = artifact("0.5");
258        let aid = ArtifactId::parse(&id('a')).unwrap();
259        let opts = vec!["a".to_string(), "b".to_string()];
260        let check = BindingCheck {
261            artifact: &aid,
262            task: "support.topic",
263            question: "topic",
264            options: &opts,
265        };
266        let d: BTreeMap<String, f64> = [("a".to_string(), 1.0), ("b".to_string(), 0.0)].into();
267        let p = apply(
268            &art,
269            &art.id().unwrap(),
270            &check,
271            CalibrationInput::Distribution(&d),
272            tol(),
273        )
274        .unwrap();
275        assert_eq!(p.masses(), &[1.0, 0.0]);
276    }
277
278    #[test]
279    fn a_mismatched_binding_is_refused() {
280        let art = artifact("1");
281        let other = ArtifactId::parse(&id('b')).unwrap();
282        let aid = ArtifactId::parse(&id('a')).unwrap();
283        let opts = vec!["a".to_string(), "b".to_string()];
284        let reordered = vec!["b".to_string(), "a".to_string()];
285        for check in [
286            BindingCheck {
287                artifact: &other,
288                task: "support.topic",
289                question: "topic",
290                options: &opts,
291            },
292            BindingCheck {
293                artifact: &aid,
294                task: "other",
295                question: "topic",
296                options: &opts,
297            },
298            BindingCheck {
299                artifact: &aid,
300                task: "support.topic",
301                question: "other",
302                options: &opts,
303            },
304            BindingCheck {
305                artifact: &aid,
306                task: "support.topic",
307                question: "topic",
308                options: &reordered,
309            },
310        ] {
311            assert!(matches!(
312                check_binding(&art, &check),
313                Err(CalibrationError::Binding(_))
314            ));
315        }
316    }
317
318    #[test]
319    fn temperature_bounds_are_enforced() {
320        let aid = ArtifactId::parse(&id('a')).unwrap();
321        let opts = vec!["a".to_string(), "b".to_string()];
322        let check = BindingCheck {
323            artifact: &aid,
324            task: "support.topic",
325            question: "topic",
326            options: &opts,
327        };
328        assert!(check_binding(&artifact("1000"), &check).is_ok());
329        assert!(check_binding(&artifact("1000.000000001"), &check).is_err());
330        assert!(check_binding(&artifact("0"), &check).is_err());
331        assert!(check_binding(&artifact("-1"), &check).is_err());
332    }
333}
rustev_core/calibrate.rs

rustev_core/
calibrate.rs
