ta_lib_in_rust/indicators/math/
mod.rs

1use polars::prelude::*;
2
3/// Vector arithmetic addition
4///
5/// # Arguments
6///
7/// * `df` - DataFrame containing the data
8/// * `col1` - First column name
9/// * `col2` - Second column name
10///
11/// # Returns
12///
13/// Returns a PolarsResult containing the addition Series
14pub fn calculate_add(df: &DataFrame, col1: &str, col2: &str) -> PolarsResult<Series> {
15    if !df.schema().contains(col1) || !df.schema().contains(col2) {
16        return Err(PolarsError::ComputeError(
17            format!("Addition requires both {col1} and {col2} columns").into()
18        ));
19    }
20    
21    let series1 = df.column(col1)?.f64()?;
22    let series2 = df.column(col2)?.f64()?;
23    
24    let result = series1 + series2;
25    
26    Ok(result.with_name(format!("{col1}_add_{col2}").into()).into())
27}
28
29/// Vector arithmetic subtraction
30///
31/// # Arguments
32///
33/// * `df` - DataFrame containing the data
34/// * `col1` - First column name (minuend)
35/// * `col2` - Second column name (subtrahend)
36///
37/// # Returns
38///
39/// Returns a PolarsResult containing the subtraction Series
40pub fn calculate_sub(df: &DataFrame, col1: &str, col2: &str) -> PolarsResult<Series> {
41    if !df.schema().contains(col1) || !df.schema().contains(col2) {
42        return Err(PolarsError::ComputeError(
43            format!("Subtraction requires both {col1} and {col2} columns").into()
44        ));
45    }
46    
47    let series1 = df.column(col1)?.f64()?;
48    let series2 = df.column(col2)?.f64()?;
49    
50    let result = series1 - series2;
51    
52    Ok(result.with_name(format!("{col1}_sub_{col2}").into()).into())
53}
54
55/// Vector arithmetic multiplication
56///
57/// # Arguments
58///
59/// * `df` - DataFrame containing the data
60/// * `col1` - First column name
61/// * `col2` - Second column name
62///
63/// # Returns
64///
65/// Returns a PolarsResult containing the multiplication Series
66pub fn calculate_mult(df: &DataFrame, col1: &str, col2: &str) -> PolarsResult<Series> {
67    if !df.schema().contains(col1) || !df.schema().contains(col2) {
68        return Err(PolarsError::ComputeError(
69            format!("Multiplication requires both {col1} and {col2} columns").into()
70        ));
71    }
72    
73    let series1 = df.column(col1)?.f64()?;
74    let series2 = df.column(col2)?.f64()?;
75    
76    let result = series1 * series2;
77    
78    Ok(result.with_name(format!("{col1}_mult_{col2}").into()).into())
79}
80
81/// Vector arithmetic division
82///
83/// # Arguments
84///
85/// * `df` - DataFrame containing the data
86/// * `col1` - First column name (numerator)
87/// * `col2` - Second column name (denominator)
88///
89/// # Returns
90///
91/// Returns a PolarsResult containing the division Series
92pub fn calculate_div(df: &DataFrame, col1: &str, col2: &str) -> PolarsResult<Series> {
93    if !df.schema().contains(col1) || !df.schema().contains(col2) {
94        return Err(PolarsError::ComputeError(
95            format!("Division requires both {col1} and {col2} columns").into()
96        ));
97    }
98    
99    let series1 = df.column(col1)?.f64()?;
100    let series2 = df.column(col2)?.f64()?;
101    
102    // Replace zeros with NaN to avoid division by zero
103    let mut div_values = Vec::with_capacity(df.height());
104    
105    for i in 0..df.height() {
106        let num = series1.get(i).unwrap_or(f64::NAN);
107        let denom = series2.get(i).unwrap_or(f64::NAN);
108        
109        if denom != 0.0 && !denom.is_nan() && !num.is_nan() {
110            div_values.push(num / denom);
111        } else {
112            div_values.push(f64::NAN);
113        }
114    }
115    
116    Ok(Series::new(format!("{col1}_div_{col2}").into(), div_values))
117}
118
119/// Find maximum value over a specified window
120///
121/// # Arguments
122///
123/// * `df` - DataFrame containing the data
124/// * `column` - Column name to calculate on
125/// * `window` - Window size for the calculation
126///
127/// # Returns
128///
129/// Returns a PolarsResult containing the MAX Series
130pub fn calculate_max(df: &DataFrame, column: &str, window: usize) -> PolarsResult<Series> {
131    if !df.schema().contains(column) {
132        return Err(PolarsError::ComputeError(
133            format!("MAX calculation requires {column} column").into()
134        ));
135    }
136    
137    let series = df.column(column)?.f64()?;
138    
139    let mut max_values = Vec::with_capacity(df.height());
140    
141    // Fill initial values with NaN
142    for i in 0..window-1 {
143        max_values.push(f64::NAN);
144    }
145    
146    // Calculate max for each window
147    for i in window-1..df.height() {
148        let mut max_val = f64::NEG_INFINITY;
149        let mut all_nan = true;
150        
151        for j in 0..window {
152            let val = series.get(i - j).unwrap_or(f64::NAN);
153            if !val.is_nan() {
154                max_val = max_val.max(val);
155                all_nan = false;
156            }
157        }
158        
159        if all_nan {
160            max_values.push(f64::NAN);
161        } else {
162            max_values.push(max_val);
163        }
164    }
165    
166    Ok(Series::new(format!("{column}_max_{window}").into(), max_values))
167}
168
169/// Find minimum value over a specified window
170///
171/// # Arguments
172///
173/// * `df` - DataFrame containing the data
174/// * `column` - Column name to calculate on
175/// * `window` - Window size for the calculation
176///
177/// # Returns
178///
179/// Returns a PolarsResult containing the MIN Series
180pub fn calculate_min(df: &DataFrame, column: &str, window: usize) -> PolarsResult<Series> {
181    if !df.schema().contains(column) {
182        return Err(PolarsError::ComputeError(
183            format!("MIN calculation requires {column} column").into()
184        ));
185    }
186    
187    let series = df.column(column)?.f64()?;
188    
189    let mut min_values = Vec::with_capacity(df.height());
190    
191    // Fill initial values with NaN
192    for i in 0..window-1 {
193        min_values.push(f64::NAN);
194    }
195    
196    // Calculate min for each window
197    for i in window-1..df.height() {
198        let mut min_val = f64::INFINITY;
199        let mut all_nan = true;
200        
201        for j in 0..window {
202            let val = series.get(i - j).unwrap_or(f64::NAN);
203            if !val.is_nan() {
204                min_val = min_val.min(val);
205                all_nan = false;
206            }
207        }
208        
209        if all_nan {
210            min_values.push(f64::NAN);
211        } else {
212            min_values.push(min_val);
213        }
214    }
215    
216    Ok(Series::new(format!("{column}_min_{window}").into(), min_values))
217}
218
219/// Calculate sum over a specified window
220///
221/// # Arguments
222///
223/// * `df` - DataFrame containing the data
224/// * `column` - Column name to calculate on
225/// * `window` - Window size for the calculation
226///
227/// # Returns
228///
229/// Returns a PolarsResult containing the SUM Series
230pub fn calculate_sum(df: &DataFrame, column: &str, window: usize) -> PolarsResult<Series> {
231    if !df.schema().contains(column) {
232        return Err(PolarsError::ComputeError(
233            format!("SUM calculation requires {column} column").into()
234        ));
235    }
236    
237    let series = df.column(column)?.f64()?;
238    
239    let mut sum_values = Vec::with_capacity(df.height());
240    
241    // Fill initial values with NaN
242    for i in 0..window-1 {
243        sum_values.push(f64::NAN);
244    }
245    
246    // Calculate sum for each window
247    for i in window-1..df.height() {
248        let mut sum = 0.0;
249        let mut all_nan = true;
250        
251        for j in 0..window {
252            let val = series.get(i - j).unwrap_or(f64::NAN);
253            if !val.is_nan() {
254                sum += val;
255                all_nan = false;
256            }
257        }
258        
259        if all_nan {
260            sum_values.push(f64::NAN);
261        } else {
262            sum_values.push(sum);
263        }
264    }
265    
266    Ok(Series::new(format!("{column}_sum_{window}").into(), sum_values))
267}
ta_lib_in_rust/indicators/math/mod.rs

ta_lib_in_rust/indicators/math/
mod.rs
