Skip to main content

chronix_analytics/forecast/
parallel.rs

1//! Parallel forecasting utilities built on rayon.
2
3use rayon::prelude::*;
4
5use crate::forecast::error::ForecastError;
6use crate::forecast::result::ForecastResult;
7use crate::forecast::traits::ForecastModel;
8
9/// Fit multiple series in parallel.
10///
11/// Each element is (timestamps, values) for one series.
12/// Returns one `Result` per series in the same order.
13///
14/// # Errors
15///
16/// Returns a single-element vec with `ForecastError` if slice lengths differ.
17pub fn parallel_fit<M: ForecastModel + Send>(
18    models: &mut [M],
19    data: &[(&[i64], &[f64])],
20) -> Vec<Result<(), ForecastError>> {
21    let _start = std::time::Instant::now();
22    if models.len() != data.len() {
23        let err_msg = format!(
24            "models.len() ({}) != data.len() ({})",
25            models.len(),
26            data.len()
27        );
28        return (0..models.len())
29            .map(|_| Err(ForecastError::InvalidInput(err_msg.clone())))
30            .collect();
31    }
32    let results: Vec<_> = models
33        .par_iter_mut()
34        .zip(data.par_iter())
35        .map(|(model, (ts, vals))| model.fit(ts, vals))
36        .collect();
37    metrics::histogram!("chronix_compute_parallel_fit_duration_seconds")
38        .record(_start.elapsed().as_secs_f64());
39    results
40}
41
42/// Predict from multiple fitted models in parallel.
43///
44/// Returns one `Result<ForecastResult>` per model in the same order.
45pub fn parallel_predict<M: ForecastModel + Sync>(
46    models: &[M],
47    horizon: usize,
48) -> Vec<Result<ForecastResult, ForecastError>> {
49    let _start = std::time::Instant::now();
50    let results: Vec<_> = models
51        .par_iter()
52        .map(|model| model.predict(horizon))
53        .collect();
54    metrics::histogram!("chronix_compute_parallel_predict_duration_seconds")
55        .record(_start.elapsed().as_secs_f64());
56    results
57}
58
59/// Fit a single model type across many series in parallel, then predict.
60///
61/// This is a convenience function that creates one model per series,
62/// fits each, and predicts the requested horizon.
63///
64/// The `factory` closure creates a fresh model instance for each series.
65pub fn parallel_fit_predict<M, F>(
66    data: &[(&[i64], &[f64])],
67    horizon: usize,
68    factory: F,
69) -> Vec<Result<ForecastResult, ForecastError>>
70where
71    M: ForecastModel + Send,
72    F: Fn() -> M + Sync,
73{
74    data.par_iter()
75        .map(|(ts, vals)| {
76            let mut model = factory();
77            model.fit(ts, vals)?;
78            model.predict(horizon)
79        })
80        .collect()
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::forecast::ses::SesModel;
87
88    fn make_series(len: usize, offset: f64) -> (Vec<i64>, Vec<f64>) {
89        let ts: Vec<i64> = (0..len as i64).map(|i| i * 1_000_000_000).collect();
90        let vals: Vec<f64> = (0..len).map(|i| i as f64 + offset).collect();
91        (ts, vals)
92    }
93
94    #[test]
95    fn parallel_fit_all_succeed() {
96        let s1 = make_series(50, 0.0);
97        let s2 = make_series(50, 100.0);
98        let s3 = make_series(50, 200.0);
99        let data: Vec<(&[i64], &[f64])> = vec![(&s1.0, &s1.1), (&s2.0, &s2.1), (&s3.0, &s3.1)];
100        let mut models = vec![
101            SesModel::new(None),
102            SesModel::new(None),
103            SesModel::new(None),
104        ];
105        let results = parallel_fit(&mut models, &data);
106        assert!(results.iter().all(std::result::Result::is_ok));
107    }
108
109    #[test]
110    fn parallel_predict_produces_forecasts() {
111        let s1 = make_series(50, 0.0);
112        let s2 = make_series(50, 100.0);
113        let data: Vec<(&[i64], &[f64])> = vec![(&s1.0, &s1.1), (&s2.0, &s2.1)];
114        let mut models = vec![SesModel::new(None), SesModel::new(None)];
115        parallel_fit(&mut models, &data);
116        let preds = parallel_predict(&models, 5);
117        assert_eq!(preds.len(), 2);
118        for p in &preds {
119            let r = p.as_ref().unwrap();
120            assert_eq!(r.values.len(), 5);
121        }
122    }
123
124    #[test]
125    fn parallel_fit_predict_convenience() {
126        let s1 = make_series(50, 0.0);
127        let s2 = make_series(50, 100.0);
128        let s3 = make_series(50, 200.0);
129        let data: Vec<(&[i64], &[f64])> = vec![(&s1.0, &s1.1), (&s2.0, &s2.1), (&s3.0, &s3.1)];
130        let results = parallel_fit_predict(&data, 10, || SesModel::new(None));
131        assert_eq!(results.len(), 3);
132        for r in &results {
133            let fc = r.as_ref().unwrap();
134            assert_eq!(fc.values.len(), 10);
135        }
136    }
137
138    #[test]
139    fn partial_failure_propagates() {
140        let good = make_series(50, 0.0);
141        let bad_ts: Vec<i64> = vec![0];
142        let bad_vals: Vec<f64> = vec![1.0];
143        let data: Vec<(&[i64], &[f64])> = vec![(&good.0, &good.1), (&bad_ts, &bad_vals)];
144        let mut models = vec![SesModel::new(None), SesModel::new(None)];
145        let results = parallel_fit(&mut models, &data);
146        assert!(results[0].is_ok());
147        assert!(results[1].is_err());
148    }
149}