Skip to main content

zsfm_core/
input.rs

1use anyhow::Context;
2
3/// Parse the `context` field of a forecast request into `[batch][variate][time]`.
4///
5/// Accepted shapes:
6/// - `[t0, t1, ...]` → batch=1, n_var=1 (flat single series)
7/// - `[[t0, t1, ...], ...]` → batch=N, n_var=1 (batch of univariate)
8/// - `[[[v0_t0, ...], [v1_t0, ...]], ...]` → batch=N, n_var=M (batch of multivariate)
9pub fn parse_mv_contexts(val: serde_json::Value) -> anyhow::Result<Vec<Vec<Vec<f32>>>> {
10    match val {
11        serde_json::Value::Array(arr) if arr.is_empty() => {
12            anyhow::bail!("context must be a non-empty array")
13        }
14        serde_json::Value::Array(arr) => {
15            let first = arr.first().unwrap();
16            if !first.is_array() {
17                // Flat: [t0, t1, ...] → batch=1, n_var=1
18                let ctx = serde_json::from_value::<Vec<f32>>(serde_json::Value::Array(arr))
19                    .context("context must be a JSON array of numbers")?;
20                Ok(vec![vec![ctx]])
21            } else if first
22                .as_array()
23                .and_then(|a| a.first())
24                .map(|v| v.is_array())
25                .unwrap_or(false)
26            {
27                // 3D: [batch][variate][time]
28                arr.into_iter()
29                    .enumerate()
30                    .map(|(i, batch_item)| {
31                        serde_json::from_value::<Vec<Vec<f32>>>(batch_item)
32                            .with_context(|| format!("context[{i}] must be an array of variate arrays"))
33                    })
34                    .collect()
35            } else {
36                // 2D: [batch][time] → each series is n_var=1
37                arr.into_iter()
38                    .enumerate()
39                    .map(|(i, v)| {
40                        let series = serde_json::from_value::<Vec<f32>>(v)
41                            .with_context(|| format!("context[{i}] must be an array of numbers"))?;
42                        Ok(vec![series])
43                    })
44                    .collect()
45            }
46        }
47        _ => anyhow::bail!("context must be a JSON array"),
48    }
49}