Update

pynapple-org · May 23, 2024 · 730aa31 · 730aa31
1 parent 8f328f4
commit 730aa31
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/pynapple/core/_jitted_functions.py b/pynapple/core/_jitted_functions.py
@@ -9,7 +9,7 @@
 def jitrestrict(time_array, starts, ends):
     n = len(time_array)
     m = len(starts)
-    ix = np.zeros(n, dtype="int")
+    ix = np.zeros(n, dtype=np.int64)
 
     k = 0
     t = 0
@@ -47,11 +47,12 @@ def jitrestrict(time_array, starts, ends):
 def jitrestrict_with_count(time_array, starts, ends):
     n = len(time_array)
     m = len(starts)
-    ix = np.zeros(n, dtype=np.bool_)
+    ix = np.zeros(n, dtype=np.int64)
     count = np.zeros(m, dtype=np.int64)
 
     k = 0
     t = 0
+    x = 0
 
     while ends[k] < time_array[t]:
         k += 1
@@ -60,9 +61,6 @@ def jitrestrict_with_count(time_array, starts, ends):
         # Outside
         while t < n:
             if time_array[t] >= starts[k]:
-                # ix[t] = True
-                # count[k] += 1
-                # t += 1
                 break
             t += 1
 
@@ -72,16 +70,17 @@ def jitrestrict_with_count(time_array, starts, ends):
                 k += 1
                 break
             else:
-                ix[t] = True
+                ix[x] = t
                 count[k] += 1
+                x += 1
             t += 1
 
         if k == m:
             break
         if t == n:
             break
 
-    return ix, count
+    return ix[0:x], count
 
 
 @jit(nopython=True)
@@ -160,6 +159,7 @@ def jitcount(time_array, starts, ends, bin_size):
                         break
 
                 lbound += bin_size
+                lbound = np.round(lbound, 9)
                 b += 1
         t = maxt
         k += 1