Skip to content

Removing gil for numba (mostly in peak detection) #4309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yger wants to merge 2 commits into
base: main
Choose a base branch
from
+18 −18
Open

Removing gil for numba (mostly in peak detection) #4309

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9519598
Removing gil for numba speedup
yger Jan 9, 2026
7cedd15
WIP
yger Jan 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 12 additions & 12 deletions src/spikeinterface/sortingcomponents/clustering/isosplit_isocut.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@
##########################
# isocut zone

@numba.jit(nopython=True)
@numba.jit(nopython=True, nogil=True)
def jisotonic5(x, weights):
N = x.shape[0]

Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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.
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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")
Expand Down
6 changes: 3 additions & 3 deletions src/spikeinterface/sortingcomponents/matching/tdc_peeler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand Down Expand Up @@ -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
):
Expand Down Expand Up @@ -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
):
Expand Down
4 changes: 2 additions & 2 deletions src/spikeinterface/sortingcomponents/peak_detection/locally_exclusive.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand All @@ -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
):
Expand Down
2 changes: 1 addition & 1 deletion src/spikeinterface/sortingcomponents/peak_detection/matched_filtering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand Down