Skip to main content

henad_explore/
schema.rs

1//! A model's parameters, stats and actions as JSON, for tools that build sweeps from outside.
2
3use serde_json::{Map, Value, json};
4
5use henad_compute::entry::ModelEntry;
6use henad_core::explore::fingerprint::schema_hash;
7use henad_core::metadata::Backend;
8use henad_core::params::{ParamApply, ParamDescriptor, ParamFormat, ParamKind};
9
10use crate::probe::ProbeReport;
11
12/// Version of the object [`schema_json`] returns.
13pub const SCHEMA_VERSION: u64 = 1;
14
15/// Returns the parameters, stats and actions of `entry` as one JSON object.
16///
17/// The object carries `stat_columns`, the columns a sweep writes, only when `probe` is given. A choice's default is
18/// its option name, with its index beside it as `default_index`.
19pub fn schema_json(entry: &ModelEntry, probe: Option<&ProbeReport>) -> Value {
20    let params: Vec<Value> = entry
21        .param_descriptors()
22        .iter()
23        .enumerate()
24        .map(|(index, descriptor)| param_json(index, descriptor))
25        .collect();
26    let stats: Vec<Value> = entry
27        .stat_descriptors()
28        .iter()
29        .map(|stat| json!({ "label": stat.label, "color": stat.color }))
30        .collect();
31    let actions: Vec<Value> = entry
32        .action_descriptors()
33        .iter()
34        .enumerate()
35        .map(|(index, action)| json!({ "index": index, "id": action.id, "label": action.label }))
36        .collect();
37    let mut schema = json!({
38        "schema_version": SCHEMA_VERSION,
39        "model": entry.id(),
40        "name": entry.name(),
41        "backend": backend_name(entry.metadata().backend),
42        "schema_hash": format!("{:016x}", schema_hash(&entry.schema())),
43        "params": params,
44        "stats": stats,
45        "actions": actions,
46        "seed": { "type": "u64" },
47    });
48    if let Some(probe) = probe {
49        let columns: Vec<&str> = (0..probe.columns.len()).map(|i| probe.columns.name(i)).collect();
50        schema["stat_columns"] = json!(columns);
51    }
52    schema
53}
54
55/// Returns `backend` as written in JSON and CSV, `cpu` or `gpu`.
56pub fn backend_name(backend: Backend) -> &'static str {
57    match backend {
58        Backend::Cpu => "cpu",
59        Backend::Gpu => "gpu",
60    }
61}
62
63fn param_json(index: usize, descriptor: &ParamDescriptor) -> Value {
64    let mut param = Map::new();
65    param.insert("index".to_owned(), json!(index));
66    param.insert("id".to_owned(), json!(descriptor.id));
67    param.insert("label".to_owned(), json!(descriptor.label));
68    let apply = match descriptor.apply {
69        ParamApply::Live => "live",
70        ParamApply::OnReload => "reload",
71    };
72    param.insert("apply".to_owned(), json!(apply));
73    let format = match descriptor.format {
74        ParamFormat::Plain => "plain",
75        ParamFormat::Percent => "percent",
76    };
77    param.insert("format".to_owned(), json!(format));
78    let fields = match descriptor.kind {
79        ParamKind::F32 {
80            min,
81            max,
82            default,
83            step,
84        } => {
85            let mut fields = json!({
86                "kind": "f32",
87                "min": f32_json(min),
88                "max": f32_json(max),
89                "default": f32_json(default),
90            });
91            if let Some(step) = step {
92                fields["step"] = f32_json(step);
93            }
94            fields
95        }
96        ParamKind::U32 { min, max, default } => {
97            json!({ "kind": "u32", "min": min, "max": max, "default": default })
98        }
99        ParamKind::Bool { default } => json!({ "kind": "bool", "default": default }),
100        ParamKind::Choice { options, default } => json!({
101            "kind": "choice",
102            "options": options,
103            "default": options.get(default),
104            "default_index": default,
105        }),
106    };
107    if let Value::Object(fields) = fields {
108        param.extend(fields);
109    }
110    Value::Object(param)
111}
112
113/// Returns `value` as the JSON number with the shortest decimal form that parses back to `value`.
114///
115/// A plain conversion widens to `f64` first and writes `0.025` as `0.02500000037252903`.
116pub(crate) fn f32_json(value: f32) -> Value {
117    value
118        .to_string()
119        .parse::<f64>()
120        .ok()
121        .and_then(serde_json::Number::from_f64)
122        .map_or(Value::Null, Value::Number)
123}
124
125#[cfg(test)]
126mod tests {
127    use henad_models::example_models;
128    use serde_json::json;
129
130    use henad_core::explore::fingerprint::schema_hash;
131
132    use super::{f32_json, schema_json};
133    use crate::probe::ProbeReport;
134
135    /// Each example model's `schema_hash` as Henad 0.2.0 wrote it, at commit 773a7a5.
136    ///
137    /// `crates/henad-models/tests/fixtures/docs/schema-hashes-0.2.0.md` describes the procedure that recorded them.
138    const SCHEMA_HASHES_0_2_0: [(&str, &str); 10] = [
139        ("sir", "6ff1dc3971fd0a96"),
140        ("boids", "ae99f3e0d37c8d3d"),
141        ("game_of_life", "aba7303a467b06bf"),
142        ("ants", "f6eb31e76efdf4cf"),
143        ("virus_network", "77993c7b4047b8bd"),
144        ("team_assembly", "d1724ce65dada5d6"),
145        ("gpu_game_of_life", "461ad8e0d063a398"),
146        ("gpu_sir", "99f14d30367e3743"),
147        ("gpu_boids", "f7328729019009c0"),
148        ("gpu_ants", "b58e8a5a5b6a829e"),
149    ];
150
151    /// Checks that every example model hashes as it did in 0.2.0, so every folder written since still resumes.
152    ///
153    /// A schema reads only declarations, and the GPU models are checked without a device.
154    #[test]
155    fn schema_hashes_are_unchanged_since_0_2_0() {
156        let models = example_models();
157        for (id, recorded) in SCHEMA_HASHES_0_2_0 {
158            let entry = models.get(id).unwrap_or_else(|| panic!("{id} is registered"));
159            let current = format!("{:016x}", schema_hash(&entry.schema()));
160            assert_eq!(current, recorded, "{id}'s schema hash moved since 0.2.0");
161        }
162    }
163
164    #[test]
165    fn every_descriptor_is_listed_with_its_kind_and_bounds() {
166        let models = example_models();
167        for entry in models.iter().filter(|entry| entry.gpu_needs().is_none()) {
168            let schema = schema_json(entry, None);
169            let params = schema["params"].as_array().expect("params is a list");
170            assert_eq!(params.len(), entry.param_descriptors().len(), "{}", entry.id());
171            for (index, param) in params.iter().enumerate() {
172                assert_eq!(param["index"], json!(index));
173                assert_eq!(param["id"], json!(entry.param_descriptors()[index].id));
174                assert!(param["kind"].is_string() && param["default"] != json!(null), "{param}");
175            }
176            assert_eq!(schema["model"], json!(entry.id()));
177            assert_eq!(schema["schema_hash"].as_str().map(str::len), Some(16));
178            assert!(schema.get("stat_columns").is_none(), "no probe, no columns");
179        }
180    }
181
182    #[test]
183    fn a_probe_adds_the_stat_columns() {
184        let entry = henad_models::example_models()
185            .get("sir")
186            .cloned()
187            .expect("sir is registered");
188        let defaults: Vec<_> = entry
189            .param_descriptors()
190            .iter()
191            .map(|descriptor| descriptor.kind.default_value())
192            .collect();
193        let probe = ProbeReport::build(&entry, None, &defaults, None).expect("sir builds at its defaults");
194        let schema = schema_json(&entry, Some(&probe));
195        assert_eq!(schema["stat_columns"], json!(["Susceptible", "Infected", "Recovered"]));
196        let grid_width = &schema["params"][0];
197        assert_eq!(grid_width["id"], json!("grid_width"));
198        assert_eq!(grid_width["kind"], json!("u32"));
199        assert_eq!(grid_width["apply"], json!("reload"));
200    }
201
202    #[test]
203    fn an_f32_is_written_in_its_shortest_form() {
204        assert_eq!(f32_json(0.025).to_string(), "0.025");
205        assert_eq!(f32_json(1.0).to_string(), "1.0");
206        assert_eq!(f32_json(f32::NAN), json!(null));
207    }
208}