ENH: Support groupby.ewm operations (#37878)

Author: mroeschke

asv_bench/benchmarks/rolling.py

@@ -225,4 +225,17 @@ def time_rolling_offset(self, method):
         getattr(self.groupby_roll_offset, method)()
 
 
+class GroupbyEWM:
+
+    params = ["cython", "numba"]
+    param_names = ["engine"]
+
+    def setup(self, engine):
+        df = pd.DataFrame({"A": range(50), "B": range(50)})
+        self.gb_ewm = df.groupby("A").ewm(com=1.0)
+
+    def time_groupby_mean(self, engine):
+        self.gb_ewm.mean(engine=engine)
+
+
 from .pandas_vb_common import setup  # noqa: F401 isort:skip

doc/source/user_guide/window.rst

@@ -43,7 +43,7 @@ Concept                         Method             Returned Object
 Rolling window                  ``rolling``        ``Rolling``                   Yes                          Yes
 Weighted window                 ``rolling``        ``Window``                    No                           No
 Expanding window                ``expanding``      ``Expanding``                 No                           Yes
-Exponentially Weighted window   ``ewm``            ``ExponentialMovingWindow``   No                           No
+Exponentially Weighted window   ``ewm``            ``ExponentialMovingWindow``   No                           Yes (as of version 1.2)
 =============================   =================  ===========================   ===========================  ========================
 
 As noted above, some operations support specifying a window based on a time offset:

doc/source/whatsnew/v1.2.0.rst

@@ -204,6 +204,23 @@ example where the index name is preserved:
 The same is true for :class:`MultiIndex`, but the logic is applied separately on a
 level-by-level basis.
 
+.. _whatsnew_120.groupby_ewm:
+
+Groupby supports EWM operations directly
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+:class:`DataFrameGroupBy` now supports exponentially weighted window operations directly (:issue:`16037`).
+
+.. ipython:: python
+
+    df = pd.DataFrame({'A': ['a', 'b', 'a', 'b'], 'B': range(4)})
+    df
+    df.groupby('A').ewm(com=1.0).mean()
+
+Additionally ``mean`` supports execution via `Numba <https://numba.pydata.org/>`__ with
+the  ``engine`` and ``engine_kwargs`` arguments. Numba must be installed as an optional dependency
+to use this feature.
+
 .. _whatsnew_120.enhancements.other:
 
 Other enhancements

pandas/_libs/window/aggregations.pyx

@@ -1496,15 +1496,17 @@ def roll_weighted_var(float64_t[:] values, float64_t[:] weights,
 # ----------------------------------------------------------------------
 # Exponentially weighted moving average
 
-def ewma_time(const float64_t[:] vals, int minp, ndarray[int64_t] times,
-              int64_t halflife):
+def ewma_time(const float64_t[:] vals, int64_t[:] start, int64_t[:] end,
+              int minp, ndarray[int64_t] times, int64_t halflife):
     """
     Compute exponentially-weighted moving average using halflife and time
     distances.
 
     Parameters
     ----------
     vals : ndarray[float_64]
+    start: ndarray[int_64]
+    end: ndarray[int_64]
     minp : int
     times : ndarray[int64]
     halflife : int64
@@ -1552,17 +1554,20 @@ def ewma_time(const float64_t[:] vals, int minp, ndarray[int64_t] times,
     return output
 
 
-def ewma(float64_t[:] vals, float64_t com, bint adjust, bint ignore_na, int minp):
+def ewma(float64_t[:] vals, int64_t[:] start, int64_t[:] end, int minp,
+         float64_t com, bint adjust, bint ignore_na):
     """
     Compute exponentially-weighted moving average using center-of-mass.
 
     Parameters
     ----------
     vals : ndarray (float64 type)
+    start: ndarray (int64 type)
+    end: ndarray (int64 type)
+    minp : int
     com : float64
     adjust : int
     ignore_na : bool
-    minp : int
 
     Returns
     -------
@@ -1620,19 +1625,21 @@ def ewma(float64_t[:] vals, float64_t com, bint adjust, bint ignore_na, int minp
 # Exponentially weighted moving covariance
 
 
-def ewmcov(float64_t[:] input_x, float64_t[:] input_y,
-           float64_t com, bint adjust, bint ignore_na, int minp, bint bias):
+def ewmcov(float64_t[:] input_x, int64_t[:] start, int64_t[:] end, int minp,
+           float64_t[:] input_y, float64_t com, bint adjust, bint ignore_na, bint bias):
     """
     Compute exponentially-weighted moving variance using center-of-mass.
 
     Parameters
     ----------
     input_x : ndarray (float64 type)
+    start: ndarray (int64 type)
+    end: ndarray (int64 type)
+    minp : int
     input_y : ndarray (float64 type)
     com : float64
     adjust : int
     ignore_na : bool
-    minp : int
     bias : int
 
     Returns

pandas/core/groupby/base.py

@@ -192,6 +192,7 @@ def _gotitem(self, key, ndim, subset=None):
         "describe",
         "dtypes",
         "expanding",
+        "ewm",
         "filter",
         "get_group",
         "groups",

pandas/core/groupby/groupby.py

@@ -1859,6 +1859,16 @@ def expanding(self, *args, **kwargs):
 
         return ExpandingGroupby(self, *args, **kwargs)
 
+    @Substitution(name="groupby")
+    @Appender(_common_see_also)
+    def ewm(self, *args, **kwargs):
+        """
+        Return an ewm grouper, providing ewm functionality per group.
+        """
+        from pandas.core.window import ExponentialMovingWindowGroupby
+
+        return ExponentialMovingWindowGroupby(self, *args, **kwargs)
+
     def _fill(self, direction, limit=None):
         """
         Shared function for `pad` and `backfill` to call Cython method.

pandas/core/window/__init__.py

@@ -1,3 +1,6 @@
-from pandas.core.window.ewm import ExponentialMovingWindow  # noqa:F401
+from pandas.core.window.ewm import (  # noqa:F401
+    ExponentialMovingWindow,
+    ExponentialMovingWindowGroupby,
+)
 from pandas.core.window.expanding import Expanding, ExpandingGroupby  # noqa:F401
 from pandas.core.window.rolling import Rolling, RollingGroupby, Window  # noqa:F401

pandas/core/window/ewm.py

@@ -14,8 +14,20 @@
 from pandas.core.dtypes.common import is_datetime64_ns_dtype
 
 import pandas.core.common as common
-from pandas.core.window.common import _doc_template, _shared_docs, zsqrt
-from pandas.core.window.rolling import BaseWindow, flex_binary_moment
+from pandas.core.util.numba_ import maybe_use_numba
+from pandas.core.window.common import (
+    _doc_template,
+    _shared_docs,
+    flex_binary_moment,
+    zsqrt,
+)
+from pandas.core.window.indexers import (
+    BaseIndexer,
+    ExponentialMovingWindowIndexer,
+    GroupbyIndexer,
+)
+from pandas.core.window.numba_ import generate_numba_groupby_ewma_func
+from pandas.core.window.rolling import BaseWindow, BaseWindowGroupby, dispatch
 
 if TYPE_CHECKING:
     from pandas import Series
@@ -219,14 +231,16 @@ def __init__(
         ignore_na: bool = False,
         axis: int = 0,
         times: Optional[Union[str, np.ndarray, FrameOrSeries]] = None,
+        **kwargs,
     ):
-        self.com: Optional[float]
         self.obj = obj
         self.min_periods = max(int(min_periods), 1)
         self.adjust = adjust
         self.ignore_na = ignore_na
         self.axis = axis
         self.on = None
+        self.center = False
+        self.closed = None
         if times is not None:
             if isinstance(times, str):
                 times = self._selected_obj[times]
@@ -245,7 +259,7 @@ def __init__(
             if common.count_not_none(com, span, alpha) > 0:
                 self.com = get_center_of_mass(com, span, None, alpha)
             else:
-                self.com = None
+                self.com = 0.0
         else:
             if halflife is not None and isinstance(halflife, (str, datetime.timedelta)):
                 raise ValueError(
@@ -260,6 +274,12 @@ def __init__(
     def _constructor(self):
         return ExponentialMovingWindow
 
+    def _get_window_indexer(self) -> BaseIndexer:
+        """
+        Return an indexer class that will compute the window start and end bounds
+        """
+        return ExponentialMovingWindowIndexer()
+
     _agg_see_also_doc = dedent(
         """
     See Also
@@ -299,27 +319,6 @@ def aggregate(self, func, *args, **kwargs):
 
     agg = aggregate
 
-    def _apply(self, func):
-        """
-        Rolling statistical measure using supplied function. Designed to be
-        used with passed-in Cython array-based functions.
-
-        Parameters
-        ----------
-        func : str/callable to apply
-
-        Returns
-        -------
-        y : same type as input argument
-        """
-
-        def homogeneous_func(values: np.ndarray):
-            if values.size == 0:
-                return values.copy()
-            return np.apply_along_axis(func, self.axis, values)
-
-        return self._apply_blockwise(homogeneous_func)
-
     @Substitution(name="ewm", func_name="mean")
     @Appender(_doc_template)
     def mean(self, *args, **kwargs):
@@ -336,7 +335,6 @@ def mean(self, *args, **kwargs):
             window_func = self._get_roll_func("ewma_time")
             window_func = partial(
                 window_func,
-                minp=self.min_periods,
                 times=self.times,
                 halflife=self.halflife,
             )
@@ -347,7 +345,6 @@ def mean(self, *args, **kwargs):
                 com=self.com,
                 adjust=self.adjust,
                 ignore_na=self.ignore_na,
-                minp=self.min_periods,
             )
         return self._apply(window_func)
 
@@ -371,13 +368,19 @@ def var(self, bias: bool = False, *args, **kwargs):
         Exponential weighted moving variance.
         """
         nv.validate_window_func("var", args, kwargs)
+        window_func = self._get_roll_func("ewmcov")
+        window_func = partial(
+            window_func,
+            com=self.com,
+            adjust=self.adjust,
+            ignore_na=self.ignore_na,
+            bias=bias,
+        )
 
-        def f(arg):
-            return window_aggregations.ewmcov(
-                arg, arg, self.com, self.adjust, self.ignore_na, self.min_periods, bias
-            )
+        def var_func(values, begin, end, min_periods):
+            return window_func(values, begin, end, min_periods, values)
 
-        return self._apply(f)
+        return self._apply(var_func)
 
     @Substitution(name="ewm", func_name="cov")
     @Appender(_doc_template)
@@ -419,11 +422,13 @@ def _get_cov(X, Y):
             Y = self._shallow_copy(Y)
             cov = window_aggregations.ewmcov(
                 X._prep_values(),
+                np.array([0], dtype=np.int64),
+                np.array([0], dtype=np.int64),
+                self.min_periods,
                 Y._prep_values(),
                 self.com,
                 self.adjust,
                 self.ignore_na,
-                self.min_periods,
                 bias,
             )
             return wrap_result(X, cov)
@@ -470,7 +475,15 @@ def _get_corr(X, Y):
 
             def _cov(x, y):
                 return window_aggregations.ewmcov(
-                    x, y, self.com, self.adjust, self.ignore_na, self.min_periods, 1
+                    x,
+                    np.array([0], dtype=np.int64),
+                    np.array([0], dtype=np.int64),
+                    self.min_periods,
+                    y,
+                    self.com,
+                    self.adjust,
+                    self.ignore_na,
+                    1,
                 )
 
             x_values = X._prep_values()
@@ -485,3 +498,78 @@ def _cov(x, y):
         return flex_binary_moment(
             self._selected_obj, other._selected_obj, _get_corr, pairwise=bool(pairwise)
         )
+
+
+class ExponentialMovingWindowGroupby(BaseWindowGroupby, ExponentialMovingWindow):
+    """
+    Provide an exponential moving window groupby implementation.
+    """
+
+    def _get_window_indexer(self) -> GroupbyIndexer:
+        """
+        Return an indexer class that will compute the window start and end bounds
+
+        Returns
+        -------
+        GroupbyIndexer
+        """
+        window_indexer = GroupbyIndexer(
+            groupby_indicies=self._groupby.indices,
+            window_indexer=ExponentialMovingWindowIndexer,
+        )
+        return window_indexer
+
+    var = dispatch("var", bias=False)
+    std = dispatch("std", bias=False)
+    cov = dispatch("cov", other=None, pairwise=None, bias=False)
+    corr = dispatch("corr", other=None, pairwise=None)
+
+    def mean(self, engine=None, engine_kwargs=None):
+        """
+        Parameters
+        ----------
+        engine : str, default None
+            * ``'cython'`` : Runs mean through C-extensions from cython.
+            * ``'numba'`` : Runs mean through JIT compiled code from numba.
+              Only available when ``raw`` is set to ``True``.
+            * ``None`` : Defaults to ``'cython'`` or globally setting
+              ``compute.use_numba``
+
+              .. versionadded:: 1.2.0
+
+        engine_kwargs : dict, default None
+            * For ``'cython'`` engine, there are no accepted ``engine_kwargs``
+            * For ``'numba'`` engine, the engine can accept ``nopython``, ``nogil``
+              and ``parallel`` dictionary keys. The values must either be ``True`` or
+              ``False``. The default ``engine_kwargs`` for the ``'numba'`` engine is
+              ``{'nopython': True, 'nogil': False, 'parallel': False}``.
+
+              .. versionadded:: 1.2.0
+
+        Returns
+        -------
+        Series or DataFrame
+            Return type is determined by the caller.
+        """
+        if maybe_use_numba(engine):
+            groupby_ewma_func = generate_numba_groupby_ewma_func(
+                engine_kwargs,
+                self.com,
+                self.adjust,
+                self.ignore_na,
+            )
+            return self._apply(
+                groupby_ewma_func,
+                numba_cache_key=(lambda x: x, "groupby_ewma"),
+            )
+        elif engine in ("cython", None):
+            if engine_kwargs is not None:
+                raise ValueError("cython engine does not accept engine_kwargs")
+
+            def f(x):
+                x = self._shallow_copy(x, groupby=self._groupby)
+                return x.mean()
+
+            return self._groupby.apply(f)
+        else:
+            raise ValueError("engine must be either 'numba' or 'cython'")