@@ -14,7 +14,7 @@ def calculate_variable_window_bounds(
14
14
int64_t num_values ,
15
15
int64_t window_size ,
16
16
object min_periods , # unused but here to match get_window_bounds signature
17
- object center , # unused but here to match get_window_bounds signature
17
+ bint center ,
18
18
object closed ,
19
19
const int64_t[:] index
20
20
):
@@ -32,8 +32,8 @@ def calculate_variable_window_bounds(
32
32
min_periods : object
33
33
ignored, exists for compatibility
34
34
35
- center : object
36
- ignored, exists for compatibility
35
+ center : bint
36
+ center the rolling window on the current observation
37
37
38
38
closed : str
39
39
string of side of the window that should be closed
@@ -46,7 +46,8 @@ def calculate_variable_window_bounds(
46
46
(ndarray[int64], ndarray[int64])
47
47
"""
48
48
cdef:
49
- bint left_closed = False , right_closed = False
49
+ bint left_closed = False
50
+ bint right_closed = False
50
51
ndarray[int64_t, ndim= 1 ] start, end
51
52
int64_t start_bound, end_bound, index_growth_sign = 1
52
53
Py_ssize_t i, j
@@ -77,14 +78,27 @@ def calculate_variable_window_bounds(
77
78
# right endpoint is open
78
79
else :
79
80
end[0 ] = 0
81
+ if center:
82
+ for j in range (0 , num_values + 1 ):
83
+ if (index[j] == index[0 ] + index_growth_sign * window_size / 2 and
84
+ right_closed):
85
+ end[0 ] = j + 1
86
+ break
87
+ elif index[j] >= index[0 ] + index_growth_sign * window_size / 2 :
88
+ end[0 ] = j
89
+ break
80
90
81
91
with nogil:
82
92
83
93
# start is start of slice interval (including)
84
94
# end is end of slice interval (not including)
85
95
for i in range (1 , num_values):
86
- end_bound = index[i]
87
- start_bound = index[i] - index_growth_sign * window_size
96
+ if center:
97
+ end_bound = index[i] + index_growth_sign * window_size / 2
98
+ start_bound = index[i] - index_growth_sign * window_size / 2
99
+ else :
100
+ end_bound = index[i]
101
+ start_bound = index[i] - index_growth_sign * window_size
88
102
89
103
# left endpoint is closed
90
104
if left_closed:
@@ -98,14 +112,27 @@ def calculate_variable_window_bounds(
98
112
start[i] = j
99
113
break
100
114
115
+ # for centered window advance the end bound until we are
116
+ # outside the constraint
117
+ if center:
118
+ for j in range (end[i - 1 ], num_values + 1 ):
119
+ if j == num_values:
120
+ end[i] = j
121
+ elif ((index[j] - end_bound) * index_growth_sign == 0 and
122
+ right_closed):
123
+ end[i] = j + 1
124
+ break
125
+ elif (index[j] - end_bound) * index_growth_sign >= 0 :
126
+ end[i] = j
127
+ break
101
128
# end bound is previous end
102
129
# or current index
103
- if (index[end[i - 1 ]] - end_bound) * index_growth_sign <= 0 :
130
+ elif (index[end[i - 1 ]] - end_bound) * index_growth_sign <= 0 :
104
131
end[i] = i + 1
105
132
else :
106
133
end[i] = end[i - 1 ]
107
134
108
135
# right endpoint is open
109
- if not right_closed:
136
+ if not right_closed and not center :
110
137
end[i] -= 1
111
138
return start, end
0 commit comments