fixup

pola-rs · Mar 3, 2024 · 0b61bc7 · 0b61bc7
1 parent 29ebc9e
commit 0b61bc7
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 28 deletions.
diff --git a/polars_xdt/functions.py b/polars_xdt/functions.py
@@ -836,8 +836,8 @@ def ewma_by_time(
     halflife: timedelta,
     adjust: bool = True,
 ) -> pl.Expr:
-    """
-    Calculated time-based exponentially weighted moving average.
+    r"""
+    Calculate time-based exponentially weighted moving average.
 
     Given observations :math:`x_1, x_2, \ldots, x_n` at times
     :math:`t_1, t_2, \ldots, t_n`, the **unadjusted** EWMA is calculated as
@@ -846,11 +846,11 @@ def ewma_by_time(
 
             y_0 &= x_0
 
-            \\alpha_i &= exp(-\\lambda(t_i - t_{i-1}))
+            \alpha_i &= exp(-\lambda(t_i - t_{i-1}))
+
+            y_i &= \alpha_i x_i + (1 - \alpha_i) y_{i-1}; \quad i > 0
 
-            y_i &= \\alpha_i x_i + (1 - \\alpha_i) y_{i-1}; \\quad i > 0
-    
-    where :math:`\\lambda` equals :math:`\\ln(2) / \\text{halflife}`.
+    where :math:`\lambda` equals :math:`\ln(2) / \text{halflife}`.
 
     Parameters
     ----------
@@ -863,23 +863,33 @@ def ewma_by_time(
     adjust
         Whether to adjust the result to account for the bias towards the
         initial value. Defaults to True.
-    
+
     Returns
     -------
     pl.Expr
         Float64
-    
+
     Examples
     --------
     >>> import polars as pl
     >>> import polars_xdt as xdt
     >>> from datetime import date, timedelta
     >>> df = pl.DataFrame(
     ...     {
-    ...      'values': [0, 1, 2, None, 4],
-    ...       'times': [date(2020, 1, 1), date(2020, 1, 3), date(2020, 1, 10), date(2020, 1, 15), date(2020, 1, 17)]})
+    ...         "values": [0, 1, 2, None, 4],
+    ...         "times": [
+    ...             date(2020, 1, 1),
+    ...             date(2020, 1, 3),
+    ...             date(2020, 1, 10),
+    ...             date(2020, 1, 15),
+    ...             date(2020, 1, 17),
+    ...         ],
+    ...     }
+    ... )
     >>> df.with_columns(
-    ...     ewma = xdt.ewma_by_time("values", times="times", halflife=timedelta(days=4)),
+    ...     ewma=xdt.ewma_by_time(
+    ...         "values", times="times", halflife=timedelta(days=4)
+    ...     ),
     ... )
     shape: (5, 3)
     ┌────────┬────────────┬──────────┐
@@ -893,6 +903,7 @@ def ewma_by_time(
     │ null   ┆ 2020-01-15 ┆ null     │
     │ 4      ┆ 2020-01-17 ┆ 3.233686 │
     └────────┴────────────┴──────────┘
+
     """
     times = parse_into_expr(times)
     halflife_us = (

diff --git a/src/ewma_by_time.rs b/src/ewma_by_time.rs
@@ -1,17 +1,16 @@
 use polars::prelude::*;
 use polars_arrow::array::PrimitiveArray;
-use pyo3_polars::export::polars_core::export::num::{Float, Pow};
+use pyo3_polars::export::polars_core::export::num::Pow;
 
 pub(crate) fn impl_ewma_by_time_float(
     times: &Int64Chunked,
     values: &Float64Chunked,
     halflife: i64,
     adjust: bool,
     time_unit: TimeUnit,
-) -> Float64Chunked
-{
+) -> Float64Chunked {
     let mut out = Vec::with_capacity(times.len());
-    if values.len() == 0 {
+    if values.is_empty() {
         return Float64Chunked::full_null("", times.len());
     }
 
@@ -36,7 +35,7 @@ pub(crate) fn impl_ewma_by_time_float(
                     let result: f64;
                     if adjust {
                         alpha *= Pow::pow(0.5, delta_time as f64 / halflife as f64);
-                        result = (value + alpha * prev_result) / ((1. + alpha));
+                        result = (value + alpha * prev_result) / (1. + alpha);
                         alpha += 1.;
                     } else {
                         // equivalent to:
@@ -56,7 +55,6 @@ pub(crate) fn impl_ewma_by_time_float(
     Float64Chunked::from(arr)
 }
 
-
 pub(crate) fn impl_ewma_by_time(
     times: &Int64Chunked,
     values: &Series,
@@ -72,13 +70,13 @@ pub(crate) fn impl_ewma_by_time(
         DataType::Int64 | DataType::Int32 => {
             let values = values.cast(&DataType::Float64).unwrap();
             let values = values.f64().unwrap();
-            impl_ewma_by_time_float(times, &values, halflife, adjust, time_unit).into_series()
+            impl_ewma_by_time_float(times, values, halflife, adjust, time_unit).into_series()
         }
         DataType::Float32 => {
             // todo: preserve Float32 in this case
             let values = values.cast(&DataType::Float64).unwrap();
             let values = values.f64().unwrap();
-            impl_ewma_by_time_float(times, &values, halflife, adjust, time_unit).into_series()
+            impl_ewma_by_time_float(times, values, halflife, adjust, time_unit).into_series()
         }
         dt => panic!("Expected values to be signed numeric, got {:?}", dt),
     }

diff --git a/src/expressions.rs b/src/expressions.rs
@@ -180,13 +180,27 @@ fn ewma_by_time(inputs: &[Series], kwargs: EwmTimeKwargs) -> PolarsResult<Series
     match &inputs[0].dtype() {
         DataType::Datetime(_, _) => {
             let time = &inputs[0].datetime().unwrap();
-            Ok(impl_ewma_by_time(&time.0, values, kwargs.halflife, kwargs.adjust, time.time_unit()).into_series())
+            Ok(impl_ewma_by_time(
+                &time.0,
+                values,
+                kwargs.halflife,
+                kwargs.adjust,
+                time.time_unit(),
+            )
+            .into_series())
         }
         DataType::Date => {
             let binding = &inputs[0].cast(&DataType::Datetime(TimeUnit::Milliseconds, None))?;
             let time = binding.datetime().unwrap();
-            Ok(impl_ewma_by_time(&time.0, values, kwargs.halflife, kwargs.adjust, time.time_unit()).into_series())
+            Ok(impl_ewma_by_time(
+                &time.0,
+                values,
+                kwargs.halflife,
+                kwargs.adjust,
+                time.time_unit(),
+            )
+            .into_series())
         }
-        _ => polars_bail!(InvalidOperation: "First argument should be a date or datetime type.")
+        _ => polars_bail!(InvalidOperation: "First argument should be a date or datetime type."),
     }
 }
diff --git a/tests/test_ewma_by_time.py b/tests/test_ewma_by_time.py
@@ -3,16 +3,34 @@
 import polars_xdt as xdt
 from datetime import datetime, timedelta
 
+
 def test_ewma_by_time():
     df = pl.DataFrame(
         {
-         'values': [0., 1, 2, None, 4],
-          'times': [datetime(2020, 1, 1), datetime(2020, 1, 3), datetime(2020, 1, 10), datetime(2020, 1, 15), datetime(2020, 1, 17)]})
+            "values": [0.0, 1, 2, None, 4],
+            "times": [
+                datetime(2020, 1, 1),
+                datetime(2020, 1, 3),
+                datetime(2020, 1, 10),
+                datetime(2020, 1, 15),
+                datetime(2020, 1, 17),
+            ],
+        }
+    )
     result = df.select(
-        ewma = xdt.ewma_by_time("values", times="times", halflife=timedelta(days=4)),
+        ewma=xdt.ewma_by_time(
+            "values", times="times", halflife=timedelta(days=4)
+        ),
     )
     expected = pl.DataFrame(
         {
-            'ewma': [0.0, 0.585786437626905, 1.52388878049859, None, 3.2336858398518338]
-        })
-    assert_frame_equal(result, expected)
+            "ewma": [
+                0.0,
+                0.585786437626905,
+                1.52388878049859,
+                None,
+                3.2336858398518338,
+            ]
+        }
+    )
+    assert_frame_equal(result, expected)