Better window minp

pandas-dev · Dec 27, 2017 · a25e793 · a25e793
1 parent b8d1f2b
commit a25e793
Showing 1 changed file with 15 additions and 16 deletions.
diff --git a/pandas/_libs/window.pyx b/pandas/_libs/window.pyx
@@ -223,11 +223,12 @@ cdef class VariableWindowIndexer(WindowIndexer):
 
     """
     def __init__(self, ndarray input, int64_t win, int64_t minp,
-                 bint left_closed, bint right_closed, ndarray index):
+                 bint left_closed, bint right_closed, ndarray index,
+                 object floor=None):
 
         self.is_variable = 1
         self.N = len(index)
-        self.minp = _check_minp(win, minp, self.N)
+        self.minp = _check_minp(win, minp, self.N, floor=floor)
 
         self.start = np.empty(self.N, dtype='int64')
         self.start.fill(-1)
@@ -342,7 +343,7 @@ def get_window_indexer(input, win, minp, index, closed,
 
     if index is not None:
         indexer = VariableWindowIndexer(input, win, minp, left_closed,
-                                        right_closed, index)
+                                        right_closed, index, floor=floor)
     elif use_mock:
         indexer = MockFixedWindowIndexer(input, win, minp, left_closed,
                                          right_closed, index, floor)
@@ -409,11 +410,10 @@ def roll_count(ndarray[double_t] input, int64_t win, int64_t minp,
 # Rolling sum
 
 
-cdef inline double calc_sum(int64_t minp, int64_t nobs, double sum_x,
-                            bint no_min=False) nogil:
+cdef inline double calc_sum(int64_t minp, int64_t nobs, double sum_x) nogil:
     cdef double result
 
-    if no_min or nobs >= minp:
+    if nobs >= minp:
         result = sum_x
     else:
         result = NaN
@@ -444,17 +444,16 @@ def roll_sum(ndarray[double_t] input, int64_t win, int64_t minp,
         double val, prev_x, sum_x = 0
         int64_t s, e
         int64_t nobs = 0, i, j, N
-        bint is_variable, no_min
+        bint is_variable
         ndarray[int64_t] start, end
         ndarray[double_t] output
 
-    if minp == 0:
-        no_min = True
-    else:
-        no_min = False
+    print('minp', minp)
     start, end, N, win, minp, is_variable = get_window_indexer(input, win,
                                                                minp, index,
-                                                               closed)
+                                                               closed,
+                                                               floor=0)
+    print('minp', minp)
     output = np.empty(N, dtype=float)
 
     # for performance we are going to iterate
@@ -488,27 +487,27 @@ def roll_sum(ndarray[double_t] input, int64_t win, int64_t minp,
                     for j in range(end[i - 1], e):
                         add_sum(input[j], &nobs, &sum_x)
 
-                output[i] = calc_sum(minp, nobs, sum_x, no_min)
+                output[i] = calc_sum(minp, nobs, sum_x)
 
     else:
 
         # fixed window
 
         with nogil:
 
-            for i in range(0, minp - 1):
+            for i in range(0, min(1, minp) - 1):
                 add_sum(input[i], &nobs, &sum_x)
                 output[i] = NaN
 
-            for i in range(minp - 1, N):
+            for i in range(min(1, minp) - 1, N):
                 val = input[i]
                 add_sum(val, &nobs, &sum_x)
 
                 if i > win - 1:
                     prev_x = input[i - win]
                     remove_sum(prev_x, &nobs, &sum_x)
 
-                output[i] = calc_sum(minp, nobs, sum_x, no_min)
+                output[i] = calc_sum(minp, nobs, sum_x)
 
     return output