1use prospicio_core::{Error, Result};
9use serde_json::{Map, Number, Value, json};
10
11use crate::{Family, Link};
12
13pub 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
23pub 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
40pub 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
49pub 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
57pub 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
62pub 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
71pub 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
89pub 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
104pub 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
119pub 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
129pub fn nums(xs: &[f64]) -> Value {
131 Value::Array(xs.iter().map(|&x| num(x)).collect())
132}
133
134pub 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
148pub 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
158pub fn uint(o: &Map<String, Value>, key: &str) -> Result<u64> {
160 field(o, key)?.as_u64().ok_or_else(|| bad(key))
161}
162
163pub 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
169pub fn object<'a>(v: &'a Value, what: &str) -> Result<&'a Map<String, Value>> {
171 v.as_object().ok_or_else(|| bad(what))
172}
173
174pub fn bad(what: &str) -> Error {
176 Error::Data(format!("model artifact field {what:?} is invalid"))
177}