use crate::traits::{Error, Fit, Result, Transform};
use polars::prelude::*;
pub struct Difference {
fitted: bool,
columns: Vec<String>,
period: i64,
pct_change: bool,
}
impl Difference {
pub fn new(columns: &[&str], period: i64, pct_change: bool) -> Self {
Self {
fitted: false,
columns: columns.iter().map(|s| s.to_string()).collect(),
period,
pct_change,
}
}
pub fn diff(columns: &[&str], period: i64) -> Self {
Self::new(columns, period, false)
}
pub fn pct_change(columns: &[&str], period: i64) -> Self {
Self::new(columns, period, true)
}
}
impl Fit<DataFrame, DataFrame> for Difference {
type Output = ();
fn fit(&mut self, x: DataFrame, _y: DataFrame) -> Result<()> {
if self.columns.is_empty() {
return Err(Error::InvalidInput(
"Difference: at least one column is required.".into(),
));
}
for col in &self.columns {
if x.column(col.as_str()).is_err() {
return Err(Error::InvalidInput(format!(
"Difference: column '{}' not found.",
col
)));
}
}
self.fitted = true;
Ok(())
}
}
impl Transform<DataFrame> for Difference {
type Output = DataFrame;
fn transform(&self, x: DataFrame) -> Result<DataFrame> {
if !self.fitted {
return Err(Error::NotFitted("Difference".into()));
}
let mut out = x.clone();
for col in &self.columns {
let s = out.column(col.as_str()).unwrap().clone();
let suffix = if self.pct_change { "pct" } else { "diff" };
let col_name = format!("{}_{}_{}", col, suffix, self.period);
let orig = s.f64().map_err(|_| {
Error::InvalidInput(format!("Difference: column '{}' is not f64", col))
})?;
let shifted = s.shift(self.period);
let shifted = shifted.f64().map_err(|_| {
Error::InvalidInput(format!("Difference: column '{}' is not f64", col))
})?;
let result: ChunkedArray<Float64Type> = orig
.iter()
.zip(shifted.iter())
.map(|(a, b)| match (a, b) {
(Some(va), Some(vb)) => {
if self.pct_change {
if vb.abs() > f64::EPSILON {
Some((va - vb) / vb)
} else {
None
}
} else {
Some(va - vb)
}
}
_ => None,
})
.collect();
out.with_column(
result
.into_series()
.with_name(col_name.as_str().into())
.into(),
)
.map_err(|e| Error::Computation(e.to_string()))?;
}
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_relative_eq;
#[test]
fn test_diff() {
let vals = Column::from(Series::new("x".into(), &[10.0f64, 20.0, 30.0, 40.0]));
let df = DataFrame::new(4, vec![vals]).unwrap();
let mut d = Difference::diff(&["x"], 1);
let y = df.clone();
d.fit(df.clone(), y).unwrap();
let result = d.transform(df).unwrap();
let diffed = result.column("x_diff_1").unwrap().f64().unwrap();
assert!(diffed.get(0).is_none());
assert_relative_eq!(diffed.get(1).unwrap(), 10.0, epsilon = 1e-6);
assert_relative_eq!(diffed.get(2).unwrap(), 10.0, epsilon = 1e-6);
}
#[test]
fn test_pct_change() {
let vals = Column::from(Series::new("x".into(), &[100.0f64, 110.0, 121.0]));
let df = DataFrame::new(3, vec![vals]).unwrap();
let mut d = Difference::pct_change(&["x"], 1);
let y = df.clone();
d.fit(df.clone(), y).unwrap();
let result = d.transform(df).unwrap();
let pct = result.column("x_pct_1").unwrap().f64().unwrap();
assert!(pct.get(0).is_none());
assert_relative_eq!(pct.get(1).unwrap(), 0.1, epsilon = 1e-6);
}
}