1"""Indexer objects for computing start/end window bounds for rolling operations"""
2
3from __future__ import annotations
4
5from datetime import timedelta
6from typing import TYPE_CHECKING
7
8import numpy as np
9
10from pandas._libs.tslibs import BaseOffset
11from pandas._libs.window.indexers import calculate_variable_window_bounds
12from pandas.util._decorators import set_module
13
14from pandas.core.dtypes.common import ensure_platform_int
15
16from pandas.core.indexes.datetimes import DatetimeIndex
17
18from pandas.tseries.offsets import Nano
19
20if TYPE_CHECKING:
21 from collections.abc import Hashable
22 from typing import Any
23
24
25@set_module("pandas.api.indexers")
26class BaseIndexer:
27 """
28 Base class for window bounds calculations.
29
30 Subclasses should implement ``get_window_bounds`` to define custom
31 windowing logic for use with :meth:`DataFrame.rolling` or
32 :meth:`Series.rolling`.
33
34 Parameters
35 ----------
36 index_array : np.ndarray, default None
37 Array-like structure representing the indices for the data points.
38 If None, the default indices are assumed. This can be useful for
39 handling non-uniform indices in data, such as in time series
40 with irregular timestamps.
41 window_size : int, default 0
42 Size of the moving window. This is the number of observations used
43 for calculating the statistic. The default is to consider all
44 observations within the window.
45 **kwargs
46 Additional keyword arguments passed to the subclass's methods.
47
48 See Also
49 --------
50 DataFrame.rolling : Provides rolling window calculations on dataframe.
51 Series.rolling : Provides rolling window calculations on series.
52
53 Examples
54 --------
55 >>> from pandas.api.indexers import BaseIndexer
56 >>> class CustomIndexer(BaseIndexer):
57 ... def get_window_bounds(self, num_values, min_periods, center, closed, step):
58 ... start = np.arange(num_values, dtype=np.int64)
59 ... end = np.arange(num_values, dtype=np.int64) + self.window_size
60 ... return start, end
61 >>> df = pd.DataFrame({"values": range(5)})
62 >>> indexer = CustomIndexer(window_size=2)
63 >>> df.rolling(indexer).sum()
64 values
65 0 1.0
66 1 3.0
67 2 5.0
68 3 7.0
69 4 4.0
70 """
71
72 def __init__(
73 self, index_array: np.ndarray | None = None, window_size: int = 0, **kwargs
74 ) -> None:
75 self.index_array = index_array
76 self.window_size = window_size
77 # Set user defined kwargs as attributes that can be used in get_window_bounds
78 for key, value in kwargs.items():
79 setattr(self, key, value)
80
81 def get_window_bounds(
82 self,
83 num_values: int = 0,
84 min_periods: int | None = None,
85 center: bool | None = None,
86 closed: str | None = None,
87 step: int | None = None,
88 ) -> tuple[np.ndarray, np.ndarray]:
89 """
90 Computes the bounds of a window.
91
92 Parameters
93 ----------
94 num_values : int, default 0
95 number of values that will be aggregated over
96 min_periods : int, default None
97 min_periods passed from the top level rolling API
98 center : bool, default None
99 center passed from the top level rolling API
100 closed : str, default None
101 closed passed from the top level rolling API
102 step : int, default None
103 step passed from the top level rolling API
104
105 Returns
106 -------
107 A tuple of ndarray[int64]s, indicating the boundaries of each
108 window
109 """
110 raise NotImplementedError
111
112
113class FixedWindowIndexer(BaseIndexer):
114 """Creates window boundaries that are of fixed length."""
115
116 def get_window_bounds(
117 self,
118 num_values: int = 0,
119 min_periods: int | None = None,
120 center: bool | None = None,
121 closed: str | None = None,
122 step: int | None = None,
123 ) -> tuple[np.ndarray, np.ndarray]:
124 """
125 Computes the bounds of a window.
126
127 Parameters
128 ----------
129 num_values : int, default 0
130 number of values that will be aggregated over
131 min_periods : int, default None
132 min_periods passed from the top level rolling API
133 center : bool, default None
134 center passed from the top level rolling API
135 closed : str, default None
136 closed passed from the top level rolling API
137 step : int, default None
138 step passed from the top level rolling API
139
140 Returns
141 -------
142 A tuple of ndarray[int64]s, indicating the boundaries of each
143 window
144 """
145 if center or self.window_size == 0:
146 offset = (self.window_size - 1) // 2
147 else:
148 offset = 0
149
150 end = np.arange(1 + offset, num_values + 1 + offset, step, dtype="int64")
151 start = end - self.window_size
152 if closed in ["left", "both"]:
153 start -= 1
154 if closed in ["left", "neither"]:
155 end -= 1
156
157 end = np.clip(end, 0, num_values)
158 start = np.clip(start, 0, num_values)
159
160 return start, end
161
162
163class VariableWindowIndexer(BaseIndexer):
164 """Creates window boundaries that are of variable length, namely for time series."""
165
166 def get_window_bounds(
167 self,
168 num_values: int = 0,
169 min_periods: int | None = None,
170 center: bool | None = None,
171 closed: str | None = None,
172 step: int | None = None,
173 ) -> tuple[np.ndarray, np.ndarray]:
174 """
175 Computes the bounds of a window.
176
177 Parameters
178 ----------
179 num_values : int, default 0
180 number of values that will be aggregated over
181 min_periods : int, default None
182 min_periods passed from the top level rolling API
183 center : bool, default None
184 center passed from the top level rolling API
185 closed : str, default None
186 closed passed from the top level rolling API
187 step : int, default None
188 step passed from the top level rolling API
189
190 Returns
191 -------
192 A tuple of ndarray[int64]s, indicating the boundaries of each
193 window
194 """
195 assert self.index_array is not None
196 if (index_length := len(self.index_array)) < num_values:
197 raise ValueError(
198 "Variable rolling window requires the index to be at least as long "
199 f"as the 'other' index. Got {index_length} < {num_values}. "
200 "Please align 'other' to the rolling object's index using "
201 "reindex_like() or similar method."
202 )
203 # error: Argument 4 to "calculate_variable_window_bounds" has incompatible
204 # type "Optional[bool]"; expected "bool"
205 return calculate_variable_window_bounds(
206 num_values,
207 self.window_size,
208 min_periods,
209 center, # type: ignore[arg-type]
210 closed,
211 self.index_array,
212 )
213
214
215@set_module("pandas.api.indexers")
216class VariableOffsetWindowIndexer(BaseIndexer):
217 """
218 Calculate window boundaries based on a non-fixed offset such as a BusinessDay.
219
220 Unlike the default fixed-size window, this indexer uses a
221 :class:`DateOffset` to determine variable-width window boundaries
222 based on the datetime index of the data.
223
224 Parameters
225 ----------
226 index_array : np.ndarray, default 0
227 Array-like structure specifying the indices for data points.
228 This parameter is currently not used.
229
230 window_size : int, optional, default 0
231 Specifies the number of data points in each window.
232 This parameter is currently not used.
233
234 index : DatetimeIndex, optional
235 ``DatetimeIndex`` of the labels of each observation.
236
237 offset : BaseOffset, optional
238 ``DateOffset`` representing the size of the window.
239
240 **kwargs
241 Additional keyword arguments passed to the parent class ``BaseIndexer``.
242
243 See Also
244 --------
245 api.indexers.BaseIndexer : Base class for all indexers.
246 DataFrame.rolling : Rolling window calculations on DataFrames.
247 offsets : Module providing various time offset classes.
248
249 Examples
250 --------
251 >>> from pandas.api.indexers import VariableOffsetWindowIndexer
252 >>> df = pd.DataFrame(range(10), index=pd.date_range("2020", periods=10))
253 >>> offset = pd.offsets.BDay(1)
254 >>> indexer = VariableOffsetWindowIndexer(index=df.index, offset=offset)
255 >>> df
256 0
257 2020-01-01 0
258 2020-01-02 1
259 2020-01-03 2
260 2020-01-04 3
261 2020-01-05 4
262 2020-01-06 5
263 2020-01-07 6
264 2020-01-08 7
265 2020-01-09 8
266 2020-01-10 9
267 >>> df.rolling(indexer).sum()
268 0
269 2020-01-01 0.0
270 2020-01-02 1.0
271 2020-01-03 2.0
272 2020-01-04 3.0
273 2020-01-05 7.0
274 2020-01-06 12.0
275 2020-01-07 6.0
276 2020-01-08 7.0
277 2020-01-09 8.0
278 2020-01-10 9.0
279 """
280
281 def __init__(
282 self,
283 index_array: np.ndarray | None = None,
284 window_size: int = 0,
285 index: DatetimeIndex | None = None,
286 offset: BaseOffset | None = None,
287 **kwargs,
288 ) -> None:
289 super().__init__(index_array, window_size, **kwargs)
290 if not isinstance(index, DatetimeIndex):
291 raise ValueError("index must be a DatetimeIndex.")
292 self.index = index
293 if not isinstance(offset, BaseOffset):
294 raise ValueError("offset must be a DateOffset-like object.")
295 self.offset = offset
296
297 def get_window_bounds(
298 self,
299 num_values: int = 0,
300 min_periods: int | None = None,
301 center: bool | None = None,
302 closed: str | None = None,
303 step: int | None = None,
304 ) -> tuple[np.ndarray, np.ndarray]:
305 """
306 Computes the bounds of a window.
307
308 Parameters
309 ----------
310 num_values : int, default 0
311 number of values that will be aggregated over
312 min_periods : int, default None
313 min_periods passed from the top level rolling API
314 center : bool, default None
315 center passed from the top level rolling API
316 closed : str, default None
317 closed passed from the top level rolling API
318 step : int, default None
319 step passed from the top level rolling API
320
321 Returns
322 -------
323 A tuple of ndarray[int64]s, indicating the boundaries of each
324 window
325 """
326 if step is not None:
327 raise NotImplementedError("step not implemented for variable offset window")
328 if num_values <= 0:
329 return np.empty(0, dtype="int64"), np.empty(0, dtype="int64")
330
331 # if windows is variable, default is 'right', otherwise default is 'both'
332 if closed is None:
333 closed = "right" if self.index is not None else "both"
334
335 right_closed = closed in ["right", "both"]
336 left_closed = closed in ["left", "both"]
337
338 if self.index[num_values - 1] < self.index[0]:
339 index_growth_sign = -1
340 else:
341 index_growth_sign = 1
342 offset_diff = index_growth_sign * self.offset
343
344 start = np.empty(num_values, dtype="int64")
345 start.fill(-1)
346 end = np.empty(num_values, dtype="int64")
347 end.fill(-1)
348
349 start[0] = 0
350
351 # right endpoint is closed
352 if right_closed:
353 end[0] = 1
354 # right endpoint is open
355 else:
356 end[0] = 0
357
358 zero = timedelta(0)
359 # start is start of slice interval (including)
360 # end is end of slice interval (not including)
361 for i in range(1, num_values):
362 end_bound = self.index[i]
363 start_bound = end_bound - offset_diff
364
365 # left endpoint is closed
366 if left_closed:
367 start_bound -= Nano(1)
368
369 # advance the start bound until we are
370 # within the constraint
371 start[i] = i
372 for j in range(start[i - 1], i):
373 start_diff = (self.index[j] - start_bound) * index_growth_sign
374 if start_diff > zero:
375 start[i] = j
376 break
377
378 # end bound is previous end
379 # or current index
380 end_diff = (self.index[end[i - 1]] - end_bound) * index_growth_sign
381 if end_diff == zero and not right_closed:
382 end[i] = end[i - 1] + 1
383 elif end_diff <= zero:
384 end[i] = i + 1
385 else:
386 end[i] = end[i - 1]
387
388 # right endpoint is open
389 if not right_closed:
390 end[i] -= 1
391
392 return start, end
393
394
395class ExpandingIndexer(BaseIndexer):
396 """Calculate expanding window bounds, mimicking df.expanding()"""
397
398 def get_window_bounds(
399 self,
400 num_values: int = 0,
401 min_periods: int | None = None,
402 center: bool | None = None,
403 closed: str | None = None,
404 step: int | None = None,
405 ) -> tuple[np.ndarray, np.ndarray]:
406 """
407 Computes the bounds of a window.
408
409 Parameters
410 ----------
411 num_values : int, default 0
412 number of values that will be aggregated over
413 min_periods : int, default None
414 min_periods passed from the top level rolling API
415 center : bool, default None
416 center passed from the top level rolling API
417 closed : str, default None
418 closed passed from the top level rolling API
419 step : int, default None
420 step passed from the top level rolling API
421
422 Returns
423 -------
424 A tuple of ndarray[int64]s, indicating the boundaries of each
425 window
426 """
427 return (
428 np.zeros(num_values, dtype=np.int64),
429 np.arange(1, num_values + 1, dtype=np.int64),
430 )
431
432
433@set_module("pandas.api.indexers")
434class FixedForwardWindowIndexer(BaseIndexer):
435 """
436 Creates window boundaries for fixed-length windows that include the current row.
437
438 This indexer produces forward-looking windows, where each window starts at
439 the current row and extends ``window_size`` rows ahead, unlike the default
440 backward-looking behavior of :meth:`DataFrame.rolling`.
441
442 Parameters
443 ----------
444 index_array : np.ndarray, default None
445 Array-like structure representing the indices for the data points.
446 If None, the default indices are assumed. This can be useful for
447 handling non-uniform indices in data, such as in time series
448 with irregular timestamps.
449 window_size : int, default 0
450 Size of the moving window. This is the number of observations used
451 for calculating the statistic. The default is to consider all
452 observations within the window.
453 **kwargs
454 Additional keyword arguments passed to the subclass's methods.
455
456 See Also
457 --------
458 DataFrame.rolling : Provides rolling window calculations.
459 api.indexers.VariableWindowIndexer : Calculate window bounds based on
460 variable-sized windows.
461
462 Examples
463 --------
464 >>> df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
465 >>> df
466 B
467 0 0.0
468 1 1.0
469 2 2.0
470 3 NaN
471 4 4.0
472
473 >>> indexer = pd.api.indexers.FixedForwardWindowIndexer(window_size=2)
474 >>> df.rolling(window=indexer, min_periods=1).sum()
475 B
476 0 1.0
477 1 3.0
478 2 2.0
479 3 4.0
480 4 4.0
481 """
482
483 def get_window_bounds(
484 self,
485 num_values: int = 0,
486 min_periods: int | None = None,
487 center: bool | None = None,
488 closed: str | None = None,
489 step: int | None = None,
490 ) -> tuple[np.ndarray, np.ndarray]:
491 """
492 Computes the bounds of a window.
493
494 Parameters
495 ----------
496 num_values : int, default 0
497 number of values that will be aggregated over
498 min_periods : int, default None
499 min_periods passed from the top level rolling API
500 center : bool, default None
501 center passed from the top level rolling API
502 closed : str, default None
503 closed passed from the top level rolling API
504 step : int, default None
505 step passed from the top level rolling API
506
507 Returns
508 -------
509 A tuple of ndarray[int64]s, indicating the boundaries of each
510 window
511 """
512 if center:
513 raise ValueError("Forward-looking windows can't have center=True")
514 if closed is not None:
515 raise ValueError(
516 "Forward-looking windows don't support setting the closed argument"
517 )
518 if step is None:
519 step = 1
520
521 start = np.arange(0, num_values, step, dtype="int64")
522 end = start + self.window_size
523 if self.window_size:
524 end = np.clip(end, 0, num_values)
525
526 return start, end
527
528
529class GroupByIndexer(BaseIndexer):
530 """Calculate bounds to compute groupby rolling, mimicking df.groupby().rolling()"""
531
532 def __init__(
533 self,
534 index_array: np.ndarray | None = None,
535 window_size: int | BaseIndexer = 0,
536 groupby_indices: dict[Hashable, np.ndarray] | None = None,
537 window_indexer: type[BaseIndexer] = BaseIndexer,
538 indexer_kwargs: dict[str, Any] | None = None,
539 **kwargs,
540 ) -> None:
541 """
542 Parameters
543 ----------
544 index_array : np.ndarray or None
545 np.ndarray of the index of the original object that we are performing
546 a chained groupby operation over. This index has been pre-sorted relative to
547 the groups
548 window_size : int or BaseIndexer
549 window size during the windowing operation
550 groupby_indices : dict or None
551 dict of {group label: [positional index of rows belonging to the group]}
552 window_indexer : BaseIndexer
553 BaseIndexer class determining the start and end bounds of each group
554 indexer_kwargs : dict or None
555 Custom kwargs to be passed to window_indexer
556 **kwargs :
557 keyword arguments that will be available when get_window_bounds is called
558 """
559 self.groupby_indices = groupby_indices or {}
560 self.window_indexer = window_indexer
561 self.indexer_kwargs = indexer_kwargs.copy() if indexer_kwargs else {}
562 super().__init__(
563 index_array=index_array,
564 window_size=self.indexer_kwargs.pop("window_size", window_size),
565 **kwargs,
566 )
567
568 def get_window_bounds(
569 self,
570 num_values: int = 0,
571 min_periods: int | None = None,
572 center: bool | None = None,
573 closed: str | None = None,
574 step: int | None = None,
575 ) -> tuple[np.ndarray, np.ndarray]:
576 """
577 Computes the bounds of a window.
578
579 Parameters
580 ----------
581 num_values : int, default 0
582 number of values that will be aggregated over
583 min_periods : int, default None
584 min_periods passed from the top level rolling API
585 center : bool, default None
586 center passed from the top level rolling API
587 closed : str, default None
588 closed passed from the top level rolling API
589 step : int, default None
590 step passed from the top level rolling API
591
592 Returns
593 -------
594 A tuple of ndarray[int64]s, indicating the boundaries of each
595 window
596 """
597 # 1) For each group, get the indices that belong to the group
598 # 2) Use the indices to calculate the start & end bounds of the window
599 # 3) Append the window bounds in group order
600 start_arrays = []
601 end_arrays = []
602 window_indices_start = 0
603 for indices in self.groupby_indices.values():
604 index_array: np.ndarray | None
605
606 if self.index_array is not None:
607 index_array = self.index_array.take(ensure_platform_int(indices))
608 else:
609 index_array = self.index_array
610 indexer = self.window_indexer(
611 index_array=index_array,
612 window_size=self.window_size,
613 **self.indexer_kwargs,
614 )
615 start, end = indexer.get_window_bounds(
616 len(indices), min_periods, center, closed, step
617 )
618 start = start.astype(np.int64)
619 end = end.astype(np.int64)
620 assert len(start) == len(end), (
621 "these should be equal in length from get_window_bounds"
622 )
623 # Cannot use groupby_indices as they might not be monotonic with the object
624 # we're rolling over
625 window_indices = np.arange(
626 window_indices_start, window_indices_start + len(indices)
627 )
628 window_indices_start += len(indices)
629 # Extend as we'll be slicing window like [start, end)
630 window_indices = np.append(window_indices, [window_indices[-1] + 1]).astype(
631 np.int64, copy=False
632 )
633 start_arrays.append(window_indices.take(ensure_platform_int(start)))
634 end_arrays.append(window_indices.take(ensure_platform_int(end)))
635 if len(start_arrays) == 0:
636 return np.array([], dtype=np.int64), np.array([], dtype=np.int64)
637 start = np.concatenate(start_arrays)
638 end = np.concatenate(end_arrays)
639 return start, end
640
641
642class ExponentialMovingWindowIndexer(BaseIndexer):
643 """Calculate ewm window bounds (the entire window)"""
644
645 def get_window_bounds(
646 self,
647 num_values: int = 0,
648 min_periods: int | None = None,
649 center: bool | None = None,
650 closed: str | None = None,
651 step: int | None = None,
652 ) -> tuple[np.ndarray, np.ndarray]:
653 """
654 Computes the bounds of a window.
655
656 Parameters
657 ----------
658 num_values : int, default 0
659 number of values that will be aggregated over
660 min_periods : int, default None
661 min_periods passed from the top level rolling API
662 center : bool, default None
663 center passed from the top level rolling API
664 closed : str, default None
665 closed passed from the top level rolling API
666 step : int, default None
667 step passed from the top level rolling API
668
669 Returns
670 -------
671 A tuple of ndarray[int64]s, indicating the boundaries of each
672 window
673 """
674 return np.array([0], dtype=np.int64), np.array([num_values], dtype=np.int64)