diff --git a/ta_lib/core/src/ops.rs b/ta_lib/core/src/ops.rs index f69807f0..1692ac97 100644 --- a/ta_lib/core/src/ops.rs +++ b/ta_lib/core/src/ops.rs @@ -2,18 +2,29 @@ use crate::series::Series; use std::ops::{Add, Div, Mul, Neg, Sub}; impl Series { - pub fn add_series(&self, rhs: &Series) -> Series { + fn binary_op_series(&self, rhs: &Series, op: F) -> Series + where + F: Fn(&f64, &f64) -> f64, + { self.zip_with(rhs, |a, b| match (a, b) { - (Some(a_val), Some(b_val)) => Some(*a_val + *b_val), + (Some(a_val), Some(b_val)) => Some(op(a_val, b_val)), _ => None, }) } + fn unary_op_scalar(&self, scalar: f64, op: F) -> Series + where + F: Fn(&f64, f64) -> f64, + { + self.fmap(|val| val.map(|v| op(v, scalar))) + } + + pub fn add_series(&self, rhs: &Series) -> Series { + self.binary_op_series(rhs, |a, b| a + b) + } + pub fn mul_series(&self, rhs: &Series) -> Series { - self.zip_with(rhs, |a, b| match (a, b) { - (Some(a_val), Some(b_val)) => Some(*a_val * *b_val), - _ => None, - }) + self.binary_op_series(rhs, |a, b| a * b) } pub fn div_series(&self, rhs: &Series) -> Series { @@ -28,7 +39,7 @@ impl Series { None } } else { - Some(*a_val / *b_val) + Some(a_val / b_val) } } _ => None, @@ -36,40 +47,35 @@ impl Series { } pub fn sub_series(&self, rhs: &Series) -> Series { - self.zip_with(rhs, |a, b| match (a, b) { - (Some(a_val), Some(b_val)) => Some(*a_val - *b_val), - _ => None, - }) + self.binary_op_series(rhs, |a, b| a - b) } pub fn add_scalar(&self, scalar: f64) -> Series { - self.fmap(|val| val.map(|v| v + scalar)) + self.unary_op_scalar(scalar, |v, s| v + s) } pub fn mul_scalar(&self, scalar: f64) -> Series { - self.fmap(|val| val.map(|v| v * scalar)) + self.unary_op_scalar(scalar, |v, s| v * s) } pub fn div_scalar(&self, scalar: f64) -> Series { - self.fmap(|val| { - val.map(|v| { - if *v == 0.0 { - if scalar > 0.0 { - std::f64::INFINITY - } else if scalar < 0.0 { - std::f64::NEG_INFINITY - } else { - 0.0 - } + self.unary_op_scalar(scalar, |v, s| { + if *v == 0.0 { + if s > 0.0 { + std::f64::INFINITY + } else if s < 0.0 { + std::f64::NEG_INFINITY } else { - v / scalar + 0.0 } - }) + } else { + v / s + } }) } pub fn sub_scalar(&self, scalar: f64) -> Series { - self.fmap(|val| val.map(|v| v - scalar)) + self.unary_op_scalar(scalar, |v, s| v - s) } pub fn neg(&self) -> Series { @@ -78,18 +84,19 @@ impl Series { } impl Series { - pub fn mul_series(&self, rhs: &Series) -> Series { + fn bool_op_series(&self, rhs: &Series, op: F) -> Series + where + F: Fn(&bool, &f64) -> f64, + { self.zip_with(rhs, |b, a| match (b, a) { - (Some(b_val), Some(a_val)) => { - if *b_val { - Some(*a_val) - } else { - Some(0.0) - } - } + (Some(b_val), Some(a_val)) => Some(op(b_val, a_val)), _ => None, }) } + + pub fn mul_series(&self, rhs: &Series) -> Series { + self.bool_op_series(rhs, |b, a| if *b { *a } else { 0.0 }) + } } impl Add> for &Series {