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