chronix_analytics/forecast/
ses.rs

1//! Simple Exponential Smoothing (SES) forecast model.
2
3use crate::forecast::error::ForecastError;
4use crate::forecast::result::{ForecastResult, ModelParams, ModelType};
5use crate::forecast::traits::ForecastModel;
6use crate::forecast::util::{median_interval, validate_input};
7
8/// Simple Exponential Smoothing model.
9///
10/// Flat-forecast at last smoothed level. Alpha is auto-optimized via Brent's
11/// method (1-D line search) minimizing MSE if not specified.
12#[derive(serde::Serialize, serde::Deserialize)]
13pub struct SesModel {
14    alpha: Option<f64>,
15    params: ModelParams,
16    last_ts: i64,
17    interval_ns: i64,
18    n_points: usize,
19    fitted: bool,
20}
21
22impl SesModel {
23    /// Creates a new SES model. If `alpha` is `None`, it will be auto-optimized.
24    pub fn new(alpha: Option<f64>) -> Self {
25        Self {
26            alpha,
27            params: ModelParams::Ses {
28                alpha: 0.0,
29                level: 0.0,
30                residual_std: 0.0,
31            },
32            last_ts: 0,
33            interval_ns: 0,
34            n_points: 0,
35            fitted: false,
36        }
37    }
38
39    /// Find optimal alpha by Brent's method minimizing MSE.
40    fn optimize_alpha(values: &[f64]) -> f64 {
41        let result = crate::forecast::optimizer::minimize_brent(
42            |alpha| Self::compute_mse(values, alpha),
43            0.01,
44            0.99,
45            1e-8,
46            100,
47        );
48        result.params[0].clamp(0.01, 0.99)
49    }
50
51    fn compute_mse(values: &[f64], alpha: f64) -> f64 {
52        if values.len() < 2 {
53            return 0.0;
54        }
55        let one_minus = 1.0 - alpha;
56        let mut level = values[0];
57        let mut sse = 0.0;
58        for &value in values.iter().skip(1) {
59            let err = value - level;
60            sse += err * err;
61            level = alpha * value + one_minus * level;
62        }
63        sse / (values.len() - 1) as f64
64    }
65}
66
67impl ForecastModel for SesModel {
68    #[tracing::instrument(skip_all, level = "debug")]
69    fn fit(&mut self, timestamps: &[i64], values: &[f64]) -> Result<(), ForecastError> {
70        let _start = std::time::Instant::now();
71        validate_input(timestamps, values)?;
72        if values.len() < 2 {
73            return Err(ForecastError::InsufficientData {
74                min: 2,
75                got: values.len(),
76            });
77        }
78
79        let alpha = self.alpha.unwrap_or_else(|| Self::optimize_alpha(values));
80        let one_minus = 1.0 - alpha;
81
82        // Compute smoothed level and residual statistics
83        let mut level = values[0];
84        let mut sse = 0.0;
85        for &value in values.iter().skip(1) {
86            let err = value - level;
87            sse += err * err;
88            level = alpha * value + one_minus * level;
89        }
90
91        // Correct degrees of freedom — we have n-1 residuals and
92        // 1 estimated parameter (α), so divisor = max(n-2, 1).
93        let dof = (values.len() as f64 - 2.0).max(1.0);
94        let residual_std = (sse / dof).sqrt();
95
96        self.interval_ns = median_interval(timestamps);
97        self.last_ts = timestamps[timestamps.len() - 1];
98        self.n_points = values.len();
99
100        self.params = ModelParams::Ses {
101            alpha,
102            level,
103            residual_std,
104        };
105        self.fitted = true;
106        metrics::histogram!("chronix_forecast_fit_duration_seconds", "model_type" => "ses")
107            .record(_start.elapsed().as_secs_f64());
108        Ok(())
109    }
110
111    #[tracing::instrument(skip_all, level = "debug")]
112    fn predict(&self, horizon: usize) -> Result<ForecastResult, ForecastError> {
113        let _start = std::time::Instant::now();
114        if !self.fitted {
115            return Err(ForecastError::NotFitted);
116        }
117
118        let ModelParams::Ses {
119            alpha,
120            level,
121            residual_std,
122        } = &self.params
123        else {
124            return Err(ForecastError::InvalidInput(
125                "unexpected model params variant".into(),
126            ));
127        };
128
129        let z = 1.96; // 95% confidence
130        let mut values = Vec::with_capacity(horizon);
131        let mut timestamps = Vec::with_capacity(horizon);
132        let mut lower = Vec::with_capacity(horizon);
133        let mut upper = Vec::with_capacity(horizon);
134
135        for h in 1..=horizon {
136            let ts = self
137                .last_ts
138                .saturating_add((h as i64).saturating_mul(self.interval_ns));
139            let width = z * residual_std * (1.0 + (h as f64 - 1.0) * alpha * alpha).sqrt();
140            values.push(*level);
141            timestamps.push(ts);
142            lower.push(level - width);
143            upper.push(level + width);
144        }
145
146        metrics::histogram!("chronix_forecast_predict_duration_seconds", "model_type" => "ses")
147            .record(_start.elapsed().as_secs_f64());
148        Ok(ForecastResult {
149            values,
150            timestamps,
151            confidence_lower: lower,
152            confidence_upper: upper,
153            confidence_level: 0.95,
154        })
155    }
156
157    fn update(&mut self, timestamp: i64, value: f64) -> Result<(), ForecastError> {
158        if !self.fitted {
159            return Err(ForecastError::NotFitted);
160        }
161
162        let ModelParams::Ses {
163            alpha,
164            ref mut level,
165            ..
166        } = self.params
167        else {
168            return Err(ForecastError::InvalidInput(
169                "unexpected model params variant".into(),
170            ));
171        };
172
173        *level = alpha * value + (1.0 - alpha) * *level;
174        self.last_ts = timestamp;
175        self.n_points += 1;
176        Ok(())
177    }
178
179    fn model_type(&self) -> ModelType {
180        ModelType::Ses
181    }
182
183    fn params(&self) -> &ModelParams {
184        &self.params
185    }
186}
187
188impl crate::forecast::storage::ModelStore for SesModel {}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn ses_constant_series() {
196        let ts: Vec<i64> = (0..100).map(|i| i * 1_000_000_000).collect();
197        let vals = vec![42.0; 100];
198        let mut model = SesModel::new(None);
199        model.fit(&ts, &vals).unwrap();
200        let result = model.predict(10).unwrap();
201        for &v in &result.values {
202            assert!((v - 42.0).abs() < 1e-6);
203        }
204    }
205
206    #[test]
207    fn ses_trending_series() {
208        let ts: Vec<i64> = (0..100).map(|i| i * 1_000_000_000).collect();
209        let vals: Vec<f64> = (0..100).map(|i| i as f64 * 2.0).collect();
210        let mut model = SesModel::new(Some(0.8));
211        model.fit(&ts, &vals).unwrap();
212        let result = model.predict(5).unwrap();
213        // SES can't track trend perfectly, but forecast should be near last value
214        let last = vals.last().unwrap();
215        assert!((result.values[0] - last).abs() < 20.0);
216    }
217
218    #[test]
219    fn ses_auto_alpha() {
220        let ts: Vec<i64> = (0..50).map(|i| i * 1_000_000_000).collect();
221        let vals: Vec<f64> = (0..50).map(|i| (i as f64).sin() * 10.0 + 50.0).collect();
222        let mut model = SesModel::new(None);
223        model.fit(&ts, &vals).unwrap();
224        let ModelParams::Ses { alpha, .. } = model.params() else {
225            panic!()
226        };
227        assert!(*alpha > 0.0 && *alpha <= 1.0);
228    }
229
230    #[test]
231    fn ses_update() {
232        let ts: Vec<i64> = (0..10).map(|i| i * 1_000_000_000).collect();
233        let vals = vec![10.0; 10];
234        let mut model = SesModel::new(Some(0.5));
235        model.fit(&ts, &vals).unwrap();
236
237        // Update with a high value — level should shift
238        model.update(10_000_000_000, 100.0).unwrap();
239        let ModelParams::Ses { level, .. } = model.params() else {
240            panic!()
241        };
242        assert!(*level > 10.0 && *level < 100.0);
243    }
244
245    #[test]
246    fn ses_confidence_intervals_widen() {
247        let ts: Vec<i64> = (0..50).map(|i| i * 1_000_000_000).collect();
248        let vals: Vec<f64> = (0..50).map(|i| i as f64 + (i as f64 * 0.5).sin()).collect();
249        let mut model = SesModel::new(Some(0.3));
250        model.fit(&ts, &vals).unwrap();
251        let result = model.predict(10).unwrap();
252        // Confidence intervals should widen with horizon
253        let width_1 = result.confidence_upper[0] - result.confidence_lower[0];
254        let width_10 = result.confidence_upper[9] - result.confidence_lower[9];
255        assert!(width_10 > width_1);
256    }
257
258    #[test]
259    fn ses_insufficient_data() {
260        let mut model = SesModel::new(None);
261        assert!(model.fit(&[0], &[1.0]).is_err());
262    }
263
264    #[test]
265    fn ses_not_fitted() {
266        let model = SesModel::new(None);
267        assert!(model.predict(10).is_err());
268    }
269
270    #[test]
271    fn metrics_recorded_on_fit_predict() {
272        // Metrics macros are no-ops without a recorder, but we verify the
273        // code paths execute without panic — the `metrics` crate guarantees
274        // this is safe even without a subscriber installed.
275        let ts: Vec<i64> = (0..100).map(|i| i * 1_000_000_000).collect();
276        let vals: Vec<f64> = (0..100).map(|i| 50.0 + (i as f64) * 0.1).collect();
277        let mut model = SesModel::new(Some(0.3));
278        model.fit(&ts, &vals).unwrap();
279        let result = model.predict(10).unwrap();
280        assert_eq!(result.values.len(), 10);
281        // If metrics crate panicked, we wouldn't reach here
282    }
283}
chronix_analytics/forecast/ses.rs

chronix_analytics/forecast/
ses.rs
