Add power transformation logic to forecaster transforms

crates/augurs-forecaster/Cargo.toml

@@ -16,9 +16,10 @@ bench = false
 augurs-core.workspace = true
 itertools.workspace = true
 thiserror.workspace = true
+argmin = "0.10.0"
 
 [dev-dependencies]
-augurs.workspace = true
+augurs = { workspace = true, features = ["mstl", "ets", "forecaster"]}
 augurs-testing.workspace = true
 
 [lints]

crates/augurs-forecaster/src/lib.rs

@@ -3,12 +3,15 @@
 mod data;
 mod error;
 mod forecaster;
+mod power_transforms;
 pub mod transforms;
 
 pub use data::Data;
 pub use error::Error;
 pub use forecaster::Forecaster;
 pub use transforms::Transform;
-pub(crate) use transforms::Transforms;
+pub use transforms::Transforms;
+pub use power_transforms::optimize_lambda;
+
 
 type Result<T> = std::result::Result<T, Error>;

crates/augurs-forecaster/src/power_transforms.rs

@@ -0,0 +1,63 @@
+use argmin::core::*;
+use argmin::solver::brent::BrentOpt;
+use crate::transforms::box_cox;
+
+fn box_cox_log_likelihood(data: &[f64], lambda: f64) -> f64 {
+    let n = data.len() as f64;
+    let transformed_data: Vec<f64> = data.iter().map(|&x| box_cox(x, lambda)).collect();
+    let mean_transformed: f64 = transformed_data.iter().copied().sum::<f64>() / n;
+    let variance: f64 = transformed_data.iter().map(|&x| (x - mean_transformed).powi(2)).sum::<f64>() / n;
+
+    // Avoid log(0) by ensuring variance is positive
+    let log_likelihood = -0.5 * n * variance.ln() + (lambda - 1.0) * data.iter().map(|&x| x.ln()).sum::<f64>();
+    log_likelihood
+}
+
+#[derive(Clone)]
+struct BoxCoxProblem {
+    data: Vec<f64>,
+}
+
+impl CostFunction for BoxCoxProblem {
+    type Param = f64;
+    type Output = f64;
+
+    // The goal is to minimize the negative log-likelihood
+    fn cost(&self, lambda: &Self::Param) -> Result<Self::Output, Error> {
+        Ok(-box_cox_log_likelihood(&self.data, *lambda))
+    }
+}
+
+/// Optimize the lambda parameter for the Box-Cox transformation
+pub fn optimize_lambda(data: &[f64]) -> f64 {
+    let cost = BoxCoxProblem { data: data.to_vec() };
+    let init_param = 0.5;
+    let solver = BrentOpt::new(-2.0, 2.0);
+
+    let res = Executor::new(cost, solver)
+        .configure(|state| state.param(init_param).max_iters(100))
+        .run()
+        .unwrap();
+    return res.state.best_param.unwrap();
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use augurs_testing::assert_approx_eq;
+
+    #[test]
+    fn correct_optimal_lambda() {
+        let data = &[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7];
+        let got = optimize_lambda(data);
+        assert_approx_eq!(got, 0.7123778635679304);
+    }
+
+    #[test]
+    fn test_boxcox_llf() {
+        let data = &[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7];
+        let lambda = 1.0;
+        let got = box_cox_log_likelihood(data, lambda);
+        assert_approx_eq!(got, 11.266065387038703);
+    }
+}

crates/augurs-forecaster/src/transforms.rs

@@ -7,15 +7,21 @@ use augurs_core::{
     Forecast,
 };
 
+/// Transforms and Transform implementations.
+/// 
+/// The `Transforms` struct is a collection of `Transform` instances that can be applied to a time series.
+/// The `Transform` enum represents a single transformation that can be applied to a time series.
 #[derive(Debug, Default)]
-pub(crate) struct Transforms(Vec<Transform>);
+pub struct Transforms(Vec<Transform>);
 
 impl Transforms {
-    pub(crate) fn new(transforms: Vec<Transform>) -> Self {
+    /// create a new `Transforms` instance with the given transforms.
+    pub fn new(transforms: Vec<Transform>) -> Self {
         Self(transforms)
     }
 
-    pub(crate) fn transform<'a, T>(&'a self, input: T) -> Box<dyn Iterator<Item = f64> + 'a>
+    /// Apply the transformations to the given time series.
+    pub fn transform<'a, T>(&'a self, input: T) -> Box<dyn Iterator<Item = f64> + 'a>
     where
         T: Iterator<Item = f64> + 'a,
     {
@@ -24,7 +30,8 @@ impl Transforms {
             .fold(Box::new(input) as _, |y, t| t.transform(y))
     }
 
-    pub(crate) fn inverse_transform(&self, forecast: Forecast) -> Forecast {
+    /// Apply the inverse transformations to the given forecast.
+    pub fn inverse_transform(&self, forecast: Forecast) -> Forecast {
         self.0
             .iter()
             .rev()
@@ -46,6 +53,8 @@ pub enum Transform {
     Logit,
     /// Log transform.
     Log,
+    /// Box-Cox transform.
+    BoxCox(f64),
 }
 
 impl Transform {
@@ -81,7 +90,16 @@ impl Transform {
         Self::Log
     }
 
-    pub(crate) fn transform<'a, T>(&'a self, input: T) -> Box<dyn Iterator<Item = f64> + 'a>
+
+    /// Create a new Box-Cox transform.
+    /// 
+    /// This transform applies the Box-Cox transformation to each item.
+    pub fn boxcox(lambda: f64) -> Self {
+        Self::BoxCox(lambda)
+    }
+
+    /// Apply the transformation to the given time series.
+    pub fn transform<'a, T>(&'a self, input: T) -> Box<dyn Iterator<Item = f64> + 'a>
     where
         T: Iterator<Item = f64> + 'a,
     {
@@ -90,10 +108,12 @@ impl Transform {
             Self::MinMaxScaler(params) => Box::new(input.min_max_scale(params.clone())),
             Self::Logit => Box::new(input.logit()),
             Self::Log => Box::new(input.log()),
+            Self::BoxCox(lambda) => Box::new(input.boxcox(*lambda)),
         }
     }
 
-    pub(crate) fn inverse_transform<'a, T>(&'a self, input: T) -> Box<dyn Iterator<Item = f64> + 'a>
+    /// Apply the inverse transformation to the given time series.
+    pub fn inverse_transform<'a, T>(&'a self, input: T) -> Box<dyn Iterator<Item = f64> + 'a>
     where
         T: Iterator<Item = f64> + 'a,
     {
@@ -102,10 +122,12 @@ impl Transform {
             Self::MinMaxScaler(params) => Box::new(input.inverse_min_max_scale(params.clone())),
             Self::Logit => Box::new(input.logistic()),
             Self::Log => Box::new(input.exp()),
+            Self::BoxCox(lambda) => Box::new(input.inverse_boxcox(*lambda)),
         }
     }
 
-    pub(crate) fn inverse_transform_forecast(&self, mut f: Forecast) -> Forecast {
+    /// Apply the inverse transformations to the given forecast.
+    pub fn inverse_transform_forecast(&self, mut f: Forecast) -> Forecast {
         f.point = self.inverse_transform(f.point.into_iter()).collect();
         if let Some(mut intervals) = f.intervals.take() {
             intervals.lower = self
@@ -371,6 +393,83 @@ trait ExpExt: Iterator<Item = f64> {
 
 impl<T> ExpExt for T where T: Iterator<Item = f64> {}
 
+
+/// Returns the Box-Cox transformation of the given value.
+/// Assumes x > 0.
+pub fn box_cox(x: f64, lambda: f64) -> f64 {
+    if lambda == 0.0 {
+        x.ln()
+    } else {
+        (x.powf(lambda) - 1.0) / lambda
+    }
+}
+
+/// An iterator adapter that applies the Box-Cox transformation to each item.
+#[derive(Clone, Debug)]
+struct BoxCox<T> {
+    inner: T,
+    lambda: f64,
+}
+
+impl<T> Iterator for BoxCox<T>
+where
+    T: Iterator<Item = f64>,
+{
+    type Item = f64;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next().map(|x| box_cox(x, self.lambda))
+    }
+}
+
+trait BoxCoxExt: Iterator<Item = f64> {
+    fn boxcox(self, lambda: f64) -> BoxCox<Self>
+    where
+        Self: Sized,
+    {
+        BoxCox { inner: self, lambda }
+    }
+}
+
+impl<T> BoxCoxExt for T where T: Iterator<Item = f64> {}
+
+/// Returns the inverse Box-Cox transformation of the given value.
+fn inverse_box_cox(y: f64, lambda: f64) -> f64 {
+    if lambda == 0.0 {
+        y.exp()
+    } else {
+        (y * lambda + 1.0).powf(1.0 / lambda)
+    }
+}
+
+/// An iterator adapter that applies the inverse Box-Cox transformation to each item.
+#[derive(Clone, Debug)]
+struct InverseBoxCox<T> {
+    inner: T,
+    lambda: f64,
+}
+
+impl<T> Iterator for InverseBoxCox<T>
+where
+    T: Iterator<Item = f64>,
+{
+    type Item = f64;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next().map(|y| inverse_box_cox(y, self.lambda))
+    }
+}
+
+trait InverseBoxCoxExt: Iterator<Item = f64> {
+    fn inverse_boxcox(self, lambda: f64) -> InverseBoxCox<Self>
+    where
+        Self: Sized,
+    {
+        InverseBoxCox { inner: self, lambda }
+    }
+}
+
+impl<T> InverseBoxCoxExt for T where T: Iterator<Item = f64> {}
+
+
 #[cfg(test)]
 mod test {
     use augurs_testing::{assert_all_close, assert_approx_eq};
@@ -502,4 +601,22 @@ mod test {
         assert_approx_eq!(params.scaled_min, 0.0);
         assert_approx_eq!(params.scaled_max, 1.0);
     }
+
+    #[test]
+    fn boxcox_test() {
+        let data = vec![1.0, 2.0, 3.0];
+        let lambda = 0.5;
+        let expected = vec![0.0, 0.8284271247461903, 1.4641016151377544];
+        let actual: Vec<_> = data.into_iter().boxcox(lambda).collect();
+        assert_all_close(&expected, &actual);
+    }
+
+    #[test]
+    fn inverse_boxcox_test() {
+        let data = vec![0.0, 0.5_f64.ln(), 1.0_f64.ln()];
+        let lambda = 0.5;
+        let expected = vec![1.0, 0.426966072919605, 1.0];
+        let actual: Vec<_> = data.into_iter().inverse_boxcox(lambda).collect();
+        assert_all_close(&expected, &actual);
+    }
 }