ENH: Allow users to definite their own window bound calculations in rolling

pandas/_libs/window/indexers.pyx

@@ -1,140 +1,78 @@
 # cython: boundscheck=False, wraparound=False, cdivision=True
 
-from typing import Tuple
-
 import numpy as np
 from numpy cimport ndarray, int64_t
 
-# ----------------------------------------------------------------------
-# The indexer objects for rolling
-# These define start/end indexers to compute offsets
-
-
-class FixedWindowIndexer:
-    """
-    create a fixed length window indexer object
-    that has start & end, that point to offsets in
-    the index object; these are defined based on the win
-    arguments
-
-    Parameters
-    ----------
-    values: ndarray
-        values data array
-    win: int64_t
-        window size
-    index: object
-        index of the values
-    closed: string
-        closed behavior
-    """
-    def __init__(self, ndarray values, int64_t win, object closed, object index=None):
-        cdef:
-            ndarray[int64_t, ndim=1] start_s, start_e, end_s, end_e
-            int64_t N = len(values)
-
-        start_s = np.zeros(win, dtype='int64')
-        start_e = np.arange(win, N, dtype='int64') - win + 1
-        self.start = np.concatenate([start_s, start_e])[:N]
-
-        end_s = np.arange(win, dtype='int64') + 1
-        end_e = start_e + win
-        self.end = np.concatenate([end_s, end_e])[:N]
-
-    def get_window_bounds(self) -> Tuple[np.ndarray, np.ndarray]:
-        return self.start, self.end
-
-
-class VariableWindowIndexer:
-    """
-    create a variable length window indexer object
-    that has start & end, that point to offsets in
-    the index object; these are defined based on the win
-    arguments
-
-    Parameters
-    ----------
-    values: ndarray
-        values data array
-    win: int64_t
-        window size
-    index: ndarray
-        index of the values
-    closed: string
-        closed behavior
-    """
-    def __init__(self, ndarray values, int64_t win, object closed, ndarray index):
-        cdef:
-            bint left_closed = False
-            bint right_closed = False
-            int64_t N = len(index)
-
-        # if windows is variable, default is 'right', otherwise default is 'both'
-        if closed is None:
-            closed = 'right' if index is not None else 'both'
-
-        if closed in ['right', 'both']:
-            right_closed = True
-
-        if closed in ['left', 'both']:
-            left_closed = True
-
-        self.start, self.end = self.build(index, win, left_closed, right_closed, N)
-
-    @staticmethod
-    def build(const int64_t[:] index, int64_t win, bint left_closed,
-              bint right_closed, int64_t N) -> Tuple[np.ndarray, np.ndarray]:
-
-        cdef:
-            ndarray[int64_t] start, end
-            int64_t start_bound, end_bound
-            Py_ssize_t i, j
-
-        start = np.empty(N, dtype='int64')
-        start.fill(-1)
-        end = np.empty(N, dtype='int64')
-        end.fill(-1)
-
-        start[0] = 0
-
-        # right endpoint is closed
-        if right_closed:
-            end[0] = 1
-        # right endpoint is open
-        else:
-            end[0] = 0
-
-        with nogil:
-
-            # start is start of slice interval (including)
-            # end is end of slice interval (not including)
-            for i in range(1, N):
-                end_bound = index[i]
-                start_bound = index[i] - win
-
-                # left endpoint is closed
-                if left_closed:
-                    start_bound -= 1
-
-                # advance the start bound until we are
-                # within the constraint
-                start[i] = i
-                for j in range(start[i - 1], i):
-                    if index[j] > start_bound:
-                        start[i] = j
-                        break
-
-                # end bound is previous end
-                # or current index
-                if index[end[i - 1]] <= end_bound:
-                    end[i] = i + 1
-                else:
-                    end[i] = end[i - 1]
-
-                # right endpoint is open
-                if not right_closed:
-                    end[i] -= 1
-        return start, end
-
-    def get_window_bounds(self) -> Tuple[np.ndarray, np.ndarray]:
-        return self.start, self.end
+# Cython routines for window indexers
+
+def calculate_variable_window_bounds(
+    int64_t num_values,
+    int64_t window_size,
+    object min_periods,
+    object center,
+    object closed,
+    object win_type,
+    const int64_t[:] index
+):
+    cdef:
+        bint left_closed = False
+        bint right_closed = False
+        ndarray[int64_t, ndim=1] start, end
+        int64_t start_bound, end_bound
+        Py_ssize_t i, j
+
+    # if windows is variable, default is 'right', otherwise default is 'both'
+    if closed is None:
+        closed = 'right' if index is not None else 'both'
+
+    if closed in ['right', 'both']:
+        right_closed = True
+
+    if closed in ['left', 'both']:
+        left_closed = True
+
+    start = np.empty(num_values, dtype='int64')
+    start.fill(-1)
+    end = np.empty(num_values, dtype='int64')
+    end.fill(-1)
+
+    start[0] = 0
+
+    # right endpoint is closed
+    if right_closed:
+        end[0] = 1
+    # right endpoint is open
+    else:
+        end[0] = 0
+
+    with nogil:
+
+        # start is start of slice interval (including)
+        # end is end of slice interval (not including)
+        for i in range(1, num_values):
+            end_bound = index[i]
+            start_bound = index[i] - window_size
+
+            # left endpoint is closed
+            if left_closed:
+                start_bound -= 1
+
+            # advance the start bound until we are
+            # within the constraint
+            start[i] = i
+            for j in range(start[i - 1], i):
+                if index[j] > start_bound:
+                    start[i] = j
+                    break
+
+            # end bound is previous end
+            # or current index
+            if index[end[i - 1]] <= end_bound:
+                end[i] = i + 1
+            else:
+                end[i] = end[i - 1]
+
+            # right endpoint is open
+            if not right_closed:
+                end[i] -= 1
+    return start, end

pandas/api/__init__.py

@@ -1,2 +1,2 @@
 """ public toolkit API """
-from . import extensions, types  # noqa
+from . import extensions, indexers, types  # noqa

pandas/api/indexers/__init__.py

@@ -0,0 +1,2 @@
+"""Public API for Rolling Window Indexers"""
+from pandas.core.window.indexers import BaseIndexer  # noqa: F401

pandas/core/window/indexers.py

@@ -0,0 +1,103 @@
+"""Indexer objects for computing start/end window bounds for rolling operations"""
+from typing import Optional, Tuple
+
+import numpy as np
+
+from pandas._libs.window.indexers import calculate_variable_window_bounds
+from pandas.util._decorators import Appender
+
+get_window_bounds_doc = """
+Computes the bounds of a window.
+
+Parameters
+----------
+num_values : int, default 0
+    number of values that will be aggregated over
+window_size : int, default 0
+    the number of rows in a window
+min_periods : int, default None
+    min_periods passed from the top level rolling API
+center : bool, default None
+    center passed from the top level rolling API
+closed : str, default None
+    closed passed from the top level rolling API
+win_type : str, default None
+    win_type passed from the top level rolling API
+
+Returns
+-------
+A tuple of ndarray[int64]s, indicating the boundaries of each
+window
+"""
+
+
+class BaseIndexer:
+    """Base class for window bounds calculations"""
+
+    def __init__(
+        self, index: Optional[np.ndarray] = None, **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        **kwargs :
+            keyword argument that will be available when get_window_bounds is called
+        """
+        self.index = index
+        self.__dict__.update(kwargs)
+
+    @Appender(get_window_bounds_doc)
+    def get_window_bounds(
+        self,
+        num_values: int = 0,
+        window_size: int = 0,
+        min_periods: Optional[int] = None,
+        center: Optional[bool] = None,
+        closed: Optional[str] = None,
+        win_type: Optional[str] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+
+        raise NotImplementedError
+
+
+class FixedWindowIndexer(BaseIndexer):
+    """Creates window boundaries that are of fixed length."""
+
+    @Appender(get_window_bounds_doc)
+    def get_window_bounds(
+        self,
+        num_values: int = 0,
+        window_size: int = 0,
+        min_periods: Optional[int] = None,
+        center: Optional[bool] = None,
+        closed: Optional[str] = None,
+        win_type: Optional[str] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+
+        start_s = np.zeros(window_size, dtype="int64")
+        start_e = np.arange(window_size, num_values, dtype="int64") - window_size + 1
+        start = np.concatenate([start_s, start_e])[:num_values]
+
+        end_s = np.arange(window_size, dtype="int64") + 1
+        end_e = start_e + window_size
+        end = np.concatenate([end_s, end_e])[:num_values]
+        return start, end
+
+
+class VariableWindowIndexer(BaseIndexer):
+    """Creates window boundaries that are of variable length, namely for time series."""
+
+    @Appender(get_window_bounds_doc)
+    def get_window_bounds(
+        self,
+        num_values: int = 0,
+        window_size: int = 0,
+        min_periods: Optional[int] = None,
+        center: Optional[bool] = None,
+        closed: Optional[str] = None,
+        win_type: Optional[str] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+
+        return calculate_variable_window_bounds(
+            num_values, window_size, min_periods, center, closed, win_type, self.index
+        )