diff --git a/src/spikeinterface/sortingcomponents/clustering/isosplit_isocut.py b/src/spikeinterface/sortingcomponents/clustering/isosplit_isocut.py index 12ac25dd26..f6f0fd9c94 100644 --- a/src/spikeinterface/sortingcomponents/clustering/isosplit_isocut.py +++ b/src/spikeinterface/sortingcomponents/clustering/isosplit_isocut.py @@ -47,7 +47,7 @@ ########################## # isocut zone - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def jisotonic5(x, weights): N = x.shape[0] @@ -100,7 +100,7 @@ def jisotonic5(x, weights): return y, MSE - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def updown_arange(num_bins, dtype=np.int_): num_bins_1 = int(np.ceil(num_bins / 2)) num_bins_2 = num_bins - num_bins_1 @@ -111,7 +111,7 @@ def updown_arange(num_bins, dtype=np.int_): ) ) - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def compute_ks4(counts1, counts2): c1s = counts1.sum() c2s = counts2.sum() @@ -123,7 +123,7 @@ def compute_ks4(counts1, counts2): ks *= np.sqrt((c1s + c2s) / 2) return ks - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def compute_ks5(counts1, counts2): best_ks = -np.inf length = counts1.size @@ -138,7 +138,7 @@ def compute_ks5(counts1, counts2): return best_ks, best_length - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def up_down_isotonic_regression(x, weights=None): # determine switch point _, mse1 = jisotonic5(x, weights) @@ -153,14 +153,14 @@ def up_down_isotonic_regression(x, weights=None): return np.hstack((y1, y2)) - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def down_up_isotonic_regression(x, weights=None): return -up_down_isotonic_regression(-x, weights=weights) # num_bins_factor = 1 float_0 = np.array([0.0]) - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def isocut(samples): # , sample_weights=None isosplit6 not handle weight anymore """ Compute a dip-test to check if 1-d samples are unimodal or not. @@ -464,7 +464,7 @@ def ensure_continuous_labels(labels): if HAVE_NUMBA: - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def compute_centroids_and_covmats(X, centroids, covmats, labels, label_set, to_compute_mask): ## manual loop with numba to be faster @@ -498,7 +498,7 @@ def compute_centroids_and_covmats(X, centroids, covmats, labels, label_set, to_c if to_compute_mask[i] and count[i] > 0: covmats[i, :, :] /= count[i] - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def get_pairs_to_compare(centroids, comparisons_made, active_labels_mask): n = centroids.shape[0] @@ -526,7 +526,7 @@ def get_pairs_to_compare(centroids, comparisons_made, active_labels_mask): return pairs - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def compute_distances(centroids, comparisons_made, active_labels_mask): n = centroids.shape[0] dists = np.zeros((n, n), dtype=centroids.dtype) @@ -548,7 +548,7 @@ def compute_distances(centroids, comparisons_made, active_labels_mask): return dists - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def merge_test(X1, X2, centroid1, centroid2, covmat1, covmat2, isocut_threshold): if X1.size == 0 or X2.size == 0: @@ -584,7 +584,7 @@ def merge_test(X1, X2, centroid1, centroid2, covmat1, covmat2, isocut_threshold) return do_merge, L12 - @numba.jit(nopython=True) + @numba.jit(nopython=True, nogil=True) def compare_pairs(X, labels, pairs, centroids, covmats, min_cluster_size, isocut_threshold): clusters_changed_mask = np.zeros(centroids.shape[0], dtype="bool") diff --git a/src/spikeinterface/sortingcomponents/matching/tdc_peeler.py b/src/spikeinterface/sortingcomponents/matching/tdc_peeler.py index 6377fea566..031b112fd5 100644 --- a/src/spikeinterface/sortingcomponents/matching/tdc_peeler.py +++ b/src/spikeinterface/sortingcomponents/matching/tdc_peeler.py @@ -902,7 +902,7 @@ def fit_one_amplitude_with_neighbors( if HAVE_NUMBA: from numba import jit, prange - @jit(nopython=True) + @jit(nopython=True, nogil=True) def construct_prediction_sparse( spikes, traces, sparse_templates_array, template_sparsity_mask, wanted_channel_mask, nbefore, additive ): @@ -932,7 +932,7 @@ def construct_prediction_sparse( if template_sparsity_mask[cluster_index, chan]: chan_in_template += 1 - @jit(nopython=True) + @jit(nopython=True, nogil=True) def numba_sparse_distance( wf, sparse_templates_array, template_sparsity_mask, wanted_channel_mask, possible_clusters ): @@ -968,7 +968,7 @@ def numba_sparse_distance( distances[i] = sum_dist return distances - @jit(nopython=True) + @jit(nopython=True, nogil=True) def numba_best_shift_sparse( traces, sparse_template, sample_index, nbefore, possible_shifts, distances_shift, chan_sparsity ): diff --git a/src/spikeinterface/sortingcomponents/peak_detection/locally_exclusive.py b/src/spikeinterface/sortingcomponents/peak_detection/locally_exclusive.py index e68660ce83..dfc64e51b1 100644 --- a/src/spikeinterface/sortingcomponents/peak_detection/locally_exclusive.py +++ b/src/spikeinterface/sortingcomponents/peak_detection/locally_exclusive.py @@ -140,7 +140,7 @@ def detect_peaks_numba_locally_exclusive_on_chunk( return peak_sample_ind, peak_chan_ind - @numba.jit(nopython=True, parallel=False) + @numba.jit(nopython=True, parallel=False, nogil=True) def _numba_detect_peak_pos( traces, traces_center, peak_mask, exclude_sweep_size, abs_thresholds, peak_sign, neighbours_mask ): @@ -165,7 +165,7 @@ def _numba_detect_peak_pos( break return peak_mask - @numba.jit(nopython=True, parallel=False) + @numba.jit(nopython=True, parallel=False, nogil=True) def _numba_detect_peak_neg( traces, traces_center, peak_mask, exclude_sweep_size, abs_thresholds, peak_sign, neighbours_mask ): diff --git a/src/spikeinterface/sortingcomponents/peak_detection/matched_filtering.py b/src/spikeinterface/sortingcomponents/peak_detection/matched_filtering.py index 46e0709d9e..4c04bd979c 100644 --- a/src/spikeinterface/sortingcomponents/peak_detection/matched_filtering.py +++ b/src/spikeinterface/sortingcomponents/peak_detection/matched_filtering.py @@ -176,7 +176,7 @@ def get_convolved_traces(self, traces): if HAVE_NUMBA: import numba - @numba.jit(nopython=True, parallel=False) + @numba.jit(nopython=True, parallel=False, nogil=True) def _numba_detect_peak_matched_filtering( traces, traces_center, peak_mask, exclude_sweep_size, abs_thresholds, peak_sign, neighbours_mask, num_channels ):