augurs/ets.rs
1//! Bindings for AutoETS model search.
2use numpy::PyReadonlyArrayDyn;
3use pyo3::{exceptions::PyException, prelude::*};
4
5use crate::Forecast;
6
7/// Automatic exponential smoothing model search.
8#[derive(Debug)]
9#[pyclass]
10pub struct AutoETS {
11 inner: augurs_ets::AutoETS,
12}
13
14#[pymethods]
15impl AutoETS {
16 /// Create a new `AutoETS` model search instance.
17 ///
18 /// # Errors
19 ///
20 /// If the `spec` string is invalid, this function returns an error.
21 #[new]
22 #[allow(clippy::too_many_arguments)]
23 pub fn new(season_length: usize, spec: String) -> PyResult<Self> {
24 let inner = augurs_ets::AutoETS::new(season_length, spec.as_str())
25 .map_err(|e| PyException::new_err(e.to_string()))?;
26 Ok(Self { inner })
27 }
28
29 fn __repr__(&self) -> String {
30 format!(
31 "AutoETS(spec=\"{}\", season_length={})",
32 self.inner.spec(),
33 self.inner.season_length()
34 )
35 }
36
37 /// Search for the best model, fitting it to the data.
38 ///
39 /// The model will be stored on the inner `AutoETS` instance, after which
40 /// forecasts can be produced using its `predict` method.
41 ///
42 /// # Errors
43 ///
44 /// If no model can be found, or if any parameters are invalid, this function
45 /// returns an error.
46 pub fn fit(&mut self, y: PyReadonlyArrayDyn<'_, f64>) -> PyResult<()> {
47 self.inner
48 .fit(y.as_slice()?)
49 .map_err(|e| PyException::new_err(e.to_string()))?;
50 Ok(())
51 }
52
53 /// Predict the next `horizon` values using the best model, optionally including
54 /// prediction intervals at the specified level.
55 ///
56 /// `level` should be a float between 0 and 1 representing the confidence level.
57 ///
58 /// # Errors
59 ///
60 /// This function will return an error if no model has been fit yet (using [`AutoETS::fit`]).
61 pub fn predict(&self, horizon: usize, level: Option<f64>) -> PyResult<Forecast> {
62 self.inner
63 .predict(horizon, level)
64 .map(Forecast::from)
65 .map_err(|e| PyException::new_err(e.to_string()))
66 }
67
68 /// Get the in-sample predictions for the model, optionally including
69 /// prediction intervals at the specified level.
70 ///
71 /// `level` should be a float between 0 and 1 representing the confidence level.
72 ///
73 /// # Errors
74 ///
75 /// This function will return an error if no model has been fit yet (using [`AutoETS::fit`]).
76 pub fn predict_in_sample(&self, level: Option<f64>) -> PyResult<Forecast> {
77 self.inner
78 .predict_in_sample(level)
79 .map(Forecast::from)
80 .map_err(|e| PyException::new_err(e.to_string()))
81 }
82}