diff --git a/src/spikeinterface/comparison/comparisontools.py b/src/spikeinterface/comparison/comparisontools.py index 356976e47c..55e3424af9 100644 --- a/src/spikeinterface/comparison/comparisontools.py +++ b/src/spikeinterface/comparison/comparisontools.py @@ -563,6 +563,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: @@ -594,76 +663,62 @@ 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 """ - unit1_ids = sorting1.get_unit_ids() - unit2_ids = sorting2.get_unit_ids() - labels_st1 = dict() - labels_st2 = dict() + compute_pair_match_indices = get_optimized_compute_pair_match_indices() - # copy spike trains for faster access from extractors with memmapped data + # initialize labels dict of lists for each sorting with UNPAIRED 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: - 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" - else: - lab_st1 = np.array(["FN"] * len(sts)) - labels_st1[u1][seg_index] = lab_st1 + if u2 == -1 or u2 == "": + continue + st1 = sorting1.get_unit_spike_train(u1, seg_index) + 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] 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: 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 a32bef49ae..5af7e61ac7 100644 --- a/src/spikeinterface/comparison/tests/test_comparisontools.py +++ b/src/spikeinterface/comparison/tests/test_comparisontools.py @@ -1,27 +1,28 @@ 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): +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 @@ -377,6 +378,57 @@ 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"]) + + # 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. + # 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(): sorting1, sorting2 = make_sorting( [100, 200, 300, 400],