zsfm_core/
json.rs

1use std::collections::BTreeMap;
2use std::time::{SystemTime, UNIX_EPOCH};
3
4use serde::Serialize;
5
6use crate::forecast::QuantileMatrix;
7
8#[derive(Serialize)]
9#[serde(untagged)]
10pub enum ForecastOutput {
11    Univariate {
12        point: Vec<f32>,
13        quantiles: BTreeMap<String, Vec<f32>>,
14    },
15    Multivariate {
16        variates: Vec<VariateForecast>,
17    },
18}
19
20#[derive(Serialize)]
21pub struct VariateForecast {
22    pub point: Vec<f32>,
23    pub quantiles: BTreeMap<String, Vec<f32>>,
24}
25
26/// Turn one forecast's `QuantileMatrix` into a `ForecastOutput`, choosing the
27/// univariate (flat `point`/`quantiles`) or multivariate (`variates: [...]`)
28/// shape based on how many variates the matrix actually has.
29///
30/// `quantile_levels[i]` is the probability level for `qmat[i]`; `median_idx`
31/// is the row holding the point forecast (usually the index of level 0.5).
32pub fn quantile_matrix_to_output(
33    qmat: &QuantileMatrix,
34    quantile_levels: &[f32],
35    median_idx: usize,
36) -> ForecastOutput {
37    let n_var = qmat.first().map(|q| q.len()).unwrap_or(0);
38
39    let quantiles_for = |vi: usize| -> BTreeMap<String, Vec<f32>> {
40        let mut quantiles = BTreeMap::new();
41        for (qi, &level) in quantile_levels.iter().enumerate() {
42            if let Some(row) = qmat.get(qi) {
43                if let Some(q) = row.get(vi) {
44                    quantiles.insert(format!("{level:.2}"), q.clone());
45                }
46            }
47        }
48        quantiles
49    };
50    let point_for = |vi: usize| -> Vec<f32> {
51        qmat.get(median_idx).and_then(|v| v.get(vi)).cloned().unwrap_or_default()
52    };
53
54    if n_var <= 1 {
55        ForecastOutput::Univariate {
56            point: point_for(0),
57            quantiles: quantiles_for(0),
58        }
59    } else {
60        let variates = (0..n_var)
61            .map(|vi| VariateForecast { point: point_for(vi), quantiles: quantiles_for(vi) })
62            .collect();
63        ForecastOutput::Multivariate { variates }
64    }
65}
66
67/// Wrap one or more `ForecastOutput`s in the shared response envelope every
68/// model's `infer` CLI subcommand (and later, Python binding) returns.
69pub fn forecast_response_json(
70    model_name: &str,
71    context_length: usize,
72    forecast_length: usize,
73    outputs: Vec<ForecastOutput>,
74) -> anyhow::Result<String> {
75    #[derive(Serialize)]
76    struct ForecastResponse {
77        id: String,
78        object: &'static str,
79        created: u64,
80        model: String,
81        choices: Vec<Choice>,
82        usage: Usage,
83    }
84    #[derive(Serialize)]
85    struct Choice {
86        index: usize,
87        forecast: ForecastOutput,
88        finish_reason: &'static str,
89    }
90    #[derive(Serialize)]
91    struct Usage {
92        context_length: usize,
93        forecast_length: usize,
94    }
95
96    let created = SystemTime::now()
97        .duration_since(UNIX_EPOCH)
98        .unwrap_or_default()
99        .as_secs();
100
101    let resp = ForecastResponse {
102        id: format!("forecast-{created:016x}"),
103        object: "forecast",
104        created,
105        model: model_name.to_string(),
106        choices: outputs
107            .into_iter()
108            .enumerate()
109            .map(|(i, forecast)| Choice { index: i, forecast, finish_reason: "stop" })
110            .collect(),
111        usage: Usage { context_length, forecast_length },
112    };
113
114    Ok(serde_json::to_string_pretty(&resp)?)
115}
zsfm_core/json.rs

zsfm_core/
json.rs
