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}