use polars::prelude::*;
pub fn calculate_obv(df: &DataFrame) -> PolarsResult<Series> {
if !df.schema().contains("close") || !df.schema().contains("volume") {
return Err(PolarsError::ComputeError(
"OBV calculation requires both close and volume columns".into(),
));
}
let close = df.column("close")?.f64()?;
let volume = df.column("volume")?.f64()?;
let mut obv = Vec::with_capacity(df.height());
obv.push(volume.get(0).unwrap_or(0.0));
for i in 1..df.height() {
let curr_close = close.get(i).unwrap_or(0.0);
let prev_close = close.get(i - 1).unwrap_or(0.0);
let curr_volume = volume.get(i).unwrap_or(0.0);
if curr_close > prev_close {
obv.push(obv[i - 1] + curr_volume);
} else if curr_close < prev_close {
obv.push(obv[i - 1] - curr_volume);
} else {
obv.push(obv[i - 1]);
}
}
Ok(Series::new("obv".into(), obv))
}