diff --git a/src/value_estimation_team/indicators/linear_regression.rs b/src/value_estimation_team/indicators/linear_regression.rs index 2009589..87861f0 100644 --- a/src/value_estimation_team/indicators/linear_regression.rs +++ b/src/value_estimation_team/indicators/linear_regression.rs @@ -14,12 +14,14 @@ use tokio::{fs::*, io::AsyncWriteExt, sync::Mutex, time::*}; #[derive(Clone, Debug)] pub struct LrData { pub lr_value: f64, // linear regression value + pub r_squared: f64, pub close_time: i64, } impl LrData { fn new() -> LrData { let a = LrData { lr_value: 0.0, + r_squared: 0.0, close_time: 0, }; a @@ -49,10 +51,10 @@ pub async fn linear_regression( if vec.len() >= length { let rt_price_data = vec.clone(); task_vec.push(tokio::spawn(async move { - // Calculate prediction of linear regression let mut lr_data_vec: Vec = Vec::new(); for window in rt_price_data.windows(length) { + // Calculate prediction of linear regression let mut lr_data = LrData::new(); let x: Vec = (0..length).map(|x| x as f64).collect(); let y: Vec = window.iter().map(|x| x.close_price).collect(); @@ -67,7 +69,18 @@ pub async fn linear_regression( let intercept = y_mean - slope * x_mean; let linreg = intercept + slope * (length as f64 - 1.0 - offset as f64); + + // Calculate R-Squared + let ss_tot: f64 = y.iter().map(|y_i| (y_i - y_mean).powi(2)).sum(); + let ss_res: f64 = y.iter().zip(x.iter()).map(|(y_i, x_i)| { + let y_pred = intercept + slope * x_i; + (y_i - y_pred).powi(2) + }).sum(); + let r_squared = 1.0 - (ss_res / ss_tot); + let r_squared_rounded = (r_squared * 100.0).round() / 100.0; + lr_data.lr_value = linreg; + lr_data.r_squared = r_squared_rounded; lr_data.close_time = window.last().unwrap().close_time; lr_data_vec.push(lr_data.clone()); }