diff --git a/librosa/util/utils.py b/librosa/util/utils.py index 8eef988edb..94578db071 100644 --- a/librosa/util/utils.py +++ b/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`.