diff --git a/crates/augurs-forecaster/Cargo.toml b/crates/augurs-forecaster/Cargo.toml index e48c7ebf..73f341f5 100644 --- a/crates/augurs-forecaster/Cargo.toml +++ b/crates/augurs-forecaster/Cargo.toml @@ -13,12 +13,13 @@ description = "A high-level API for the augurs forecasting library." bench = false [dependencies] +argmin = "0.10.0" augurs-core.workspace = true itertools.workspace = true thiserror.workspace = true [dev-dependencies] -augurs.workspace = true +augurs = { workspace = true, features = ["mstl", "ets", "forecaster"]} augurs-testing.workspace = true [lints] diff --git a/crates/augurs-forecaster/src/forecaster.rs b/crates/augurs-forecaster/src/forecaster.rs index 6c27b0f3..5ad62957 100644 --- a/crates/augurs-forecaster/src/forecaster.rs +++ b/crates/augurs-forecaster/src/forecaster.rs @@ -124,4 +124,46 @@ mod test { let forecasts = forecaster.predict(4, None).unwrap(); assert_all_approx_eq(&forecasts.point, &[5.0, 5.0, 5.0, 5.0]); } + + #[test] + fn test_forecaster_power_positive() { + let data = &[1.0_f64, 2.0, 3.0, 4.0, 5.0]; + let got = Transform::power_transform(data); + assert!(got.is_ok()); + let transforms = vec![got.unwrap()]; + let model = MSTLModel::new(vec![2], NaiveTrend::new()); + let mut forecaster = Forecaster::new(model).with_transforms(transforms); + forecaster.fit(data).unwrap(); + let forecasts = forecaster.predict(4, None).unwrap(); + assert_all_approx_eq( + &forecasts.point, + &[ + 5.084499064884572, + 5.000000030329821, + 5.084499064884572, + 5.000000030329821, + ], + ); + } + + #[test] + fn test_forecaster_power_non_positive() { + let data = &[0.0, 2.0, 3.0, 4.0, 5.0]; + let got = Transform::power_transform(data); + assert!(got.is_ok()); + let transforms = vec![got.unwrap()]; + let model = MSTLModel::new(vec![2], NaiveTrend::new()); + let mut forecaster = Forecaster::new(model).with_transforms(transforms); + forecaster.fit(data).unwrap(); + let forecasts = forecaster.predict(4, None).unwrap(); + assert_all_approx_eq( + &forecasts.point, + &[ + 5.205557727170964, + 5.000000132803496, + 5.205557727170964, + 5.000000132803496, + ], + ); + } } diff --git a/crates/augurs-forecaster/src/lib.rs b/crates/augurs-forecaster/src/lib.rs index d8c2c689..7628d550 100644 --- a/crates/augurs-forecaster/src/lib.rs +++ b/crates/augurs-forecaster/src/lib.rs @@ -3,6 +3,7 @@ mod data; mod error; mod forecaster; +mod power_transforms; pub mod transforms; pub use data::Data; diff --git a/crates/augurs-forecaster/src/power_transforms.rs b/crates/augurs-forecaster/src/power_transforms.rs new file mode 100644 index 00000000..ee99f828 --- /dev/null +++ b/crates/augurs-forecaster/src/power_transforms.rs @@ -0,0 +1,222 @@ +use crate::transforms::box_cox; +use crate::transforms::yeo_johnson; +use argmin::core::{CostFunction, Error, Executor}; +use argmin::solver::brent::BrentOpt; + +fn box_cox_log_likelihood(data: &[f64], lambda: f64) -> Result { + let n = data.len() as f64; + if n == 0.0 { + return Err(Error::msg("Data must not be empty")); + } + if data.iter().any(|&x| x <= 0.0) { + return Err(Error::msg("All data must be greater than 0")); + } + let transformed_data: Result, _> = data.iter().map(|&x| box_cox(x, lambda)).collect(); + + let transformed_data = match transformed_data { + Ok(values) => values, + Err(e) => return Err(Error::msg(e)), + }; + let mean_transformed: f64 = transformed_data.iter().copied().sum::() / n; + let variance: f64 = transformed_data + .iter() + .map(|&x| (x - mean_transformed).powi(2)) + .sum::() + / n; + + // Avoid log(0) by ensuring variance is positive + if variance <= 0.0 { + return Err(Error::msg("Variance must be positive")); + } + let log_likelihood = + -0.5 * n * variance.ln() + (lambda - 1.0) * data.iter().map(|&x| x.ln()).sum::(); + Ok(log_likelihood) +} + +fn yeo_johnson_log_likelihood(data: &[f64], lambda: f64) -> Result { + let n = data.len() as f64; + + if n == 0.0 { + return Err(Error::msg("Data array is empty")); + } + + let transformed_data: Result, _> = + data.iter().map(|&x| yeo_johnson(x, lambda)).collect(); + + let transformed_data = match transformed_data { + Ok(values) => values, + Err(e) => return Err(Error::msg(e)), + }; + + let mean = transformed_data.iter().sum::() / n; + + let variance = transformed_data + .iter() + .map(|&x| (x - mean).powi(2)) + .sum::() + / n; + + if variance <= 0.0 { + return Err(Error::msg("Variance is non-positive")); + } + + let log_sigma_squared = variance.ln(); + let log_likelihood = -n / 2.0 * log_sigma_squared; + + let additional_term: f64 = data + .iter() + .map(|&x| (x.signum() * (x.abs() + 1.0).ln())) + .sum::() + * (lambda - 1.0); + + Ok(log_likelihood + additional_term) +} + +#[derive(Clone)] +struct BoxCoxProblem<'a> { + data: &'a [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 { + box_cox_log_likelihood(self.data, *lambda).map(|ll| -ll) + } +} + +#[derive(Clone)] +struct YeoJohnsonProblem<'a> { + data: &'a [f64], +} + +impl CostFunction for YeoJohnsonProblem<'_> { + type Param = f64; + type Output = f64; + + // The goal is to minimize the negative log-likelihood + fn cost(&self, lambda: &Self::Param) -> Result { + yeo_johnson_log_likelihood(self.data, *lambda).map(|ll| -ll) + } +} + +struct OptimizationParams { + initial_param: f64, + lower_bound: f64, + upper_bound: f64, + max_iterations: u64, +} + +impl Default for OptimizationParams { + fn default() -> Self { + Self { + initial_param: 0.0, + lower_bound: -2.0, + upper_bound: 2.0, + max_iterations: 1000, + } + } +} + +fn optimize_lambda>( + cost: T, + params: OptimizationParams, +) -> Result { + let solver = BrentOpt::new(params.lower_bound, params.upper_bound); + let result = Executor::new(cost, solver) + .configure(|state| { + state + .param(params.initial_param) + .max_iters(params.max_iterations) + }) + .run(); + + result.and_then(|res| { + res.state() + .best_param + .ok_or_else(|| Error::msg("No best parameter found")) + }) +} + +/// Optimize the lambda parameter for the Box-Cox or Yeo-Johnson transformation +pub(crate) fn optimize_box_cox_lambda(data: &[f64]) -> Result { + // Use Box-Cox transformation + let cost = BoxCoxProblem { data }; + let optimization_params = OptimizationParams::default(); + optimize_lambda(cost, optimization_params) +} + +pub(crate) fn optimize_yeo_johnson_lambda(data: &[f64]) -> Result { + // Use Yeo-Johnson transformation + let cost = YeoJohnsonProblem { data }; + let optimization_params = OptimizationParams::default(); + optimize_lambda(cost, optimization_params) +} + +#[cfg(test)] +mod test { + use super::*; + use augurs_testing::assert_approx_eq; + + #[test] + fn correct_optimal_box_cox_lambda() { + let data = &[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]; + let got = optimize_box_cox_lambda(data); + assert!(got.is_ok()); + let lambda = got.unwrap(); + assert_approx_eq!(lambda, 0.7123778635679304); + } + + #[test] + fn optimal_box_cox_lambda_lambda_empty_data() { + let data = &[]; + let got = optimize_box_cox_lambda(data); + assert!(got.is_err()); + } + + #[test] + fn optimal_box_cox_lambda_non_positive_data() { + let data = &[0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]; + let got = optimize_box_cox_lambda(data); + assert!(got.is_err()); + } + + #[test] + fn correct_optimal_yeo_johnson_lambda() { + let data = &[0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]; + let got = optimize_yeo_johnson_lambda(data); + assert!(got.is_ok()); + let lambda = got.unwrap(); + assert_approx_eq!(lambda, 1.7458442076987954); + } + + #[test] + fn test_box_cox_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!(got.is_ok()); + let llf = got.unwrap(); + assert_approx_eq!(llf, 11.266065387038703); + } + + #[test] + fn test_box_cox_llf_non_positive() { + let data = &[0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]; + let lambda = 0.0; + let got = box_cox_log_likelihood(data, lambda); + assert!(got.is_err()); + } + + #[test] + fn test_yeo_johnson_llf() { + let data = &[0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]; + let lambda = 1.0; + let got = yeo_johnson_log_likelihood(data, lambda); + assert!(got.is_ok()); + let llf = got.unwrap(); + assert_approx_eq!(llf, 10.499377905819307); + } +} diff --git a/crates/augurs-forecaster/src/transforms.rs b/crates/augurs-forecaster/src/transforms.rs index 50b42be6..fe232a73 100644 --- a/crates/augurs-forecaster/src/transforms.rs +++ b/crates/augurs-forecaster/src/transforms.rs @@ -2,19 +2,28 @@ Data transformations. */ +use argmin::core::Error; use augurs_core::{ interpolate::{InterpolateExt, LinearInterpolator}, Forecast, }; +use crate::power_transforms::{optimize_box_cox_lambda, optimize_yeo_johnson_lambda}; + +/// 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); impl Transforms { + /// create a new `Transforms` instance with the given transforms. pub(crate) fn new(transforms: Vec) -> Self { Self(transforms) } + /// Apply the transformations to the given time series. pub(crate) fn transform<'a, T>(&'a self, input: T) -> Box + 'a> where T: Iterator + 'a, @@ -24,6 +33,7 @@ impl Transforms { .fold(Box::new(input) as _, |y, t| t.transform(y)) } + /// Apply the inverse transformations to the given forecast. pub(crate) fn inverse_transform(&self, forecast: Forecast) -> Forecast { self.0 .iter() @@ -46,6 +56,20 @@ pub enum Transform { Logit, /// Log transform. Log, + /// Box-Cox transform. + BoxCox { + /// The lambda parameter for the Box-Cox transformation. + /// If lambda == 0, the transformation is equivalent to the natural logarithm. + /// Otherwise, the transformation is (x^lambda - 1) / lambda. + lambda: f64, + }, + /// Yeo-Johnson transform. + YeoJohnson { + /// The lambda parameter for the Yeo-Johnson transformation. + /// If lambda == 0, the transformation is equivalent to the natural logarithm. + /// Otherwise, the transformation is ((x + 1)^lambda - 1) / lambda. + lambda: f64, + }, } impl Transform { @@ -81,6 +105,52 @@ impl Transform { Self::Log } + /// Create a new Box-Cox transform. + /// + /// This transform applies the Box-Cox transformation to each item. + /// + /// The Box-Cox transformation is defined as: + /// + /// - if lambda == 0: x.ln() + /// - otherwise: (x^lambda - 1) / lambda + pub fn boxcox(lambda: f64) -> Self { + Self::BoxCox { lambda } + } + + /// Create a new Yeo-Johnson transform. + /// + /// This transform applies the Yeo-Johnson transformation to each item. + /// + /// The Yeo-Johnson transformation is a generalization of the Box-Cox transformation that + /// supports negative values. It is defined as: + /// + /// - if lambda != 0 and x >= 0: ((x + 1)^lambda - 1) / lambda + /// - if lambda == 0 and x >= 0: (x + 1).ln() + /// - if lambda != 2 and x < 0: ((-x + 1)^2 - 1) / 2 + /// - if lambda == 2 and x < 0: (-x + 1).ln() + pub fn yeo_johnson(lambda: f64) -> Self { + Self::YeoJohnson { lambda } + } + + /// Create a power transform that optimizes the lambda parameter. + /// + /// # Algorithm Selection + /// + /// - If all values are positive: Uses Box-Cox transformation + /// - If any values are negative or zero: Uses Yeo-Johnson transformation + /// + /// # Returns + /// + /// Returns `Result` to handle optimization failures gracefully + pub fn power_transform(data: &[f64]) -> Result { + if data.iter().all(|&x| x > 0.0) { + optimize_box_cox_lambda(data).map(|lambda| Self::BoxCox { lambda }) + } else { + optimize_yeo_johnson_lambda(data).map(|lambda| Self::YeoJohnson { lambda }) + } + } + + /// Apply the transformation to the given time series. pub(crate) fn transform<'a, T>(&'a self, input: T) -> Box + 'a> where T: Iterator + 'a, @@ -90,9 +160,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.box_cox(*lambda)), + Self::YeoJohnson { lambda } => Box::new(input.yeo_johnson(*lambda)), } } + /// Apply the inverse transformation to the given time series. pub(crate) fn inverse_transform<'a, T>(&'a self, input: T) -> Box + 'a> where T: Iterator + 'a, @@ -102,10 +175,13 @@ 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_box_cox(*lambda)), + Self::YeoJohnson { lambda } => Box::new(input.inverse_yeo_johnson(*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 +447,203 @@ trait ExpExt: Iterator { impl ExpExt for T where T: Iterator {} +/// Returns the Box-Cox transformation of the given value. +/// Assumes x > 0. +pub fn box_cox(x: f64, lambda: f64) -> Result { + if x <= 0.0 { + return Err("x must be greater than 0"); + } + if lambda == 0.0 { + Ok(x.ln()) + } else { + Ok((x.powf(lambda) - 1.0) / lambda) + } +} + +/// Returns the Yeo-Johnson transformation of the given value. +pub fn yeo_johnson(x: f64, lambda: f64) -> Result { + if x.is_nan() || lambda.is_nan() { + return Err("Input values must be valid numbers."); + } + + if x >= 0.0 { + if lambda == 0.0 { + Ok((x + 1.0).ln()) + } else { + Ok(((x + 1.0).powf(lambda) - 1.0) / lambda) + } + } else if lambda == 2.0 { + Ok(-(-x + 1.0).ln()) + } else { + Ok(-((-x + 1.0).powf(2.0 - lambda) - 1.0) / (2.0 - lambda)) + } +} + +/// An iterator adapter that applies the Box-Cox transformation to each item. +#[derive(Clone, Debug)] +struct BoxCox { + inner: T, + lambda: f64, +} + +impl Iterator for BoxCox +where + T: Iterator, +{ + type Item = f64; + fn next(&mut self) -> Option { + self.inner + .next() + .map(|x| box_cox(x, self.lambda).unwrap_or(f64::NAN)) + } +} + +trait BoxCoxExt: Iterator { + fn box_cox(self, lambda: f64) -> BoxCox + where + Self: Sized, + { + BoxCox { + inner: self, + lambda, + } + } +} + +impl BoxCoxExt for T where T: Iterator {} + +/// Returns the inverse Box-Cox transformation of the given value. +fn inverse_box_cox(y: f64, lambda: f64) -> Result { + if lambda == 0.0 { + Ok(y.exp()) + } else { + let value = y * lambda + 1.0; + if value <= 0.0 { + Err("Invalid domain for inverse Box-Cox transformation") + } else { + Ok(value.powf(1.0 / lambda)) + } + } +} + +/// An iterator adapter that applies the inverse Box-Cox transformation to each item. +#[derive(Clone, Debug)] +struct InverseBoxCox { + inner: T, + lambda: f64, +} + +impl Iterator for InverseBoxCox +where + T: Iterator, +{ + type Item = f64; + fn next(&mut self) -> Option { + self.inner + .next() + .map(|y| inverse_box_cox(y, self.lambda).unwrap_or(f64::NAN)) + } +} + +trait InverseBoxCoxExt: Iterator { + fn inverse_box_cox(self, lambda: f64) -> InverseBoxCox + where + Self: Sized, + { + InverseBoxCox { + inner: self, + lambda, + } + } +} + +impl InverseBoxCoxExt for T where T: Iterator {} + +/// An iterator adapter that applies the Yeo-Johnson transformation to each item. +#[derive(Clone, Debug)] +struct YeoJohnson { + inner: T, + lambda: f64, +} + +impl Iterator for YeoJohnson +where + T: Iterator, +{ + type Item = f64; + fn next(&mut self) -> Option { + self.inner + .next() + .map(|x| yeo_johnson(x, self.lambda).unwrap_or(f64::NAN)) + } +} + +trait YeoJohnsonExt: Iterator { + fn yeo_johnson(self, lambda: f64) -> YeoJohnson + where + Self: Sized, + { + YeoJohnson { + inner: self, + lambda, + } + } +} + +impl YeoJohnsonExt for T where T: Iterator {} + +/// Returns the inverse Yeo-Johnson transformation of the given value. +fn inverse_yeo_johnson(y: f64, lambda: f64) -> f64 { + const EPSILON: f64 = 1e-6; + + if y >= 0.0 && lambda.abs() < EPSILON { + // For lambda close to 0 (positive values) + (y.exp()) - 1.0 + } else if y >= 0.0 { + // For positive values (lambda not close to 0) + (y * lambda + 1.0).powf(1.0 / lambda) - 1.0 + } else if (lambda - 2.0).abs() < EPSILON { + // For lambda close to 2 (negative values) + -(-y.exp() - 1.0) + } else { + // For negative values (lambda not close to 2) + -((-((2.0 - lambda) * y) + 1.0).powf(1.0 / (2.0 - lambda)) - 1.0) + } +} + +/// An iterator adapter that applies the inverse Yeo-Johnson transformation to each item. +#[derive(Clone, Debug)] +struct InverseYeoJohnson { + inner: T, + lambda: f64, +} + +impl Iterator for InverseYeoJohnson +where + T: Iterator, +{ + type Item = f64; + fn next(&mut self) -> Option { + self.inner + .next() + .map(|y| inverse_yeo_johnson(y, self.lambda)) + } +} + +trait InverseYeoJohnsonExt: Iterator { + fn inverse_yeo_johnson(self, lambda: f64) -> InverseYeoJohnson + where + Self: Sized, + { + InverseYeoJohnson { + inner: self, + lambda, + } + } +} + +impl InverseYeoJohnsonExt for T where T: Iterator {} + #[cfg(test)] mod test { use augurs_testing::{assert_all_close, assert_approx_eq}; @@ -502,4 +775,40 @@ mod test { assert_approx_eq!(params.scaled_min, 0.0); assert_approx_eq!(params.scaled_max, 1.0); } + + #[test] + fn box_cox_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().box_cox(lambda).collect(); + assert_all_close(&expected, &actual); + } + + #[test] + fn inverse_box_cox_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_box_cox(lambda).collect(); + assert_all_close(&expected, &actual); + } + + #[test] + fn yeo_johnson_test() { + let data = vec![-1.0, 0.0, 1.0]; + let lambda = 0.5; + let expected = vec![-1.2189514164974602, 0.0, 0.8284271247461903]; + let actual: Vec<_> = data.into_iter().yeo_johnson(lambda).collect(); + assert_all_close(&expected, &actual); + } + + #[test] + fn inverse_yeo_johnson_test() { + let data = vec![-1.2189514164974602, 0.0, 0.8284271247461903]; + let lambda = 0.5; + let expected = vec![-1.0, 0.0, 1.0]; + let actual: Vec<_> = data.into_iter().inverse_yeo_johnson(lambda).collect(); + assert_all_close(&expected, &actual); + } }