diff --git a/pandas/_libs/window/aggregations.pyx b/pandas/_libs/window/aggregations.pyx index 9e088062d7280..646444d10e416 100644 --- a/pandas/_libs/window/aggregations.pyx +++ b/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 diff --git a/pandas/core/window/ewm.py b/pandas/core/window/ewm.py index a5e30c900cae2..0e39b94574a12 100644 --- a/pandas/core/window/ewm.py +++ b/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,7 +302,7 @@ 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. """ @@ -303,20 +310,20 @@ def var(self, bias=False, *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()