Skip to main content

prospicio_models/
artifact.rs

1//! The envelope every model artifact shares (`docs/design/models.md`,
2//! "Model artifacts"): the format check and version rule, provenance, and
3//! the JSON helpers the engines' readers and writers use. Non-finite
4//! numbers are written as the strings `"NaN"`, `"inf"` and `"-inf"`; with
5//! the workspace's `float_roundtrip` every other number round-trips bit
6//! for bit.
7
8use prospicio_core::{Error, Result};
9use serde_json::{Map, Number, Value, json};
10
11use crate::{Family, Link};
12
13/// Provenance: the model, the writing crate and its version, and the hash
14/// of the training data.
15pub fn provenance(model: &str, krate: &str, version: &str, input_hash: &str) -> Value {
16    json!({
17        "model": model,
18        "versions": [[krate, version]],
19        "input_hash": input_hash,
20    })
21}
22
23/// Parses `text` and checks its format and version (at most `newest`).
24pub fn parse(text: &str, format: &str, newest: u64) -> Result<Value> {
25    let doc: Value = serde_json::from_str(text)
26        .map_err(|e| Error::Data(format!("model artifact is not JSON: {e}")))?;
27    let top = object(&doc, "artifact")?;
28    if top.get("format").and_then(Value::as_str) != Some(format) {
29        return Err(Error::Data(format!("not a {format} artifact")));
30    }
31    let version = uint(top, "format_version")?;
32    if version == 0 || version > newest {
33        return Err(Error::Data(format!(
34            "artifact format version {version}; this build reads 1 to {newest}"
35        )));
36    }
37    Ok(doc)
38}
39
40/// The provenance's `input_hash`.
41pub fn input_hash(top: &Map<String, Value>) -> Result<String> {
42    let provenance = object(field(top, "provenance")?, "provenance")?;
43    Ok(field(provenance, "input_hash")?
44        .as_str()
45        .ok_or_else(|| bad("provenance.input_hash"))?
46        .to_string())
47}
48
49/// The array of strings at `key`.
50pub fn strings(o: &Map<String, Value>, key: &str) -> Result<Vec<String>> {
51    array(o, key)?
52        .iter()
53        .map(|v| v.as_str().map(String::from).ok_or_else(|| bad(key)))
54        .collect()
55}
56
57/// The array at `key`.
58pub fn array<'a>(o: &'a Map<String, Value>, key: &str) -> Result<&'a Vec<Value>> {
59    field(o, key)?.as_array().ok_or_else(|| bad(key))
60}
61
62/// A family as `{"name", ...parameters}`.
63pub fn family_json(f: Family) -> Value {
64    match f {
65        Family::NegativeBinomial { theta } => json!({"name": f.name(), "theta": num(theta)}),
66        Family::Tweedie { power } => json!({"name": f.name(), "power": num(power)}),
67        _ => json!({"name": f.name()}),
68    }
69}
70
71/// The family [`family_json`] wrote.
72pub fn family_from(o: &Map<String, Value>) -> Result<Family> {
73    Ok(match field(o, "name")?.as_str() {
74        Some("gaussian") => Family::Gaussian,
75        Some("poisson") => Family::Poisson,
76        Some("gamma") => Family::Gamma,
77        Some("inverse_gaussian") => Family::InverseGaussian,
78        Some("binomial") => Family::Binomial,
79        Some("negative_binomial") => Family::NegativeBinomial {
80            theta: float(field(o, "theta")?, "family.theta")?,
81        },
82        Some("tweedie") => Family::Tweedie {
83            power: float(field(o, "power")?, "family.power")?,
84        },
85        _ => return Err(bad("spec.family.name")),
86    })
87}
88
89/// A link as `{"name", ...parameters}`.
90pub fn link_json(l: Link) -> Value {
91    let name = match l {
92        Link::Identity => "identity",
93        Link::Log => "log",
94        Link::Logit => "logit",
95        Link::Probit => "probit",
96        Link::Cloglog => "cloglog",
97        Link::Inverse => "inverse",
98        Link::InverseSquared => "inverse_squared",
99        Link::Power(p) => return json!({"name": "power", "power": num(p)}),
100    };
101    json!({"name": name})
102}
103
104/// The link [`link_json`] wrote.
105pub fn link_from(o: &Map<String, Value>) -> Result<Link> {
106    Ok(match field(o, "name")?.as_str() {
107        Some("identity") => Link::Identity,
108        Some("log") => Link::Log,
109        Some("logit") => Link::Logit,
110        Some("probit") => Link::Probit,
111        Some("cloglog") => Link::Cloglog,
112        Some("inverse") => Link::Inverse,
113        Some("inverse_squared") => Link::InverseSquared,
114        Some("power") => Link::Power(float(field(o, "power")?, "link.power")?),
115        _ => return Err(bad("spec.link.name")),
116    })
117}
118
119/// A number, or the string `"NaN"`, `"inf"` or `"-inf"`.
120pub fn num(x: f64) -> Value {
121    match Number::from_f64(x) {
122        Some(n) => Value::Number(n),
123        None if x.is_nan() => Value::String("NaN".into()),
124        None if x > 0.0 => Value::String("inf".into()),
125        None => Value::String("-inf".into()),
126    }
127}
128
129/// An array of [`num`]s.
130pub fn nums(xs: &[f64]) -> Value {
131    Value::Array(xs.iter().map(|&x| num(x)).collect())
132}
133
134/// The number [`num`] wrote.
135pub fn float(v: &Value, what: &str) -> Result<f64> {
136    match v {
137        Value::Number(n) => n.as_f64().ok_or_else(|| bad(what)),
138        Value::String(s) => match s.as_str() {
139            "NaN" => Ok(f64::NAN),
140            "inf" => Ok(f64::INFINITY),
141            "-inf" => Ok(f64::NEG_INFINITY),
142            _ => Err(bad(what)),
143        },
144        _ => Err(bad(what)),
145    }
146}
147
148/// The array of [`num`]s at `key`.
149pub fn floats(o: &Map<String, Value>, key: &str) -> Result<Vec<f64>> {
150    field(o, key)?
151        .as_array()
152        .ok_or_else(|| bad(key))?
153        .iter()
154        .map(|v| float(v, key))
155        .collect()
156}
157
158/// The unsigned integer at `key`.
159pub fn uint(o: &Map<String, Value>, key: &str) -> Result<u64> {
160    field(o, key)?.as_u64().ok_or_else(|| bad(key))
161}
162
163/// The value at `key`, or an error naming it.
164pub fn field<'a>(o: &'a Map<String, Value>, key: &str) -> Result<&'a Value> {
165    o.get(key)
166        .ok_or_else(|| Error::Data(format!("model artifact has no {key:?}")))
167}
168
169/// `v` as an object.
170pub fn object<'a>(v: &'a Value, what: &str) -> Result<&'a Map<String, Value>> {
171    v.as_object().ok_or_else(|| bad(what))
172}
173
174/// The error for an invalid field.
175pub fn bad(what: &str) -> Error {
176    Error::Data(format!("model artifact field {what:?} is invalid"))
177}