CLN/TYPE: EWM (#34770)

* Move min_periods validation to init

* Type signatures

* Undo unnecessary casting

* consolidate some cython type declariations

* tighten up typing and black

Co-authored-by: Matt Roeschke <[email protected]>

Author: mroeschke

pandas/_libs/window/aggregations.pyx

@@ -1759,7 +1759,7 @@ def roll_weighted_var(float64_t[:] values, float64_t[:] weights,
 # Exponentially weighted moving average
 
 
-def ewma(float64_t[:] vals, float64_t com, int adjust, bint ignore_na, int minp):
+def ewma(float64_t[:] vals, float64_t com, bint adjust, bint ignore_na, int minp):
     """
     Compute exponentially-weighted moving average using center-of-mass.
 
@@ -1777,17 +1777,14 @@ def ewma(float64_t[:] vals, float64_t com, int adjust, bint ignore_na, int minp)
     """
 
     cdef:
-        Py_ssize_t N = len(vals)
+        Py_ssize_t i, nobs, N = len(vals)
         ndarray[float64_t] output = np.empty(N, dtype=float)
         float64_t alpha, old_wt_factor, new_wt, weighted_avg, old_wt, cur
-        Py_ssize_t i, nobs
         bint is_observation
 
     if N == 0:
         return output
 
-    minp = max(minp, 1)
-
     alpha = 1. / (1. + com)
     old_wt_factor = 1. - alpha
     new_wt = 1. if adjust else alpha
@@ -1831,7 +1828,7 @@ def ewma(float64_t[:] vals, float64_t com, int adjust, bint ignore_na, int minp)
 
 
 def ewmcov(float64_t[:] input_x, float64_t[:] input_y,
-           float64_t com, int adjust, bint ignore_na, int minp, int bias):
+           float64_t com, bint adjust, bint ignore_na, int minp, bint bias):
     """
     Compute exponentially-weighted moving variance using center-of-mass.
 
@@ -1851,11 +1848,10 @@ def ewmcov(float64_t[:] input_x, float64_t[:] input_y,
     """
 
     cdef:
-        Py_ssize_t N = len(input_x), M = len(input_y)
+        Py_ssize_t i, nobs, N = len(input_x), M = len(input_y)
         float64_t alpha, old_wt_factor, new_wt, mean_x, mean_y, cov
         float64_t sum_wt, sum_wt2, old_wt, cur_x, cur_y, old_mean_x, old_mean_y
         float64_t numerator, denominator
-        Py_ssize_t i, nobs
         ndarray[float64_t] output
         bint is_observation
 
@@ -1866,8 +1862,6 @@ def ewmcov(float64_t[:] input_x, float64_t[:] input_y,
     if N == 0:
         return output
 
-    minp = max(minp, 1)
-
     alpha = 1. / (1. + com)
     old_wt_factor = 1. - alpha
     new_wt = 1. if adjust else alpha

pandas/core/window/ewm.py

@@ -1,9 +1,11 @@
 from functools import partial
 from textwrap import dedent
+from typing import Optional, Union
 
 import numpy as np
 
 import pandas._libs.window.aggregations as window_aggregations
+from pandas._typing import FrameOrSeries
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender, Substitution
 
@@ -24,7 +26,12 @@
 """
 
 
-def get_center_of_mass(comass, span, halflife, alpha) -> float:
+def get_center_of_mass(
+    comass: Optional[float],
+    span: Optional[float],
+    halflife: Optional[float],
+    alpha: Optional[float],
+) -> float:
     valid_count = com.count_not_none(comass, span, halflife, alpha)
     if valid_count > 1:
         raise ValueError("comass, span, halflife, and alpha are mutually exclusive")
@@ -114,7 +121,7 @@ class EWM(_Rolling):
           used in calculating the final weighted average of
           [:math:`x_0`, None, :math:`x_2`] are :math:`1-\alpha` and :math:`1` if
           ``adjust=True``, and :math:`1-\alpha` and :math:`\alpha` if ``adjust=False``.
-    axis : {0 or 'index', 1 or 'columns'}, default 0
+    axis : {0, 1}, default 0
         The axis to use. The value 0 identifies the rows, and 1
         identifies the columns.
 
@@ -159,18 +166,18 @@ class EWM(_Rolling):
     def __init__(
         self,
         obj,
-        com=None,
-        span=None,
-        halflife=None,
-        alpha=None,
-        min_periods=0,
-        adjust=True,
-        ignore_na=False,
-        axis=0,
+        com: Optional[float] = None,
+        span: Optional[float] = None,
+        halflife: Optional[float] = None,
+        alpha: Optional[float] = None,
+        min_periods: int = 0,
+        adjust: bool = True,
+        ignore_na: bool = False,
+        axis: int = 0,
     ):
         self.obj = obj
         self.com = get_center_of_mass(com, span, halflife, alpha)
-        self.min_periods = min_periods
+        self.min_periods = max(int(min_periods), 1)
         self.adjust = adjust
         self.ignore_na = ignore_na
         self.axis = axis
@@ -274,16 +281,16 @@ def mean(self, *args, **kwargs):
         window_func = partial(
             window_func,
             com=self.com,
-            adjust=int(self.adjust),
+            adjust=self.adjust,
             ignore_na=self.ignore_na,
-            minp=int(self.min_periods),
+            minp=self.min_periods,
         )
         return self._apply(window_func)
 
     @Substitution(name="ewm", func_name="std")
     @Appender(_doc_template)
     @Appender(_bias_template)
-    def std(self, bias=False, *args, **kwargs):
+    def std(self, bias: bool = False, *args, **kwargs):
         """
         Exponential weighted moving stddev.
         """
@@ -295,28 +302,28 @@ def std(self, bias=False, *args, **kwargs):
     @Substitution(name="ewm", func_name="var")
     @Appender(_doc_template)
     @Appender(_bias_template)
-    def var(self, bias=False, *args, **kwargs):
+    def var(self, bias: bool = False, *args, **kwargs):
         """
         Exponential weighted moving variance.
         """
         nv.validate_window_func("var", args, kwargs)
 
         def f(arg):
             return window_aggregations.ewmcov(
-                arg,
-                arg,
-                self.com,
-                int(self.adjust),
-                int(self.ignore_na),
-                int(self.min_periods),
-                int(bias),
+                arg, arg, self.com, self.adjust, self.ignore_na, self.min_periods, bias,
             )
 
         return self._apply(f)
 
     @Substitution(name="ewm", func_name="cov")
     @Appender(_doc_template)
-    def cov(self, other=None, pairwise=None, bias=False, **kwargs):
+    def cov(
+        self,
+        other: Optional[Union[np.ndarray, FrameOrSeries]] = None,
+        pairwise: Optional[bool] = None,
+        bias: bool = False,
+        **kwargs,
+    ):
         """
         Exponential weighted sample covariance.
 
@@ -350,10 +357,10 @@ def _get_cov(X, Y):
                 X._prep_values(),
                 Y._prep_values(),
                 self.com,
-                int(self.adjust),
-                int(self.ignore_na),
-                int(self.min_periods),
-                int(bias),
+                self.adjust,
+                self.ignore_na,
+                self.min_periods,
+                bias,
             )
             return X._wrap_result(cov)
 
@@ -363,7 +370,12 @@ def _get_cov(X, Y):
 
     @Substitution(name="ewm", func_name="corr")
     @Appender(_doc_template)
-    def corr(self, other=None, pairwise=None, **kwargs):
+    def corr(
+        self,
+        other: Optional[Union[np.ndarray, FrameOrSeries]] = None,
+        pairwise: Optional[bool] = None,
+        **kwargs,
+    ):
         """
         Exponential weighted sample correlation.
 
@@ -394,13 +406,7 @@ def _get_corr(X, Y):
 
             def _cov(x, y):
                 return window_aggregations.ewmcov(
-                    x,
-                    y,
-                    self.com,
-                    int(self.adjust),
-                    int(self.ignore_na),
-                    int(self.min_periods),
-                    1,
+                    x, y, self.com, self.adjust, self.ignore_na, self.min_periods, 1,
                 )
 
             x_values = X._prep_values()