diff --git a/pandas/_libs/window.pyx b/pandas/_libs/window.pyx index 46e4b17b8164c..f0b6b7a5308fe 100644 --- a/pandas/_libs/window.pyx +++ b/pandas/_libs/window.pyx @@ -1675,9 +1675,9 @@ def roll_generic(object obj, return output -def roll_window(ndarray[float64_t, ndim=1, cast=True] values, +def roll_weighted_window(ndarray[float64_t, ndim=1, cast=True] values, ndarray[float64_t, ndim=1, cast=True] weights, - int minp, bint avg=True): + int minp, name="mean"): """ Assume len(weights) << len(values) """ @@ -1686,16 +1686,19 @@ def roll_window(ndarray[float64_t, ndim=1, cast=True] values, Py_ssize_t in_i, win_i, win_n, in_n float64_t val_in, val_win, c, w + # currently supported ops + assert name in ("mean", "sum") + in_n = len(values) win_n = len(weights) output = np.zeros(in_n, dtype=float) counts = np.zeros(in_n, dtype=float) - if avg: + if name == "mean": tot_wgt = np.zeros(in_n, dtype=float) minp = _check_minp(len(weights), minp, in_n) - if avg: + if name == "mean": for win_i in range(win_n): val_win = weights[win_i] if val_win != val_win: @@ -1719,7 +1722,7 @@ def roll_window(ndarray[float64_t, ndim=1, cast=True] values, else: output[in_i] /= tot_wgt[in_i] - else: + elif name == "sum": for win_i in range(win_n): val_win = weights[win_i] if val_win != val_win: diff --git a/pandas/core/window.py b/pandas/core/window.py index 0c1f6a1a6dace..510c87d1fbf70 100644 --- a/pandas/core/window.py +++ b/pandas/core/window.py @@ -707,7 +707,7 @@ def _pop_args(win_type, arg_names, kwargs): # GH #15662. `False` makes symmetric window, rather than periodic. return sig.get_window(win_type, window, False).astype(float) - def _apply_window(self, mean=True, **kwargs): + def _apply_window(self, name, **kwargs): """ Applies a moving window of type ``window_type`` on the data. @@ -750,11 +750,11 @@ def _apply_window(self, mean=True, **kwargs): def f(arg, *args, **kwargs): minp = _use_window(self.min_periods, len(window)) - return libwindow.roll_window( + return libwindow.roll_weighted_window( np.concatenate((arg, additional_nans)) if center else arg, window, minp, - avg=mean, + name=name, ) result = np.apply_along_axis(f, self.axis, values) @@ -831,13 +831,13 @@ def aggregate(self, arg, *args, **kwargs): @Appender(_shared_docs["sum"]) def sum(self, *args, **kwargs): nv.validate_window_func("sum", args, kwargs) - return self._apply_window(mean=False, **kwargs) + return self._apply_window(name="sum", **kwargs) @Substitution(name="window") @Appender(_shared_docs["mean"]) def mean(self, *args, **kwargs): nv.validate_window_func("mean", args, kwargs) - return self._apply_window(mean=True, **kwargs) + return self._apply_window(name="mean", **kwargs) class _GroupByMixin(GroupByMixin):