ENH: Allow users to definite their own window bound calculations in rolling (pandas-dev#29878)

Author: mroeschke

doc/source/reference/window.rst

@@ -74,3 +74,14 @@ Exponentially-weighted moving window functions
    EWM.var
    EWM.corr
    EWM.cov
+
+Window Indexer
+--------------
+.. currentmodule:: pandas
+
+Base class for defining custom window boundaries.
+
+.. autosummary::
+   :toctree: api/
+
+   api.indexers.BaseIndexer

doc/source/user_guide/computation.rst

@@ -466,6 +466,64 @@ default of the index) in a DataFrame.
    dft
    dft.rolling('2s', on='foo').sum()
 
+.. _stats.custom_rolling_window:
+
+Custom window rolling
+~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 1.0
+
+In addition to accepting an integer or offset as a ``window`` argument, ``rolling`` also accepts
+a ``BaseIndexer`` subclass that allows a user to define a custom method for calculating window bounds.
+The ``BaseIndexer`` subclass will need to define a ``get_window_bounds`` method that returns
+a tuple of two arrays, the first being the starting indices of the windows and second being the
+ending indices of the windows. Additionally, ``num_values``, ``min_periods``, ``center``, ``closed``
+and will automatically be passed to ``get_window_bounds`` and the defined method must
+always accept these arguments.
+
+For example, if we have the following ``DataFrame``:
+
+.. ipython:: python
+
+   use_expanding = [True, False, True, False, True]
+   use_expanding
+   df = pd.DataFrame({'values': range(5)})
+   df
+
+and we want to use an expanding window where ``use_expanding`` is ``True`` otherwise a window of size
+1, we can create the following ``BaseIndexer``:
+
+.. code-block:: ipython
+
+   In [2]: from pandas.api.indexers import BaseIndexer
+   ...:
+   ...: class CustomIndexer(BaseIndexer):
+   ...:
+   ...:    def get_window_bounds(self, num_values, min_periods, center, closed):
+   ...:        start = np.empty(num_values, dtype=np.int64)
+   ...:        end = np.empty(num_values, dtype=np.int64)
+   ...:        for i in range(num_values):
+   ...:            if self.use_expanding[i]:
+   ...:                start[i] = 0
+   ...:                end[i] = i + 1
+   ...:            else:
+   ...:                start[i] = i
+   ...:                end[i] = i + self.window_size
+   ...:        return start, end
+   ...:
+
+   In [3]: indexer = CustomIndexer(window_size=1, use_expanding=use_expanding)
+
+   In [4]: df.rolling(indexer).sum()
+   Out[4]:
+       values
+   0     0.0
+   1     1.0
+   2     3.0
+   3     3.0
+   4    10.0
+
+
 .. _stats.rolling_window.endpoints:
 
 Rolling window endpoints

doc/source/whatsnew/v1.0.0.rst

@@ -169,6 +169,16 @@ You can use the alias ``"boolean"`` as well.
    s = pd.Series([True, False, None], dtype="boolean")
    s
 
+.. _whatsnew_1000.custom_window:
+
+Defining custom windows for rolling operations
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+We've added a :func:`pandas.api.indexers.BaseIndexer` class that allows users to define how
+window bounds are created during ``rolling`` operations. Users can define their own ``get_window_bounds``
+method on a :func:`pandas.api.indexers.BaseIndexer` subclass that will generate the start and end
+indices used for each window during the rolling aggregation. For more details and example usage, see
+the :ref:`custom window rolling documentation <stats.custom_rolling_window>`
 
 .. _whatsnew_1000.enhancements.other:
 

pandas/_libs/window/aggregations.pyx

@@ -183,7 +183,8 @@ cdef inline void remove_sum(float64_t val, int64_t *nobs, float64_t *sum_x) nogi
 
 
 def roll_sum_variable(ndarray[float64_t] values, ndarray[int64_t] start,
-                      ndarray[int64_t] end, int64_t minp):
+                      ndarray[int64_t] end, int64_t minp,
+                      bint is_monotonic_bounds=True):
     cdef:
         float64_t sum_x = 0
         int64_t s, e
@@ -198,11 +199,10 @@ def roll_sum_variable(ndarray[float64_t] values, ndarray[int64_t] start,
             s = start[i]
             e = end[i]
 
-            if i == 0:
+            if i == 0 or not is_monotonic_bounds:
 
                 # setup
-                sum_x = 0.0
-                nobs = 0
+
                 for j in range(s, e):
                     add_sum(values[j], &nobs, &sum_x)
 
@@ -218,6 +218,10 @@ def roll_sum_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             output[i] = calc_sum(minp, nobs, sum_x)
 
+            if not is_monotonic_bounds:
+                for j in range(s, e):
+                    remove_sum(values[j], &nobs, &sum_x)
+
     return output
 
 
@@ -327,7 +331,8 @@ def roll_mean_fixed(ndarray[float64_t] values, ndarray[int64_t] start,
 
 
 def roll_mean_variable(ndarray[float64_t] values, ndarray[int64_t] start,
-                       ndarray[int64_t] end, int64_t minp):
+                       ndarray[int64_t] end, int64_t minp,
+                       bint is_monotonic_bounds=True):
     cdef:
         float64_t val, sum_x = 0
         int64_t s, e
@@ -342,11 +347,9 @@ def roll_mean_variable(ndarray[float64_t] values, ndarray[int64_t] start,
             s = start[i]
             e = end[i]
 
-            if i == 0:
+            if i == 0 or not is_monotonic_bounds:
 
                 # setup
-                sum_x = 0.0
-                nobs = 0
                 for j in range(s, e):
                     val = values[j]
                     add_mean(val, &nobs, &sum_x, &neg_ct)
@@ -365,6 +368,10 @@ def roll_mean_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             output[i] = calc_mean(minp, nobs, neg_ct, sum_x)
 
+            if not is_monotonic_bounds:
+                for j in range(s, e):
+                    val = values[j]
+                    remove_mean(val, &nobs, &sum_x, &neg_ct)
     return output
 
 # ----------------------------------------------------------------------
@@ -486,7 +493,8 @@ def roll_var_fixed(ndarray[float64_t] values, ndarray[int64_t] start,
 
 
 def roll_var_variable(ndarray[float64_t] values, ndarray[int64_t] start,
-                      ndarray[int64_t] end, int64_t minp, int ddof=1):
+                      ndarray[int64_t] end, int64_t minp, int ddof=1,
+                      bint is_monotonic_bounds=True):
     """
     Numerically stable implementation using Welford's method.
     """
@@ -508,7 +516,7 @@ def roll_var_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             # Over the first window, observations can only be added
             # never removed
-            if i == 0:
+            if i == 0 or not is_monotonic_bounds:
 
                 for j in range(s, e):
                     add_var(values[j], &nobs, &mean_x, &ssqdm_x)
@@ -528,6 +536,10 @@ def roll_var_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             output[i] = calc_var(minp, ddof, nobs, ssqdm_x)
 
+            if not is_monotonic_bounds:
+                for j in range(s, e):
+                    remove_var(values[j], &nobs, &mean_x, &ssqdm_x)
+
     return output
 
 # ----------------------------------------------------------------------
@@ -629,7 +641,8 @@ def roll_skew_fixed(ndarray[float64_t] values, ndarray[int64_t] start,
 
 
 def roll_skew_variable(ndarray[float64_t] values, ndarray[int64_t] start,
-                       ndarray[int64_t] end, int64_t minp):
+                       ndarray[int64_t] end, int64_t minp,
+                       bint is_monotonic_bounds=True):
     cdef:
         float64_t val, prev
         float64_t x = 0, xx = 0, xxx = 0
@@ -648,7 +661,7 @@ def roll_skew_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             # Over the first window, observations can only be added
             # never removed
-            if i == 0:
+            if i == 0 or not is_monotonic_bounds:
 
                 for j in range(s, e):
                     val = values[j]
@@ -671,6 +684,11 @@ def roll_skew_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             output[i] = calc_skew(minp, nobs, x, xx, xxx)
 
+            if not is_monotonic_bounds:
+                for j in range(s, e):
+                    val = values[j]
+                    remove_skew(val, &nobs, &x, &xx, &xxx)
+
     return output
 
 # ----------------------------------------------------------------------
@@ -776,7 +794,8 @@ def roll_kurt_fixed(ndarray[float64_t] values, ndarray[int64_t] start,
 
 
 def roll_kurt_variable(ndarray[float64_t] values, ndarray[int64_t] start,
-                       ndarray[int64_t] end, int64_t minp):
+                       ndarray[int64_t] end, int64_t minp,
+                       bint is_monotonic_bounds=True):
     cdef:
         float64_t val, prev
         float64_t x = 0, xx = 0, xxx = 0, xxxx = 0
@@ -794,7 +813,7 @@ def roll_kurt_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             # Over the first window, observations can only be added
             # never removed
-            if i == 0:
+            if i == 0 or not is_monotonic_bounds:
 
                 for j in range(s, e):
                     add_kurt(values[j], &nobs, &x, &xx, &xxx, &xxxx)
@@ -814,6 +833,10 @@ def roll_kurt_variable(ndarray[float64_t] values, ndarray[int64_t] start,
 
             output[i] = calc_kurt(minp, nobs, x, xx, xxx, xxxx)
 
+            if not is_monotonic_bounds:
+                for j in range(s, e):
+                    remove_kurt(values[j], &nobs, &x, &xx, &xxx, &xxxx)
+
     return output
 
 
@@ -1007,7 +1030,8 @@ def roll_min_fixed(ndarray[float64_t] values, ndarray[int64_t] start,
 
 
 def roll_min_variable(ndarray[float64_t] values, ndarray[int64_t] start,
-                      ndarray[int64_t] end, int64_t minp):
+                      ndarray[int64_t] end, int64_t minp,
+                      bint is_monotonic_bounds=True):
     """
     Moving max of 1d array of any numeric type along axis=0 ignoring NaNs.
 
@@ -1400,7 +1424,10 @@ def roll_generic_variable(object obj,
                           ndarray[int64_t] start, ndarray[int64_t] end,
                           int64_t minp,
                           int offset, object func, bint raw,
-                          object args, object kwargs):
+                          object args, object kwargs,
+                          bint is_monotonic_bounds=True):
+    # is_monotonic_bounds unused since variable algorithm doesn't calculate
+    # adds/subtracts across windows, but matches other *_variable functions
     cdef:
         ndarray[float64_t] output, counts, bufarr
         ndarray[float64_t, cast=True] arr