Merge b77df89 into 27b099b

librosa/util/utils.py

@@ -12,6 +12,8 @@
 from .. import cache
 from .exceptions import ParameterError
 
+from numba import jit
+
 # Constrain STFT block sizes to 256 KB
 MAX_MEM_BLOCK = 2**8 * 2**10
 
@@ -897,13 +899,12 @@ def match_events(events_from, events_to, left=True, right=True):
     ParameterError
         If either array of input events is not the correct shape
     '''
-
     if len(events_from) == 0 or len(events_to) == 0:
         raise ParameterError('Attempting to match empty event list')
 
     # If we can't match left or right, then only strict equivalence
     # counts as a match.
-    if not (left or right) and not np.all(np.in1d(events_from, events_to)):
+    if not (left or right) and not set(events_from).issubset(set(events_to)):
             raise ParameterError('Cannot match events with left=right=False '
                                  'and events_from is not contained '
                                  'in events_to')
@@ -920,41 +921,67 @@ def match_events(events_from, events_to, left=True, right=True):
         raise ParameterError('Cannot match events with right=False '
                              'and min(events_to) > min(events_from)')
 
-    # Pre-allocate the output array
-    output = np.empty_like(events_from, dtype=np.int)
+    # array of matched items
+    output = np.empty_like(events_from, int)
 
-    # Compute how many rows we can process at once within the memory block
-    n_rows = int(MAX_MEM_BLOCK / (np.prod(output.shape[1:]) * len(events_to)
-                                  * events_from.itemsize))
+    return logic_match_events(output, events_from, events_to, left, right)
 
-    # Make sure we can at least make some progress
-    n_rows = max(1, n_rows)
 
-    # Iterate over blocks of the data
-    for bl_s in range(0, len(events_from), n_rows):
-        bl_t = min(bl_s + n_rows, len(events_from))
+@jit(nopython=True)
+def logic_match_events(output, events_from, events_to, left=True, right=True):
+    # mock dictionary for events
+    from_idx = np.argsort(events_from)
+    sorted_from = events_from[from_idx]
 
-        event_block = events_from[bl_s:bl_t]
+    to_idx = np.argsort(events_to)
+    sorted_to = events_to[to_idx]
 
-        # distance[i, j] = |events_from - events_to[j]|
-        distance = np.abs(np.subtract.outer(event_block,
-                                            events_to)).astype(np.float)
+    # find the matching indices
+    matching_indices = np.searchsorted(sorted_to, sorted_from)
 
-        # If we can't match to the right, squash all comparisons where
-        # events_to[j] > events_from[i]
-        if not right:
-            distance[np.less.outer(event_block, events_to)] = np.nan
+    # iterate over indices in matching_indices
+    for ind in range(0, len(matching_indices)):
+        left_ind = -1
+        right_ind = len(matching_indices)
 
-        # If we can't match to the left, squash all comparisons where
-        # events_to[j] < events_from[i]
-        if not left:
-            distance[np.greater.outer(event_block, events_to)] = np.nan
+        middle_ind = matching_indices[ind]
+        sorted_from_num = sorted_from[ind]
 
-        # Find the minimum distance point from whatever's left after squashing
-        output[bl_s:bl_t] = np.nanargmin(distance, axis=-1)
+        # Permitted to look to the left
+        if middle_ind > 0 and left:
+            left_ind = middle_ind - 1
 
-    return output
+        # Permitted to look to right
+        if middle_ind < len(matching_indices) - 1 and right:
+            right_ind = middle_ind + 1
+
+        # Check if left should be chosen
+        if ((left and left_ind != -1) and
+                (not right and sorted_to[middle_ind] > sorted_from_num) or
+                (abs(sorted_to[left_ind] - sorted_from_num) <
+                    abs(sorted_to[right_ind] - sorted_from_num) and
+                    abs(sorted_to[left_ind] - sorted_from_num) <
+                    abs(sorted_to[middle_ind] - sorted_from_num))):
+
+            output[ind] = to_idx[left_ind]
+
+        # Check if right should be chosen
+        elif (right and right_ind != len(matching_indices) and
+                 (abs(sorted_to[right_ind] - sorted_from_num) <
+                     abs(sorted_to[middle_ind] - sorted_from_num))):
+
+            output[ind] = to_idx[right_ind]
+
+        # Selected index wins
+        else:
+
+            output[ind] = to_idx[middle_ind]
+
+    # Undo sorting
+    solutions = np.empty_like(output)
+    solutions[from_idx] = output
 
+    return solutions
 
 def localmax(x, axis=0):
     """Find local maxima in an array `x`.