Update score labels to avoid bad assert

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="<U8") for sts in sts1[u1]]
-        labels_st1[u1] = lab_st1
-    for u2 in unit2_ids:
-        lab_st2 = [np.array(["UNPAIRED"] * len(sts), dtype="<U8") for sts in sts2[u2]]
-        labels_st2[u2] = lab_st2
+    unit1_ids = sorting1.get_unit_ids()
+    unit2_ids = sorting2.get_unit_ids()
+    labels_st1 = {u1: [] for u1 in unit1_ids}
+    labels_st2 = {u2: [] for u2 in unit2_ids}
+    for seg_index in range(num_segments):
+        for u1 in unit1_ids:
+            ns1 = len(sorting1.get_unit_spike_train(u1, seg_index))
+            labels_st1[u1].append(np.full((ns1,), "UNPAIRED", dtype="<U8"))
+        for u2 in unit2_ids:
+            ns2 = len(sorting2.get_unit_spike_train(u2, seg_index))
+            labels_st2[u2].append(np.full((ns2,), "UNPAIRED", dtype="<U8"))
 
     for seg_index in range(num_segments):
         for u1 in unit1_ids:
             u2 = unit_map12[u1]
-            sts = sts1[u1][seg_index]
-            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"
-            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"

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],