nautilus_indicators/average/
lr.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2025 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//   You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::fmt::{Debug, Display};
17
18use nautilus_model::data::Bar;
19
20use crate::indicator::Indicator;
21
22#[repr(C)]
23#[derive(Debug)]
24#[cfg_attr(
25    feature = "python",
26    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.indicators")
27)]
28pub struct LinearRegression {
29    pub period: usize,
30    pub slope: f64,
31    pub intercept: f64,
32    pub degree: f64,
33    pub cfo: f64,
34    pub r2: f64,
35    pub value: f64,
36    pub initialized: bool,
37    has_inputs: bool,
38    inputs: Vec<f64>,
39}
40
41impl Display for LinearRegression {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        write!(f, "{}({})", self.name(), self.period,)
44    }
45}
46
47impl Indicator for LinearRegression {
48    fn name(&self) -> String {
49        stringify!(LinearRegression).to_string()
50    }
51
52    fn has_inputs(&self) -> bool {
53        self.has_inputs
54    }
55
56    fn initialized(&self) -> bool {
57        self.initialized
58    }
59
60    fn handle_bar(&mut self, bar: &Bar) {
61        self.update_raw((&bar.close).into());
62    }
63
64    fn reset(&mut self) {
65        self.slope = 0.0;
66        self.intercept = 0.0;
67        self.degree = 0.0;
68        self.cfo = 0.0;
69        self.r2 = 0.0;
70        self.inputs.clear();
71        self.value = 0.0;
72        self.has_inputs = false;
73        self.initialized = false;
74    }
75}
76
77impl LinearRegression {
78    /// Creates a new [`LinearRegression`] instance.
79    #[must_use]
80    pub fn new(period: usize) -> Self {
81        Self {
82            period,
83            slope: 0.0,
84            intercept: 0.0,
85            degree: 0.0,
86            cfo: 0.0,
87            r2: 0.0,
88            value: 0.0,
89            inputs: Vec::with_capacity(period),
90            has_inputs: false,
91            initialized: false,
92        }
93    }
94
95    pub fn update_raw(&mut self, close: f64) {
96        self.inputs.push(close);
97
98        if !self.initialized {
99            self.has_inputs = true;
100            if self.inputs.len() >= self.period {
101                self.initialized = true;
102            } else {
103                return;
104            }
105        }
106
107        // let x_arr
108        let x_arr: Vec<f64> = (1..=self.period).map(|x| x as f64).collect();
109        let y_arr: Vec<f64> = self.inputs.clone();
110        let x_sum: f64 = 0.5 * self.period as f64 * (self.period as f64 + 1.0);
111        let x_mul_sum: f64 = x_sum * 2.0f64.mul_add(self.period as f64, 1.0) / 3.0;
112        let divisor: f64 = (self.period as f64).mul_add(x_mul_sum, -(x_sum * x_sum));
113        let y_sum: f64 = y_arr.iter().sum::<f64>();
114        let sum_x_y: f64 = x_arr
115            .iter()
116            .zip(y_arr.iter())
117            .map(|(x, y)| x * y)
118            .sum::<f64>();
119
120        self.slope = (self.period as f64).mul_add(sum_x_y, -(x_sum * y_sum)) / divisor;
121        self.intercept = y_sum.mul_add(x_mul_sum, -(x_sum * sum_x_y)) / divisor;
122
123        let residuals: Vec<f64> = x_arr
124            .into_iter()
125            .zip(y_arr.clone())
126            .map(|(x, y)| self.slope.mul_add(x, self.intercept) - y)
127            .collect();
128
129        self.value = residuals.last().unwrap() + y_arr.last().unwrap();
130        self.degree = 180.0 / std::f64::consts::PI * self.slope.atan();
131        self.cfo = 100.0 * residuals.last().unwrap() / y_arr.last().unwrap();
132        let mean: f64 = y_arr.iter().sum::<f64>() / y_arr.len() as f64;
133        self.r2 = 1.0
134            - residuals.iter().map(|r| r * r).sum::<f64>()
135                / y_arr.iter().map(|y| (y - mean) * (y - mean)).sum::<f64>();
136    }
137}
138
139////////////////////////////////////////////////////////////////////////////////
140// Tests
141////////////////////////////////////////////////////////////////////////////////
142#[cfg(test)]
143mod tests {
144    use nautilus_model::data::Bar;
145    use rstest::rstest;
146
147    use crate::{
148        average::lr::LinearRegression,
149        indicator::Indicator,
150        stubs::{bar_ethusdt_binance_minute_bid, indicator_lr_10},
151    };
152
153    #[rstest]
154    fn test_psl_initialized(indicator_lr_10: LinearRegression) {
155        let display_str = format!("{indicator_lr_10}");
156        assert_eq!(display_str, "LinearRegression(10)");
157        assert_eq!(indicator_lr_10.period, 10);
158        assert!(!indicator_lr_10.initialized);
159        assert!(!indicator_lr_10.has_inputs);
160    }
161
162    #[rstest]
163    fn test_value_with_one_input(mut indicator_lr_10: LinearRegression) {
164        indicator_lr_10.update_raw(1.0);
165        assert_eq!(indicator_lr_10.value, 0.0);
166    }
167
168    #[rstest]
169    fn test_value_with_three_inputs(mut indicator_lr_10: LinearRegression) {
170        indicator_lr_10.update_raw(1.0);
171        indicator_lr_10.update_raw(2.0);
172        indicator_lr_10.update_raw(3.0);
173        assert_eq!(indicator_lr_10.value, 0.0);
174    }
175
176    #[rstest]
177    fn test_value_with_ten_inputs(mut indicator_lr_10: LinearRegression) {
178        indicator_lr_10.update_raw(1.00000);
179        indicator_lr_10.update_raw(1.00010);
180        indicator_lr_10.update_raw(1.00030);
181        indicator_lr_10.update_raw(1.00040);
182        indicator_lr_10.update_raw(1.00050);
183        indicator_lr_10.update_raw(1.00060);
184        indicator_lr_10.update_raw(1.00050);
185        indicator_lr_10.update_raw(1.00040);
186        indicator_lr_10.update_raw(1.00030);
187        indicator_lr_10.update_raw(1.00010);
188        indicator_lr_10.update_raw(1.00000);
189        assert_eq!(indicator_lr_10.value, 0.800_307_272_727_272_2);
190    }
191
192    #[rstest]
193    fn test_initialized_with_required_input(mut indicator_lr_10: LinearRegression) {
194        for i in 1..10 {
195            indicator_lr_10.update_raw(f64::from(i));
196        }
197        assert!(!indicator_lr_10.initialized);
198        indicator_lr_10.update_raw(10.0);
199        assert!(indicator_lr_10.initialized);
200    }
201
202    #[rstest]
203    fn test_handle_bar(mut indicator_lr_10: LinearRegression, bar_ethusdt_binance_minute_bid: Bar) {
204        indicator_lr_10.handle_bar(&bar_ethusdt_binance_minute_bid);
205        assert_eq!(indicator_lr_10.value, 0.0);
206        assert!(indicator_lr_10.has_inputs);
207        assert!(!indicator_lr_10.initialized);
208    }
209
210    #[rstest]
211    fn test_reset(mut indicator_lr_10: LinearRegression) {
212        indicator_lr_10.update_raw(1.0);
213        indicator_lr_10.reset();
214        assert_eq!(indicator_lr_10.value, 0.0);
215        assert_eq!(indicator_lr_10.inputs.len(), 0);
216        assert_eq!(indicator_lr_10.slope, 0.0);
217        assert_eq!(indicator_lr_10.intercept, 0.0);
218        assert_eq!(indicator_lr_10.degree, 0.0);
219        assert_eq!(indicator_lr_10.cfo, 0.0);
220        assert_eq!(indicator_lr_10.r2, 0.0);
221        assert!(!indicator_lr_10.has_inputs);
222        assert!(!indicator_lr_10.initialized);
223    }
224}