nautilus_indicators/average/
lr.rsuse std::fmt::{Debug, Display};
use nautilus_model::data::Bar;
use crate::indicator::Indicator;
#[repr(C)]
#[derive(Debug)]
#[cfg_attr(
feature = "python",
pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.indicators")
)]
pub struct LinearRegression {
pub period: usize,
pub slope: f64,
pub intercept: f64,
pub degree: f64,
pub cfo: f64,
pub r2: f64,
pub value: f64,
pub initialized: bool,
has_inputs: bool,
inputs: Vec<f64>,
}
impl Display for LinearRegression {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}({})", self.name(), self.period,)
}
}
impl Indicator for LinearRegression {
fn name(&self) -> String {
stringify!(LinearRegression).to_string()
}
fn has_inputs(&self) -> bool {
self.has_inputs
}
fn initialized(&self) -> bool {
self.initialized
}
fn handle_bar(&mut self, bar: &Bar) {
self.update_raw((&bar.close).into());
}
fn reset(&mut self) {
self.slope = 0.0;
self.intercept = 0.0;
self.degree = 0.0;
self.cfo = 0.0;
self.r2 = 0.0;
self.inputs.clear();
self.value = 0.0;
self.has_inputs = false;
self.initialized = false;
}
}
impl LinearRegression {
#[must_use]
pub fn new(period: usize) -> Self {
Self {
period,
slope: 0.0,
intercept: 0.0,
degree: 0.0,
cfo: 0.0,
r2: 0.0,
value: 0.0,
inputs: Vec::with_capacity(period),
has_inputs: false,
initialized: false,
}
}
pub fn update_raw(&mut self, close: f64) {
self.inputs.push(close);
if !self.initialized {
self.has_inputs = true;
if self.inputs.len() >= self.period {
self.initialized = true;
} else {
return;
}
}
let x_arr: Vec<f64> = (1..=self.period).map(|x| x as f64).collect();
let y_arr: Vec<f64> = self.inputs.clone();
let x_sum: f64 = 0.5 * self.period as f64 * (self.period as f64 + 1.0);
let x_mul_sum: f64 = x_sum * 2.0f64.mul_add(self.period as f64, 1.0) / 3.0;
let divisor: f64 = (self.period as f64).mul_add(x_mul_sum, -(x_sum * x_sum));
let y_sum: f64 = y_arr.iter().sum::<f64>();
let sum_x_y: f64 = x_arr
.iter()
.zip(y_arr.iter())
.map(|(x, y)| x * y)
.sum::<f64>();
self.slope = (self.period as f64).mul_add(sum_x_y, -(x_sum * y_sum)) / divisor;
self.intercept = y_sum.mul_add(x_mul_sum, -(x_sum * sum_x_y)) / divisor;
let residuals: Vec<f64> = x_arr
.into_iter()
.zip(y_arr.clone())
.map(|(x, y)| self.slope.mul_add(x, self.intercept) - y)
.collect();
self.value = residuals.last().unwrap() + y_arr.last().unwrap();
self.degree = 180.0 / std::f64::consts::PI * self.slope.atan();
self.cfo = 100.0 * residuals.last().unwrap() / y_arr.last().unwrap();
let mean: f64 = y_arr.iter().sum::<f64>() / y_arr.len() as f64;
self.r2 = 1.0
- residuals.iter().map(|r| r * r).sum::<f64>()
/ y_arr.iter().map(|y| (y - mean) * (y - mean)).sum::<f64>();
}
}
#[cfg(test)]
mod tests {
use nautilus_model::data::Bar;
use rstest::rstest;
use crate::{
average::lr::LinearRegression,
indicator::Indicator,
stubs::{bar_ethusdt_binance_minute_bid, indicator_lr_10},
};
#[rstest]
fn test_psl_initialized(indicator_lr_10: LinearRegression) {
let display_str = format!("{indicator_lr_10}");
assert_eq!(display_str, "LinearRegression(10)");
assert_eq!(indicator_lr_10.period, 10);
assert!(!indicator_lr_10.initialized);
assert!(!indicator_lr_10.has_inputs);
}
#[rstest]
fn test_value_with_one_input(mut indicator_lr_10: LinearRegression) {
indicator_lr_10.update_raw(1.0);
assert_eq!(indicator_lr_10.value, 0.0);
}
#[rstest]
fn test_value_with_three_inputs(mut indicator_lr_10: LinearRegression) {
indicator_lr_10.update_raw(1.0);
indicator_lr_10.update_raw(2.0);
indicator_lr_10.update_raw(3.0);
assert_eq!(indicator_lr_10.value, 0.0);
}
#[rstest]
fn test_value_with_ten_inputs(mut indicator_lr_10: LinearRegression) {
indicator_lr_10.update_raw(1.00000);
indicator_lr_10.update_raw(1.00010);
indicator_lr_10.update_raw(1.00030);
indicator_lr_10.update_raw(1.00040);
indicator_lr_10.update_raw(1.00050);
indicator_lr_10.update_raw(1.00060);
indicator_lr_10.update_raw(1.00050);
indicator_lr_10.update_raw(1.00040);
indicator_lr_10.update_raw(1.00030);
indicator_lr_10.update_raw(1.00010);
indicator_lr_10.update_raw(1.00000);
assert_eq!(indicator_lr_10.value, 0.800_307_272_727_272_2);
}
#[rstest]
fn test_initialized_with_required_input(mut indicator_lr_10: LinearRegression) {
for i in 1..10 {
indicator_lr_10.update_raw(f64::from(i));
}
assert!(!indicator_lr_10.initialized);
indicator_lr_10.update_raw(10.0);
assert!(indicator_lr_10.initialized);
}
#[rstest]
fn test_handle_bar(mut indicator_lr_10: LinearRegression, bar_ethusdt_binance_minute_bid: Bar) {
indicator_lr_10.handle_bar(&bar_ethusdt_binance_minute_bid);
assert_eq!(indicator_lr_10.value, 0.0);
assert!(indicator_lr_10.has_inputs);
assert!(!indicator_lr_10.initialized);
}
#[rstest]
fn test_reset(mut indicator_lr_10: LinearRegression) {
indicator_lr_10.update_raw(1.0);
indicator_lr_10.reset();
assert_eq!(indicator_lr_10.value, 0.0);
assert_eq!(indicator_lr_10.inputs.len(), 0);
assert_eq!(indicator_lr_10.slope, 0.0);
assert_eq!(indicator_lr_10.intercept, 0.0);
assert_eq!(indicator_lr_10.degree, 0.0);
assert_eq!(indicator_lr_10.cfo, 0.0);
assert_eq!(indicator_lr_10.r2, 0.0);
assert!(!indicator_lr_10.has_inputs);
assert!(!indicator_lr_10.initialized);
}
}