From 0aed532327bc0dabe0e3b7d15719252ca146c231 Mon Sep 17 00:00:00 2001 From: Charlie Windolf Date: Wed, 21 Oct 2026 15:17:34 -0400 Subject: [PATCH 1/5] comparison: fix an assert that sometimes triggered in score labels, add a test, and add a numba version matching the count matrix impl because this would otherwise require a hot python loop --- .../comparison/comparisontools.py | 95 +++++++++++++++---- .../comparison/tests/test_comparisontools.py | 29 ++++++ 2 files changed, 105 insertions(+), 19 deletions(-) diff --git a/src/spikeinterface/comparison/comparisontools.py b/src/spikeinterface/comparison/comparisontools.py index 4892014575..9a8757292e 100644 --- a/src/spikeinterface/comparison/comparisontools.py +++ b/src/spikeinterface/comparison/comparisontools.py @@ -565,6 +565,75 @@ def make_hungarian_match(agreement_scores, min_score): return hungarian_match_12, hungarian_match_21 +def get_optimized_compute_pair_match_indices(): + """ + Like get_optimized_compute_matching_matrix, but for do_score_labels + Cache a compiled hot loop for that function + """ + + if hasattr(get_optimized_compute_pair_match_indices, "_cached_function"): + return get_optimized_compute_pair_match_indices._cached_function + + from numba import jit + + @jit(nopython=True, nogil=True) + def compute_pair_match_indices(spike_frames1, spike_frames2, delta_frames): + """ + Internal function used by `do_score_labels()` to match the spikes of one unit pair in one segment, + using each spike at most once. + + Each spike of the first train, in order, is matched to the earliest unmatched spike of the second train + within `delta_frames`. The matched indices of the second train are increasing, so a spike of the second + train is unmatched if and only if its index is greater than that of the last match. + + Parameters + ---------- + spike_frames1 : ndarray + An array of integer frame numbers for the first train. Must be in ascending order. + spike_frames2 : ndarray + An array of integer frame numbers for the second train. Must be in ascending order. + delta_frames : int + The inclusive upper limit on the frame difference for which two spikes are considered matching. + + Returns + ------- + match_indices1, match_indices2 : ndarray + Indices of the matched spikes in each train: `spike_frames1[match_indices1[i]]` is matched to + `spike_frames2[match_indices2[i]]`. + """ + match_indices1 = np.zeros(min(spike_frames1.size, spike_frames2.size), dtype=np.int64) + match_indices2 = np.zeros_like(match_indices1) + num_matches = 0 + + # Keeps track of which frame in the second spike train should be used as a search start for matches + second_train_search_start = 0 + for index1 in range(spike_frames1.size): + frame1 = spike_frames1[index1] + + for index2 in range(second_train_search_start, spike_frames2.size): + frame2 = spike_frames2[index2] + if frame2 < frame1 - delta_frames: + # no match move the left limit for the next loop + second_train_search_start += 1 + continue + elif frame2 > frame1 + delta_frames: + # no match stop search in train2 and continue increment in train1 + break + elif num_matches == 0 or index2 > match_indices2[num_matches - 1]: + # match with a spike that is not matched yet + match_indices1[num_matches] = index1 + match_indices2[num_matches] = index2 + num_matches += 1 + break + + return match_indices1[:num_matches], match_indices2[:num_matches] + + # Cache the compiled function + get_optimized_compute_pair_match_indices._cached_function = compute_pair_match_indices + + return compute_pair_match_indices + + def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclassification=False): """ Makes the labeling at spike level for each spike train: @@ -596,6 +665,8 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass labels_st2 : dict of lists of np.array of str Contain score labels for units of sorting 2 for each segment """ + compute_pair_match_indices = get_optimized_compute_pair_match_indices() + unit1_ids = sorting1.get_unit_ids() unit2_ids = sorting2.get_unit_ids() labels_st1 = dict() @@ -620,25 +691,11 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass if u2 != -1 and u2 != "": lab_st1 = labels_st1[u1][seg_index] lab_st2 = labels_st2[u2][seg_index] - mapped_st = sorting2.get_unit_spike_train(u2, seg_index) - times_concat = np.concatenate((sts, mapped_st)) - membership = np.concatenate((np.ones(sts.shape) * 1, np.ones(mapped_st.shape) * 2)) - indices = times_concat.argsort() - times_concat_sorted = times_concat[indices] - membership_sorted = membership[indices] - diffs = times_concat_sorted[1:] - times_concat_sorted[:-1] - inds = np.where((diffs <= delta_frames) & (membership_sorted[:-1] != membership_sorted[1:]))[0] - if len(inds) > 0: - inds2 = inds[np.where(inds[:-1] + 1 != inds[1:])[0]] + 1 - inds2 = np.concatenate((inds2, [inds[-1]])) - times_matched = times_concat_sorted[inds2] - # find and label closest spikes - ind_st1 = np.array([np.abs(sts1[u1] - tm).argmin() for tm in times_matched]) - ind_st2 = np.array([np.abs(mapped_st - tm).argmin() for tm in times_matched]) - assert len(np.unique(ind_st1)) == len(ind_st1) - assert len(np.unique(ind_st2)) == len(ind_st2) - lab_st1[ind_st1] = "TP" - lab_st2[ind_st2] = "TP" + mapped_st = sts2[u2][seg_index] + assert np.all(np.diff(sts) >= 0) and np.all(np.diff(mapped_st) >= 0) + match_indices1, match_indices2 = compute_pair_match_indices(sts, mapped_st, delta_frames) + lab_st1[match_indices1] = "TP" + lab_st2[match_indices2] = "TP" else: lab_st1 = np.array(["FN"] * len(sts)) labels_st1[u1][seg_index] = lab_st1 diff --git a/src/spikeinterface/comparison/tests/test_comparisontools.py b/src/spikeinterface/comparison/tests/test_comparisontools.py index 43288415d1..dce5437226 100644 --- a/src/spikeinterface/comparison/tests/test_comparisontools.py +++ b/src/spikeinterface/comparison/tests/test_comparisontools.py @@ -374,6 +374,35 @@ def test_do_score_labels(): ) +def test_do_score_labels_counts_spikes_once(): + # checks that spikes in each sorting are only matched once across the spike trains by do_score_labels + # for posterity: right now, make_match_count_matrix() can still count a sorting2 + # spike twice. it tries to fix that by capping each count to the total spike count, + # but that has edge cases. + # TODO: count each spike once in compute_matching_matrix() + delta_frames = 10 + + sorting1, sorting2 = make_sorting([100, 112, 500], [0, 0, 1], [101, 104, 900], [0, 0, 1]) + unit_map12 = {0: 0, 1: 1} + labels_st1, labels_st2 = do_score_labels(sorting1, sorting2, delta_frames, unit_map12) + assert_array_equal(labels_st1[0][0], ["TP", "TP"]) + assert_array_equal(labels_st2[0][0], ["TP", "TP"]) + assert_array_equal(labels_st1[1][0], ["FN"]) + assert_array_equal(labels_st2[1][0], ["FP"]) + + # these sortings work now in make_match_count_matrix + counts = make_match_count_matrix(sorting1, sorting2, delta_frames=delta_frames) + for u1, u2 in unit_map12.items(): + assert np.sum(labels_st1[u1][0] == "TP") == counts.at[u1, u2] + assert np.sum(labels_st2[u2][0] == "TP") == counts.at[u1, u2] + + # note: can't test make_match_count_matrix here, it over-counts this guy + sorting1, sorting2 = make_sorting([0, 1, 2], [0, 0, 0], [1, 5, 1000], [0, 0, 0]) + labels_st1, labels_st2 = do_score_labels(sorting1, sorting2, delta_frames, {0: 0}) + assert_array_equal(labels_st1[0][0], ["TP", "TP", "FN"]) + assert_array_equal(labels_st2[0][0], ["TP", "TP", "FP"]) + + def test_compare_spike_trains(): sorting1, sorting2 = make_sorting( [100, 200, 300, 400], From 1f5013da01c9c92095ecfb6602d913adbc8ce222 Mon Sep 17 00:00:00 2001 From: Charlie Windolf Date: Thu, 12 Nov 2026 10:16:32 -0500 Subject: [PATCH 2/5] use get_unit_spike_train Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> --- src/spikeinterface/comparison/comparisontools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/comparison/comparisontools.py b/src/spikeinterface/comparison/comparisontools.py index 9a8757292e..d549ca08aa 100644 --- a/src/spikeinterface/comparison/comparisontools.py +++ b/src/spikeinterface/comparison/comparisontools.py @@ -691,7 +691,7 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass if u2 != -1 and u2 != "": lab_st1 = labels_st1[u1][seg_index] lab_st2 = labels_st2[u2][seg_index] - mapped_st = sts2[u2][seg_index] + mapped_st = sorting2.get_unit_spike_train(u2, seg_index) assert np.all(np.diff(sts) >= 0) and np.all(np.diff(mapped_st) >= 0) match_indices1, match_indices2 = compute_pair_match_indices(sts, mapped_st, delta_frames) lab_st1[match_indices1] = "TP" From ee3daff75c9620a725b3b55537a073032d492294 Mon Sep 17 00:00:00 2001 From: Charlie Windolf Date: Thu, 12 Nov 2026 13:23:19 -0500 Subject: [PATCH 3/5] comparison: rework score labels a bit more for style and to use get_unit_spike_train This also reorganizes some of the loops which computed the same stuff repeatedly without achieving a readability benefit from doing so, you could say. Tried to make it feel more linear, although the logic in label_misclassification could still be a bit of a head scratcher. So I added a small test case about misclassification which wasn't covered. --- .../comparison/comparisontools.py | 63 ++++++++++--------- .../comparison/tests/test_comparisontools.py | 26 +++++--- 2 files changed, 50 insertions(+), 39 deletions(-) diff --git a/src/spikeinterface/comparison/comparisontools.py b/src/spikeinterface/comparison/comparisontools.py index d549ca08aa..ed5e9a76f6 100644 --- a/src/spikeinterface/comparison/comparisontools.py +++ b/src/spikeinterface/comparison/comparisontools.py @@ -667,56 +667,57 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass """ compute_pair_match_indices = get_optimized_compute_pair_match_indices() + # initialize labels dict of lists for each sorting with UNPAIRED + num_segments = sorting1.get_num_segments() unit1_ids = sorting1.get_unit_ids() unit2_ids = sorting2.get_unit_ids() - labels_st1 = dict() - labels_st2 = dict() - - # copy spike trains for faster access from extractors with memmapped data - num_segments = sorting1.get_num_segments() - sts1 = {u1: [sorting1.get_unit_spike_train(u1, seg_index) for seg_index in range(num_segments)] for u1 in unit1_ids} - sts2 = {u2: [sorting2.get_unit_spike_train(u2, seg_index) for seg_index in range(num_segments)] for u2 in unit2_ids} - - for u1 in unit1_ids: - lab_st1 = [np.array(["UNPAIRED"] * len(sts), dtype="= 0) and np.all(np.diff(mapped_st) >= 0) - match_indices1, match_indices2 = compute_pair_match_indices(sts, mapped_st, delta_frames) + st2 = sorting2.get_unit_spike_train(u2, seg_index) + match_indices1, match_indices2 = compute_pair_match_indices(st1, st2, delta_frames) lab_st1[match_indices1] = "TP" lab_st2[match_indices2] = "TP" else: - lab_st1 = np.array(["FN"] * len(sts)) + lab_st1 = np.array(["FN"] * len(st1)) labels_st1[u1][seg_index] = lab_st1 if label_misclassification: + mapped_unit2_ids = [u2 for u2 in unit2_ids if u2 in unit_map12.values] for seg_index in range(num_segments): for u1 in unit1_ids: + if unit_map12[u1] == -1 or unit_map12[u1] == "": + continue lab_st1 = labels_st1[u1][seg_index] - st1 = sts1[u1][seg_index] + st1 = sorting1.get_unit_spike_train(u1, seg_index) for l_gt, lab in enumerate(lab_st1): - if lab == "UNPAIRED": - for u2 in unit2_ids: - if u2 in unit_map12.values and unit_map12[u1] != -1: - lab_st2 = labels_st2[u2][seg_index] - n_sp = st1[l_gt] - mapped_st = sts2[u2][seg_index] - matches = np.abs(mapped_st.astype(int) - n_sp) <= delta_frames - if np.sum(matches) > 0: - if "CL" not in lab_st1[l_gt] and "CL" not in lab_st2[np.where(matches)[0][0]]: - lab_st1[l_gt] = "CL_" + str(u1) + "_" + str(u2) - lab_st2[np.where(matches)[0][0]] = "CL_" + str(u2) + "_" + str(u1) + if lab != "UNPAIRED": + continue + n_sp = st1[l_gt] + for u2 in mapped_unit2_ids: + lab_st2 = labels_st2[u2][seg_index] + st2 = sorting2.get_unit_spike_train(u2, seg_index) + matches = np.flatnonzero(np.abs(st2.astype(int) - n_sp) <= delta_frames) + if not matches.size: + continue + if "CL" not in lab_st2[matches[0]]: + lab_st1[l_gt] = "CL_" + str(u1) + "_" + str(u2) + lab_st2[matches[0]] = "CL_" + str(u2) + "_" + str(u1) + break for seg_index in range(num_segments): for u1 in unit1_ids: diff --git a/src/spikeinterface/comparison/tests/test_comparisontools.py b/src/spikeinterface/comparison/tests/test_comparisontools.py index dce5437226..b5eee7c129 100644 --- a/src/spikeinterface/comparison/tests/test_comparisontools.py +++ b/src/spikeinterface/comparison/tests/test_comparisontools.py @@ -1,21 +1,22 @@ import numpy as np +import pandas as pd from numpy.testing import assert_array_equal -from spikeinterface.extractors import NumpySorting from spikeinterface.comparison import ( + compare_spike_trains, + compute_performance, + do_confusion_matrix, do_count_event, - make_match_count_matrix, + do_count_score, + do_score_labels, make_agreement_scores, - make_possible_match, make_best_match, make_hungarian_match, - do_score_labels, - compare_spike_trains, - do_confusion_matrix, - do_count_score, - compute_performance, + make_match_count_matrix, + make_possible_match, ) from spikeinterface.core.generate import generate_sorting +from spikeinterface.extractors import NumpySorting def make_sorting(times1, labels1, times2, labels2): @@ -402,6 +403,15 @@ def test_do_score_labels_counts_spikes_once(): assert_array_equal(labels_st1[0][0], ["TP", "TP", "FN"]) assert_array_equal(labels_st2[0][0], ["TP", "TP", "FP"]) + # test label_misclassification. the spike from st1u0 is where most of the action is at. + # it should get CL'd to one spike, and that should be the one from st2u1 at time 503. + sorting1, sorting2 = make_sorting([100, 500, 900, 1300], [0, 0, 1, 2], [102, 503, 505, 901, 1301], [0, 1, 2, 1, 2]) + unit_map12 = pd.Series({0: 0, 1: 1, 2: 2}) + labels_st1, labels_st2 = do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclassification=True) + assert_array_equal(labels_st1[0][0], ["TP", "CL_0_1"]) + assert_array_equal(labels_st2[1][0], ["CL_1_0", "TP"]) + assert_array_equal(labels_st2[2][0], ["FP", "TP"]) + def test_compare_spike_trains(): sorting1, sorting2 = make_sorting( From 4c3146fee84e2a292bb86422eb641dcac9c19b3c Mon Sep 17 00:00:00 2001 From: Charlie Windolf Date: Fri, 13 Nov 2026 10:16:27 -0500 Subject: [PATCH 4/5] score labels: there was a bad edge case, make sure tests cover it --- .../comparison/comparisontools.py | 19 ++++++++----------- .../comparison/tests/test_comparisontools.py | 18 +++++++++++++----- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/src/spikeinterface/comparison/comparisontools.py b/src/spikeinterface/comparison/comparisontools.py index 4b580d81ac..55e3424af9 100644 --- a/src/spikeinterface/comparison/comparisontools.py +++ b/src/spikeinterface/comparison/comparisontools.py @@ -682,17 +682,15 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass for seg_index in range(num_segments): for u1 in unit1_ids: u2 = unit_map12[u1] + if u2 == -1 or u2 == "": + continue st1 = sorting1.get_unit_spike_train(u1, seg_index) - if u2 != -1 and u2 != "": - lab_st1 = labels_st1[u1][seg_index] - lab_st2 = labels_st2[u2][seg_index] - st2 = sorting2.get_unit_spike_train(u2, seg_index) - match_indices1, match_indices2 = compute_pair_match_indices(st1, st2, delta_frames) - lab_st1[match_indices1] = "TP" - lab_st2[match_indices2] = "TP" - else: - lab_st1 = np.array(["FN"] * len(st1)) - labels_st1[u1][seg_index] = lab_st1 + lab_st1 = labels_st1[u1][seg_index] + lab_st2 = labels_st2[u2][seg_index] + st2 = sorting2.get_unit_spike_train(u2, seg_index) + match_indices1, match_indices2 = compute_pair_match_indices(st1, st2, delta_frames) + lab_st1[match_indices1] = "TP" + lab_st2[match_indices2] = "TP" if label_misclassification: mapped_unit2_ids = [u2 for u2 in unit2_ids if u2 in unit_map12.values] @@ -721,7 +719,6 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass for u1 in unit1_ids: lab_st1 = labels_st1[u1][seg_index] lab_st1[lab_st1 == "UNPAIRED"] = "FN" - for u2 in unit2_ids: lab_st2 = labels_st2[u2][seg_index] lab_st2[lab_st2 == "UNPAIRED"] = "FP" diff --git a/src/spikeinterface/comparison/tests/test_comparisontools.py b/src/spikeinterface/comparison/tests/test_comparisontools.py index e4bb27e824..ab208335e9 100644 --- a/src/spikeinterface/comparison/tests/test_comparisontools.py +++ b/src/spikeinterface/comparison/tests/test_comparisontools.py @@ -19,10 +19,10 @@ from spikeinterface.extractors import NumpySorting -def make_sorting(times1, labels1, times2, labels2): +def make_sorting(times1, labels1, times2, labels2, ids1=None, ids2=None): sampling_frequency = 30000.0 - sorting1 = NumpySorting.from_samples_and_labels([times1], [labels1], sampling_frequency) - sorting2 = NumpySorting.from_samples_and_labels([times2], [labels2], sampling_frequency) + sorting1 = NumpySorting.from_samples_and_labels([times1], [labels1], sampling_frequency, unit_ids=ids1) + sorting2 = NumpySorting.from_samples_and_labels([times2], [labels2], sampling_frequency, unit_ids=ids2) return sorting1, sorting2 @@ -408,12 +408,20 @@ def test_do_score_labels_counts_spikes_once(): # test label_misclassification. the spike from st1u0 is where most of the action is at. # it should get CL'd to one spike, and that should be the one from st2u1 at time 503. - sorting1, sorting2 = make_sorting([100, 500, 900, 1300], [0, 0, 1, 2], [102, 503, 505, 901, 1301], [0, 1, 2, 1, 2]) - unit_map12 = pd.Series({0: 0, 1: 1, 2: 2}) + # also check empty unit handling. unit 3 in both trains are empty, unit 4 are empty too + # but are mapped to each other. + sorting1, sorting2 = make_sorting( + [100, 500, 900, 1300], [0, 0, 1, 2], [102, 503, 505, 901, 1301], [0, 1, 2, 1, 2], ids1=[0, 1, 2, 3, 4], ids2=[0, 1, 2, 3, 4] + ) + unit_map12 = pd.Series({0: 0, 1: 1, 2: 2, 3: -1, 4: 4}) labels_st1, labels_st2 = do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclassification=True) assert_array_equal(labels_st1[0][0], ["TP", "CL_0_1"]) assert_array_equal(labels_st2[1][0], ["CL_1_0", "TP"]) assert_array_equal(labels_st2[2][0], ["FP", "TP"]) + assert labels_st1[3][0].size == 0 + assert labels_st2[3][0].size == 0 + assert labels_st1[4][0].size == 0 + assert labels_st2[4][0].size == 0 def test_compare_spike_trains(): From 400bad850a366e70499c767dc3bbc9580e743833 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Nov 2026 15:16:59 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../comparison/tests/test_comparisontools.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/comparison/tests/test_comparisontools.py b/src/spikeinterface/comparison/tests/test_comparisontools.py index ab208335e9..5af7e61ac7 100644 --- a/src/spikeinterface/comparison/tests/test_comparisontools.py +++ b/src/spikeinterface/comparison/tests/test_comparisontools.py @@ -411,7 +411,12 @@ def test_do_score_labels_counts_spikes_once(): # also check empty unit handling. unit 3 in both trains are empty, unit 4 are empty too # but are mapped to each other. sorting1, sorting2 = make_sorting( - [100, 500, 900, 1300], [0, 0, 1, 2], [102, 503, 505, 901, 1301], [0, 1, 2, 1, 2], ids1=[0, 1, 2, 3, 4], ids2=[0, 1, 2, 3, 4] + [100, 500, 900, 1300], + [0, 0, 1, 2], + [102, 503, 505, 901, 1301], + [0, 1, 2, 1, 2], + ids1=[0, 1, 2, 3, 4], + ids2=[0, 1, 2, 3, 4], ) unit_map12 = pd.Series({0: 0, 1: 1, 2: 2, 3: -1, 4: 4}) labels_st1, labels_st2 = do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclassification=True)