chronix_analytics/forecast/
parallel.rs1use rayon::prelude::*;
4
5use crate::forecast::error::ForecastError;
6use crate::forecast::result::ForecastResult;
7use crate::forecast::traits::ForecastModel;
8
9pub 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
42pub 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
59pub 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}