Add r_squared value

This commit is contained in:
Sik Yoon 2024-06-06 03:58:04 +09:00
parent b4dbf9e1d2
commit 1e7be2681f

View File

@ -14,12 +14,14 @@ use tokio::{fs::*, io::AsyncWriteExt, sync::Mutex, time::*};
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct LrData { pub struct LrData {
pub lr_value: f64, // linear regression value pub lr_value: f64, // linear regression value
pub r_squared: f64,
pub close_time: i64, pub close_time: i64,
} }
impl LrData { impl LrData {
fn new() -> LrData { fn new() -> LrData {
let a = LrData { let a = LrData {
lr_value: 0.0, lr_value: 0.0,
r_squared: 0.0,
close_time: 0, close_time: 0,
}; };
a a
@ -49,10 +51,10 @@ pub async fn linear_regression(
if vec.len() >= length { if vec.len() >= length {
let rt_price_data = vec.clone(); let rt_price_data = vec.clone();
task_vec.push(tokio::spawn(async move { task_vec.push(tokio::spawn(async move {
// Calculate prediction of linear regression
let mut lr_data_vec: Vec<LrData> = Vec::new(); let mut lr_data_vec: Vec<LrData> = Vec::new();
for window in rt_price_data.windows(length) { for window in rt_price_data.windows(length) {
// Calculate prediction of linear regression
let mut lr_data = LrData::new(); let mut lr_data = LrData::new();
let x: Vec<f64> = (0..length).map(|x| x as f64).collect(); let x: Vec<f64> = (0..length).map(|x| x as f64).collect();
let y: Vec<f64> = window.iter().map(|x| x.close_price).collect(); let y: Vec<f64> = 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 intercept = y_mean - slope * x_mean;
let linreg = intercept + slope * (length as f64 - 1.0 - offset as f64); 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.lr_value = linreg;
lr_data.r_squared = r_squared_rounded;
lr_data.close_time = window.last().unwrap().close_time; lr_data.close_time = window.last().unwrap().close_time;
lr_data_vec.push(lr_data.clone()); lr_data_vec.push(lr_data.clone());
} }