Skip to content

Commit a25e793

Browse files
committed
Better window minp
1 parent b8d1f2b commit a25e793

File tree

1 file changed

+15
-16
lines changed

1 file changed

+15
-16
lines changed

pandas/_libs/window.pyx

+15-16
Original file line numberDiff line numberDiff line change
@@ -223,11 +223,12 @@ cdef class VariableWindowIndexer(WindowIndexer):
223223
224224
"""
225225
def __init__(self, ndarray input, int64_t win, int64_t minp,
226-
bint left_closed, bint right_closed, ndarray index):
226+
bint left_closed, bint right_closed, ndarray index,
227+
object floor=None):
227228

228229
self.is_variable = 1
229230
self.N = len(index)
230-
self.minp = _check_minp(win, minp, self.N)
231+
self.minp = _check_minp(win, minp, self.N, floor=floor)
231232

232233
self.start = np.empty(self.N, dtype='int64')
233234
self.start.fill(-1)
@@ -342,7 +343,7 @@ def get_window_indexer(input, win, minp, index, closed,
342343

343344
if index is not None:
344345
indexer = VariableWindowIndexer(input, win, minp, left_closed,
345-
right_closed, index)
346+
right_closed, index, floor=floor)
346347
elif use_mock:
347348
indexer = MockFixedWindowIndexer(input, win, minp, left_closed,
348349
right_closed, index, floor)
@@ -409,11 +410,10 @@ def roll_count(ndarray[double_t] input, int64_t win, int64_t minp,
409410
# Rolling sum
410411

411412

412-
cdef inline double calc_sum(int64_t minp, int64_t nobs, double sum_x,
413-
bint no_min=False) nogil:
413+
cdef inline double calc_sum(int64_t minp, int64_t nobs, double sum_x) nogil:
414414
cdef double result
415415

416-
if no_min or nobs >= minp:
416+
if nobs >= minp:
417417
result = sum_x
418418
else:
419419
result = NaN
@@ -444,17 +444,16 @@ def roll_sum(ndarray[double_t] input, int64_t win, int64_t minp,
444444
double val, prev_x, sum_x = 0
445445
int64_t s, e
446446
int64_t nobs = 0, i, j, N
447-
bint is_variable, no_min
447+
bint is_variable
448448
ndarray[int64_t] start, end
449449
ndarray[double_t] output
450450

451-
if minp == 0:
452-
no_min = True
453-
else:
454-
no_min = False
451+
print('minp', minp)
455452
start, end, N, win, minp, is_variable = get_window_indexer(input, win,
456453
minp, index,
457-
closed)
454+
closed,
455+
floor=0)
456+
print('minp', minp)
458457
output = np.empty(N, dtype=float)
459458

460459
# for performance we are going to iterate
@@ -488,27 +487,27 @@ def roll_sum(ndarray[double_t] input, int64_t win, int64_t minp,
488487
for j in range(end[i - 1], e):
489488
add_sum(input[j], &nobs, &sum_x)
490489

491-
output[i] = calc_sum(minp, nobs, sum_x, no_min)
490+
output[i] = calc_sum(minp, nobs, sum_x)
492491

493492
else:
494493

495494
# fixed window
496495

497496
with nogil:
498497

499-
for i in range(0, minp - 1):
498+
for i in range(0, min(1, minp) - 1):
500499
add_sum(input[i], &nobs, &sum_x)
501500
output[i] = NaN
502501

503-
for i in range(minp - 1, N):
502+
for i in range(min(1, minp) - 1, N):
504503
val = input[i]
505504
add_sum(val, &nobs, &sum_x)
506505

507506
if i > win - 1:
508507
prev_x = input[i - win]
509508
remove_sum(prev_x, &nobs, &sum_x)
510509

511-
output[i] = calc_sum(minp, nobs, sum_x, no_min)
510+
output[i] = calc_sum(minp, nobs, sum_x)
512511

513512
return output
514513

0 commit comments

Comments
 (0)