Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,552 @@
|
||||
//! Time series forecaster wrapper for the demo
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
use rtx_timeseries::models::{FitMetrics as RtxFitMetrics, ForecastOutput, TimeSeriesModel};
|
||||
use rtx_timeseries::{
|
||||
ARIMAModel, ExponentialSmoothingModel, NeuralProphetModel, ProphetModel,
|
||||
TransformerForecastModel,
|
||||
};
|
||||
use timeseries_shared::{
|
||||
DataPoint, FitMetrics, ForecastComponents, ForecastConfig, ForecastResult, ForecasterStatus,
|
||||
ModelType, TimeSeriesData,
|
||||
};
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{debug, info};
|
||||
|
||||
use crate::error::{ForecastError, Result};
|
||||
|
||||
/// Statistics tracking for the forecaster
|
||||
struct ForecasterStats {
|
||||
forecasts_generated: AtomicU64,
|
||||
data_points_processed: AtomicU64,
|
||||
}
|
||||
|
||||
impl ForecasterStats {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
forecasts_generated: AtomicU64::new(0),
|
||||
data_points_processed: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around the different time series model types
|
||||
enum ModelWrapper {
|
||||
Arima(ARIMAModel),
|
||||
// Note: SARIMA is handled via ARIMA with seasonal parameters
|
||||
Prophet(ProphetModel),
|
||||
ExponentialSmoothing(ExponentialSmoothingModel),
|
||||
NeuralProphet(NeuralProphetModel),
|
||||
Transformer(TransformerForecastModel),
|
||||
}
|
||||
|
||||
impl ModelWrapper {
|
||||
fn model_name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Arima(_) => "ARIMA",
|
||||
Self::Prophet(_) => "Prophet",
|
||||
Self::ExponentialSmoothing(_) => "Exponential Smoothing",
|
||||
Self::NeuralProphet(_) => "Neural Prophet",
|
||||
Self::Transformer(_) => "Transformer",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Time series forecaster that wraps various model types
|
||||
pub struct TimeSeriesForecaster {
|
||||
model: RwLock<Option<ModelWrapper>>,
|
||||
config: RwLock<Option<ForecastConfig>>,
|
||||
device: Device,
|
||||
stats: ForecasterStats,
|
||||
training_data: RwLock<Option<TimeSeriesData>>,
|
||||
fit_metrics: RwLock<Option<FitMetrics>>,
|
||||
}
|
||||
|
||||
impl TimeSeriesForecaster {
|
||||
/// Create a new forecaster
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
let device = select_device(true);
|
||||
Self {
|
||||
model: RwLock::new(None),
|
||||
config: RwLock::new(None),
|
||||
device,
|
||||
stats: ForecasterStats::new(),
|
||||
training_data: RwLock::new(None),
|
||||
fit_metrics: RwLock::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the forecaster to initial state
|
||||
pub async fn reset(&self) {
|
||||
*self.model.write().await = None;
|
||||
*self.config.write().await = None;
|
||||
*self.training_data.write().await = None;
|
||||
*self.fit_metrics.write().await = None;
|
||||
}
|
||||
|
||||
/// Fit a model to the provided data
|
||||
pub async fn fit(&self, data: TimeSeriesData, config: ForecastConfig) -> Result<FitMetrics> {
|
||||
info!(
|
||||
"Fitting {:?} model to {} data points",
|
||||
config.model_type,
|
||||
data.data.len()
|
||||
);
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// Convert data to tensors
|
||||
let values: Vec<f32> = data.data.iter().map(|p| p.value as f32).collect();
|
||||
let timestamps: Vec<f32> = data.data.iter().map(|p| p.timestamp as f32).collect();
|
||||
|
||||
let n = values.len();
|
||||
let data_tensor = Tensor::from_vec(values, &[n], &self.device)?;
|
||||
let timestamps_tensor = Tensor::from_vec(timestamps, &[n], &self.device)?;
|
||||
|
||||
// Create and fit the model
|
||||
let model = self.create_model(&config)?;
|
||||
let fit_metrics = self
|
||||
.fit_model(model, &data_tensor, ×tamps_tensor, &config)
|
||||
.await?;
|
||||
|
||||
// Store training data and fit metrics
|
||||
*self.training_data.write().await = Some(data);
|
||||
*self.config.write().await = Some(config);
|
||||
*self.fit_metrics.write().await = Some(fit_metrics.clone());
|
||||
|
||||
// Update stats
|
||||
self.stats
|
||||
.data_points_processed
|
||||
.fetch_add(n as u64, Ordering::Relaxed);
|
||||
|
||||
let elapsed = start.elapsed().as_secs_f64() * 1000.0;
|
||||
debug!("Model fitted in {:.2}ms", elapsed);
|
||||
|
||||
Ok(fit_metrics)
|
||||
}
|
||||
|
||||
/// Generate forecasts using the fitted model
|
||||
pub async fn forecast(&self) -> Result<ForecastResult> {
|
||||
let config_guard = self.config.read().await;
|
||||
let config = config_guard.as_ref().ok_or(ForecastError::NotInitialized)?;
|
||||
|
||||
let horizon = config.horizon;
|
||||
let confidence_level = config.confidence_level;
|
||||
drop(config_guard);
|
||||
|
||||
let model_guard = self.model.read().await;
|
||||
let model = model_guard.as_ref().ok_or(ForecastError::NotInitialized)?;
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// Generate forecast
|
||||
let forecast_output = self
|
||||
.generate_forecast(model, horizon, confidence_level)
|
||||
.await?;
|
||||
|
||||
// Convert to result format
|
||||
let result = self
|
||||
.convert_forecast_output(forecast_output, model.model_name(), start)
|
||||
.await?;
|
||||
|
||||
// Update stats
|
||||
self.stats
|
||||
.forecasts_generated
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Get the current status
|
||||
pub async fn status(&self) -> ForecasterStatus {
|
||||
let model_guard = self.model.read().await;
|
||||
let initialized = model_guard.is_some();
|
||||
let model_name = model_guard
|
||||
.as_ref()
|
||||
.map(|m: &ModelWrapper| m.model_name().to_string());
|
||||
|
||||
ForecasterStatus {
|
||||
initialized,
|
||||
model: model_name,
|
||||
device: self.device.to_string(),
|
||||
forecast_count: self.stats.forecasts_generated.load(Ordering::Relaxed),
|
||||
data_points_processed: self.stats.data_points_processed.load(Ordering::Relaxed),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a model is initialized
|
||||
pub async fn is_initialized(&self) -> bool {
|
||||
self.model.read().await.is_some()
|
||||
}
|
||||
|
||||
fn create_model(&self, config: &ForecastConfig) -> Result<ModelWrapper> {
|
||||
match config.model_type {
|
||||
ModelType::Arima => {
|
||||
let default_arima = timeseries_shared::ArimaConfig::default();
|
||||
let arima_config = config.arima_config.as_ref().unwrap_or(&default_arima);
|
||||
let order = (arima_config.p, arima_config.d, arima_config.q);
|
||||
Ok(ModelWrapper::Arima(ARIMAModel::new(order, None)))
|
||||
}
|
||||
ModelType::Sarima => {
|
||||
// SARIMA is implemented via ARIMA with seasonal parameters
|
||||
let default_sarima = timeseries_shared::SarimaConfig::default();
|
||||
let sarima_config = config.sarima_config.as_ref().unwrap_or(&default_sarima);
|
||||
let order = (
|
||||
sarima_config.order.p,
|
||||
sarima_config.order.d,
|
||||
sarima_config.order.q,
|
||||
);
|
||||
let seasonal = (
|
||||
sarima_config.seasonal_order.p,
|
||||
sarima_config.seasonal_order.d,
|
||||
sarima_config.seasonal_order.q,
|
||||
sarima_config.seasonal_period,
|
||||
);
|
||||
Ok(ModelWrapper::Arima(ARIMAModel::new(order, Some(seasonal))))
|
||||
}
|
||||
ModelType::Prophet => {
|
||||
let default_prophet = timeseries_shared::ProphetConfig::default();
|
||||
let prophet_config = config.prophet_config.as_ref().unwrap_or(&default_prophet);
|
||||
let rtx_config = rtx_timeseries::models::ProphetConfig {
|
||||
growth: match prophet_config.growth.as_str() {
|
||||
"logistic" => rtx_timeseries::models::GrowthType::Logistic,
|
||||
_ => rtx_timeseries::models::GrowthType::Linear,
|
||||
},
|
||||
yearly_seasonality: if prophet_config.yearly_seasonality {
|
||||
rtx_timeseries::models::SeasonalityConfig::auto()
|
||||
} else {
|
||||
rtx_timeseries::models::SeasonalityConfig::disabled()
|
||||
},
|
||||
weekly_seasonality: if prophet_config.weekly_seasonality {
|
||||
rtx_timeseries::models::SeasonalityConfig::auto()
|
||||
} else {
|
||||
rtx_timeseries::models::SeasonalityConfig::disabled()
|
||||
},
|
||||
daily_seasonality: if prophet_config.daily_seasonality {
|
||||
rtx_timeseries::models::SeasonalityConfig::auto()
|
||||
} else {
|
||||
rtx_timeseries::models::SeasonalityConfig::disabled()
|
||||
},
|
||||
changepoint_prior_scale: prophet_config.changepoint_prior_scale,
|
||||
..Default::default()
|
||||
};
|
||||
Ok(ModelWrapper::Prophet(ProphetModel::with_config(rtx_config)))
|
||||
}
|
||||
ModelType::ExponentialSmoothing => Ok(ModelWrapper::ExponentialSmoothing(
|
||||
ExponentialSmoothingModel::new(
|
||||
rtx_timeseries::models::ExponentialSmoothingConfig::default(),
|
||||
),
|
||||
)),
|
||||
ModelType::NeuralProphet => Ok(ModelWrapper::NeuralProphet(NeuralProphetModel::new(
|
||||
rtx_timeseries::models::NeuralProphetConfig::default(),
|
||||
))),
|
||||
ModelType::Transformer => Ok(ModelWrapper::Transformer(TransformerForecastModel::new(
|
||||
rtx_timeseries::models::TransformerForecastConfig::default(),
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn fit_model(
|
||||
&self,
|
||||
model: ModelWrapper,
|
||||
data: &Tensor,
|
||||
timestamps: &Tensor,
|
||||
_config: &ForecastConfig,
|
||||
) -> Result<FitMetrics> {
|
||||
let fit_metrics = match model {
|
||||
ModelWrapper::Arima(mut m) => {
|
||||
TimeSeriesModel::fit(&mut m, data, timestamps).await?;
|
||||
let metrics = m.calculate_fit_metrics(data, timestamps).await?;
|
||||
*self.model.write().await = Some(ModelWrapper::Arima(m));
|
||||
metrics
|
||||
}
|
||||
ModelWrapper::Prophet(mut m) => {
|
||||
TimeSeriesModel::fit(&mut m, data, timestamps).await?;
|
||||
let metrics = m.calculate_fit_metrics(data, timestamps).await?;
|
||||
*self.model.write().await = Some(ModelWrapper::Prophet(m));
|
||||
metrics
|
||||
}
|
||||
ModelWrapper::ExponentialSmoothing(mut m) => {
|
||||
TimeSeriesModel::fit(&mut m, data, timestamps).await?;
|
||||
let metrics = m.calculate_fit_metrics(data, timestamps).await?;
|
||||
*self.model.write().await = Some(ModelWrapper::ExponentialSmoothing(m));
|
||||
metrics
|
||||
}
|
||||
ModelWrapper::NeuralProphet(mut m) => {
|
||||
TimeSeriesModel::fit(&mut m, data, timestamps).await?;
|
||||
let metrics = m.calculate_fit_metrics(data, timestamps).await?;
|
||||
*self.model.write().await = Some(ModelWrapper::NeuralProphet(m));
|
||||
metrics
|
||||
}
|
||||
ModelWrapper::Transformer(mut m) => {
|
||||
TimeSeriesModel::fit(&mut m, data, timestamps).await?;
|
||||
let metrics = m.calculate_fit_metrics(data, timestamps).await?;
|
||||
*self.model.write().await = Some(ModelWrapper::Transformer(m));
|
||||
metrics
|
||||
}
|
||||
};
|
||||
|
||||
Ok(convert_fit_metrics(&fit_metrics))
|
||||
}
|
||||
|
||||
async fn generate_forecast(
|
||||
&self,
|
||||
model: &ModelWrapper,
|
||||
horizon: usize,
|
||||
confidence_level: f64,
|
||||
) -> Result<ForecastOutput> {
|
||||
let output = match model {
|
||||
ModelWrapper::Arima(m) => m.forecast(horizon, confidence_level).await?,
|
||||
ModelWrapper::Prophet(m) => m.forecast(horizon, confidence_level).await?,
|
||||
ModelWrapper::ExponentialSmoothing(m) => m.forecast(horizon, confidence_level).await?,
|
||||
ModelWrapper::NeuralProphet(m) => m.forecast(horizon, confidence_level).await?,
|
||||
ModelWrapper::Transformer(m) => m.forecast(horizon, confidence_level).await?,
|
||||
};
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
async fn convert_forecast_output(
|
||||
&self,
|
||||
output: ForecastOutput,
|
||||
model_name: &str,
|
||||
start: Instant,
|
||||
) -> Result<ForecastResult> {
|
||||
let training_data = self.training_data.read().await;
|
||||
let training_data = training_data
|
||||
.as_ref()
|
||||
.ok_or(ForecastError::NotInitialized)?;
|
||||
|
||||
let fit_metrics = self.fit_metrics.read().await;
|
||||
let metrics = fit_metrics
|
||||
.as_ref()
|
||||
.ok_or(ForecastError::NotInitialized)?
|
||||
.clone();
|
||||
|
||||
// Get the last timestamp from training data
|
||||
let last_timestamp = training_data
|
||||
.data
|
||||
.last()
|
||||
.map_or(0.0, |p| p.timestamp);
|
||||
|
||||
// Convert tensors to DataPoints
|
||||
let mean_vec = output.mean.to_vec()?;
|
||||
let lower_vec = output.lower.to_vec()?;
|
||||
let upper_vec = output.upper.to_vec()?;
|
||||
|
||||
let forecast: Vec<DataPoint> = mean_vec
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| DataPoint {
|
||||
timestamp: last_timestamp + (i + 1) as f64,
|
||||
value: f64::from(v),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let lower_bound: Vec<DataPoint> = lower_vec
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| DataPoint {
|
||||
timestamp: last_timestamp + (i + 1) as f64,
|
||||
value: f64::from(v),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let upper_bound: Vec<DataPoint> = upper_vec
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| DataPoint {
|
||||
timestamp: last_timestamp + (i + 1) as f64,
|
||||
value: f64::from(v),
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Get fitted values for historical data
|
||||
let fitted = training_data.data.clone();
|
||||
|
||||
// Convert components if available
|
||||
let components = output.components.map(|c| ForecastComponents {
|
||||
trend: c.trend.map(|t| {
|
||||
let vec = t.to_vec().unwrap_or_default();
|
||||
vec.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| DataPoint {
|
||||
timestamp: i as f64,
|
||||
value: f64::from(v),
|
||||
})
|
||||
.collect()
|
||||
}),
|
||||
seasonal: c.seasonal.map(|s| {
|
||||
let vec = s.to_vec().unwrap_or_default();
|
||||
vec.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| DataPoint {
|
||||
timestamp: i as f64,
|
||||
value: f64::from(v),
|
||||
})
|
||||
.collect()
|
||||
}),
|
||||
residual: c.noise.map(|r| {
|
||||
let vec = r.to_vec().unwrap_or_default();
|
||||
vec.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| DataPoint {
|
||||
timestamp: i as f64,
|
||||
value: f64::from(v),
|
||||
})
|
||||
.collect()
|
||||
}),
|
||||
});
|
||||
|
||||
let processing_time_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
|
||||
Ok(ForecastResult {
|
||||
forecast,
|
||||
lower_bound,
|
||||
upper_bound,
|
||||
fitted,
|
||||
components,
|
||||
metrics,
|
||||
processing_time_ms,
|
||||
model: model_name.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TimeSeriesForecaster {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
fn select_device(use_gpu: bool) -> Device {
|
||||
if use_gpu
|
||||
&& let Ok(device) = Device::try_default()
|
||||
&& device.is_gpu() {
|
||||
return device;
|
||||
}
|
||||
Device::cpu()
|
||||
}
|
||||
|
||||
fn convert_fit_metrics(metrics: &RtxFitMetrics) -> FitMetrics {
|
||||
FitMetrics {
|
||||
aic: metrics.aic,
|
||||
bic: metrics.bic,
|
||||
mae: metrics.mae,
|
||||
rmse: metrics.rmse,
|
||||
mape: metrics.mape,
|
||||
r_squared: metrics.r_squared,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use timeseries_shared::ArimaConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_forecaster_creation() {
|
||||
let forecaster = TimeSeriesForecaster::new();
|
||||
assert!(!forecaster.is_initialized().await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_forecaster_status() {
|
||||
let forecaster = TimeSeriesForecaster::new();
|
||||
let status = forecaster.status().await;
|
||||
assert!(!status.initialized);
|
||||
assert!(status.model.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_forecast_result_contains_actual_fit_metrics() {
|
||||
let forecaster = TimeSeriesForecaster::new();
|
||||
|
||||
let data = create_test_data();
|
||||
let config = ForecastConfig {
|
||||
model_type: ModelType::Arima,
|
||||
horizon: 5,
|
||||
confidence_level: 0.95,
|
||||
arima_config: Some(ArimaConfig { p: 1, d: 0, q: 1 }),
|
||||
sarima_config: None,
|
||||
prophet_config: None,
|
||||
use_gpu: false,
|
||||
};
|
||||
|
||||
let fit_metrics = forecaster
|
||||
.fit(data, config)
|
||||
.await
|
||||
.expect("Failed to fit model");
|
||||
|
||||
assert!(
|
||||
fit_metrics.mae > 0.0 || fit_metrics.rmse > 0.0,
|
||||
"Fit metrics returned from fit() should contain non-zero values"
|
||||
);
|
||||
|
||||
let forecast_result = forecaster
|
||||
.forecast()
|
||||
.await
|
||||
.expect("Failed to generate forecast");
|
||||
|
||||
assert!(
|
||||
forecast_result.metrics.mae > 0.0 || forecast_result.metrics.mae == 0.0,
|
||||
"Forecast result metrics MAE should be a valid number, got: {}",
|
||||
forecast_result.metrics.mae
|
||||
);
|
||||
assert!(
|
||||
forecast_result.metrics.rmse > 0.0 || forecast_result.metrics.rmse == 0.0,
|
||||
"Forecast result metrics RMSE should be a valid number, got: {}",
|
||||
forecast_result.metrics.rmse
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
forecast_result.metrics.mae, fit_metrics.mae,
|
||||
"Forecast result MAE should match the fit metrics MAE"
|
||||
);
|
||||
assert_eq!(
|
||||
forecast_result.metrics.rmse, fit_metrics.rmse,
|
||||
"Forecast result RMSE should match the fit metrics RMSE"
|
||||
);
|
||||
assert_eq!(
|
||||
forecast_result.metrics.mape, fit_metrics.mape,
|
||||
"Forecast result MAPE should match the fit metrics MAPE"
|
||||
);
|
||||
assert_eq!(
|
||||
forecast_result.metrics.aic, fit_metrics.aic,
|
||||
"Forecast result AIC should match the fit metrics AIC"
|
||||
);
|
||||
assert_eq!(
|
||||
forecast_result.metrics.bic, fit_metrics.bic,
|
||||
"Forecast result BIC should match the fit metrics BIC"
|
||||
);
|
||||
assert_eq!(
|
||||
forecast_result.metrics.r_squared, fit_metrics.r_squared,
|
||||
"Forecast result R² should match the fit metrics R²"
|
||||
);
|
||||
}
|
||||
|
||||
fn create_test_data() -> TimeSeriesData {
|
||||
let values = vec![
|
||||
100.0, 105.0, 110.0, 108.0, 112.0, 115.0, 120.0, 118.0, 125.0, 130.0, 128.0, 135.0,
|
||||
140.0, 138.0, 145.0, 150.0, 148.0, 155.0, 160.0, 158.0, 165.0, 170.0, 168.0, 175.0,
|
||||
180.0, 178.0, 185.0, 190.0, 188.0, 195.0,
|
||||
];
|
||||
|
||||
let data = values
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, value)| DataPoint {
|
||||
timestamp: i as f64,
|
||||
value,
|
||||
})
|
||||
.collect();
|
||||
|
||||
TimeSeriesData {
|
||||
data,
|
||||
name: Some("Test Series".to_string()),
|
||||
frequency: Some("daily".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user