Hide keyboard shortcuts

Hot-keys on this page

r m x p   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""" 

2from typing import Optional, Tuple 

3 

4import numpy as np 

5 

6from pandas._libs.window.indexers import calculate_variable_window_bounds 

7from pandas.util._decorators import Appender 

8 

9get_window_bounds_doc = """ 

10Computes the bounds of a window. 

11 

12Parameters 

13---------- 

14num_values : int, default 0 

15 number of values that will be aggregated over 

16window_size : int, default 0 

17 the number of rows in a window 

18min_periods : int, default None 

19 min_periods passed from the top level rolling API 

20center : bool, default None 

21 center passed from the top level rolling API 

22closed : str, default None 

23 closed passed from the top level rolling API 

24win_type : str, default None 

25 win_type passed from the top level rolling API 

26 

27Returns 

28------- 

29A tuple of ndarray[int64]s, indicating the boundaries of each 

30window 

31""" 

32 

33 

34class BaseIndexer: 

35 """Base class for window bounds calculations""" 

36 

37 def __init__( 

38 self, index_array: Optional[np.ndarray] = None, window_size: int = 0, **kwargs, 

39 ): 

40 """ 

41 Parameters 

42 ---------- 

43 **kwargs : 

44 keyword arguments that will be available when get_window_bounds is called 

45 """ 

46 self.index_array = index_array 

47 self.window_size = window_size 

48 # Set user defined kwargs as attributes that can be used in get_window_bounds 

49 for key, value in kwargs.items(): 

50 setattr(self, key, value) 

51 

52 @Appender(get_window_bounds_doc) 

53 def get_window_bounds( 

54 self, 

55 num_values: int = 0, 

56 min_periods: Optional[int] = None, 

57 center: Optional[bool] = None, 

58 closed: Optional[str] = None, 

59 ) -> Tuple[np.ndarray, np.ndarray]: 

60 

61 raise NotImplementedError 

62 

63 

64class FixedWindowIndexer(BaseIndexer): 

65 """Creates window boundaries that are of fixed length.""" 

66 

67 @Appender(get_window_bounds_doc) 

68 def get_window_bounds( 

69 self, 

70 num_values: int = 0, 

71 min_periods: Optional[int] = None, 

72 center: Optional[bool] = None, 

73 closed: Optional[str] = None, 

74 ) -> Tuple[np.ndarray, np.ndarray]: 

75 

76 start_s = np.zeros(self.window_size, dtype="int64") 

77 start_e = ( 

78 np.arange(self.window_size, num_values, dtype="int64") 

79 - self.window_size 

80 + 1 

81 ) 

82 start = np.concatenate([start_s, start_e])[:num_values] 

83 

84 end_s = np.arange(self.window_size, dtype="int64") + 1 

85 end_e = start_e + self.window_size 

86 end = np.concatenate([end_s, end_e])[:num_values] 

87 return start, end 

88 

89 

90class VariableWindowIndexer(BaseIndexer): 

91 """Creates window boundaries that are of variable length, namely for time series.""" 

92 

93 @Appender(get_window_bounds_doc) 

94 def get_window_bounds( 

95 self, 

96 num_values: int = 0, 

97 min_periods: Optional[int] = None, 

98 center: Optional[bool] = None, 

99 closed: Optional[str] = None, 

100 ) -> Tuple[np.ndarray, np.ndarray]: 

101 

102 return calculate_variable_window_bounds( 

103 num_values, self.window_size, min_periods, center, closed, self.index_array, 

104 ) 

105 

106 

107class ExpandingIndexer(BaseIndexer): 

108 """Calculate expanding window bounds, mimicking df.expanding()""" 

109 

110 @Appender(get_window_bounds_doc) 

111 def get_window_bounds( 

112 self, 

113 num_values: int = 0, 

114 min_periods: Optional[int] = None, 

115 center: Optional[bool] = None, 

116 closed: Optional[str] = None, 

117 ) -> Tuple[np.ndarray, np.ndarray]: 

118 

119 return ( 

120 np.zeros(num_values, dtype=np.int64), 

121 np.arange(1, num_values + 1, dtype=np.int64), 

122 )