speedup the collision comparison and benchmarkmatching

Author: samuelgarcia

src/spikeinterface/benchmark/benchmark_base.py

@@ -374,6 +374,8 @@ def compute_results(self, case_keys=None, verbose=False, **result_params):
 
         job_keys = []
         for key in case_keys:
+            if verbose:
+                print("### Compute result", key, "###")
             benchmark = self.benchmarks[key]
             assert benchmark is not None
             benchmark.compute_result(**result_params)

src/spikeinterface/benchmark/benchmark_matching.py

@@ -91,17 +91,20 @@ def plot_performances_ordered(self, *args, **kwargs):
 
         return plot_performances_ordered(self, *args, **kwargs)
 
-    def plot_collisions(self, case_keys=None, figsize=None):
+    def plot_collisions(self, case_keys=None, axs=None, figsize=None):
         if case_keys is None:
             case_keys = list(self.cases.keys())
         import matplotlib.pyplot as plt
 
-        fig, axs = plt.subplots(ncols=len(case_keys), nrows=1, figsize=figsize, squeeze=False)
+        if axs is None:
+            fig, axs = plt.subplots(ncols=len(case_keys), nrows=1, figsize=figsize, squeeze=False)
+            axs = axs[0, :]
+
 
         for count, key in enumerate(case_keys):
             label = self.cases[key]["label"]
             templates_array = self.get_sorting_analyzer(key).get_extension("templates").get_templates(outputs="numpy")
-            ax = axs[0, count]
+            ax = axs[count]
             plot_comparison_collision_by_similarity(
                 self.get_result(key)["gt_collision"],
                 templates_array,

src/spikeinterface/comparison/collision.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import importlib
+
 from .paircomparisons import GroundTruthComparison
 
-# keep import as we do not want to delete code below.
-# from .groundtruthstudy import GroundTruthStudy
 from .comparisontools import make_collision_events
 
 import numpy as np
 
+from tqdm.auto import tqdm
+
 
 class CollisionGTComparison(GroundTruthComparison):
     """
@@ -31,98 +33,160 @@ class CollisionGTComparison(GroundTruthComparison):
 
     """
 
-    def __init__(self, gt_sorting, tested_sorting, collision_lag=2.0, nbins=11, **kwargs):
+    def __init__(self, gt_sorting, tested_sorting, collision_lag=2.0, nbins=11, progress_bar=True, **kwargs):
         # Force compute labels
         kwargs["compute_labels"] = True
 
         if gt_sorting.get_num_segments() > 1 or tested_sorting.get_num_segments() > 1:
             raise NotImplementedError("Collision comparison is only available for mono-segment sorting objects")
 
+        self.progress_bar = progress_bar
+
         GroundTruthComparison.__init__(self, gt_sorting, tested_sorting, **kwargs)
 
+        
+
         self.collision_lag = collision_lag
         self.nbins = nbins
 
         self.detect_gt_collision()
         self.compute_all_pair_collision_bins()
 
+        
+
     def detect_gt_collision(self):
         delta = int(self.collision_lag / 1000 * self.sampling_frequency)
-        self.collision_events = make_collision_events(self.sorting1, delta)
+        self.collision_events = make_collision_events(self.sorting1, delta, progress_bar=self.progress_bar)
+
+    # def get_label_for_collision(self, gt_unit_id1, gt_unit_id2):
+    #     gt_index1 = self.sorting1.id_to_index(gt_unit_id1)
+    #     gt_index2 = self.sorting1.id_to_index(gt_unit_id2)
+    #     if gt_index1 > gt_index2:
+    #         gt_unit_id1, gt_unit_id2 = gt_unit_id2, gt_unit_id1
+    #         reversed = True
+    #     else:
+    #         reversed = False
+
+    #     # events
+    #     mask = (self.collision_events["unit_id1"] == gt_unit_id1) & (self.collision_events["unit_id2"] == gt_unit_id2)
+    #     event = self.collision_events[mask]
+
+    #     score_label1 = self._labels_st1[gt_unit_id1][0][event["index1"]]
+    #     score_label2 = self._labels_st1[gt_unit_id2][0][event["index2"]]
+    #     delta = event["delta_frame"]
+
+    #     if reversed:
+    #         score_label1, score_label2 = score_label2, score_label1
+    #         delta = -delta
+
+    #     return score_label1, score_label2, delta
+
+    # def get_label_count_per_collision_bins(self, gt_unit_id1, gt_unit_id2, bins):
+    #     score_label1, score_label2, delta = self.get_label_for_collision(gt_unit_id1, gt_unit_id2)
 
-    def get_label_for_collision(self, gt_unit_id1, gt_unit_id2):
-        gt_index1 = self.sorting1.id_to_index(gt_unit_id1)
-        gt_index2 = self.sorting1.id_to_index(gt_unit_id2)
-        if gt_index1 > gt_index2:
-            gt_unit_id1, gt_unit_id2 = gt_unit_id2, gt_unit_id1
-            reversed = True
-        else:
-            reversed = False
+    #     tp_count1 = np.zeros(bins.size - 1)
+    #     fn_count1 = np.zeros(bins.size - 1)
+    #     tp_count2 = np.zeros(bins.size - 1)
+    #     fn_count2 = np.zeros(bins.size - 1)
 
-        # events
-        mask = (self.collision_events["unit_id1"] == gt_unit_id1) & (self.collision_events["unit_id2"] == gt_unit_id2)
-        event = self.collision_events[mask]
+    #     for i in range(tp_count1.size):
+    #         l0, l1 = bins[i], bins[i + 1]
+    #         mask = (delta >= l0) & (delta < l1)
 
-        score_label1 = self._labels_st1[gt_unit_id1][0][event["index1"]]
-        score_label2 = self._labels_st1[gt_unit_id2][0][event["index2"]]
-        delta = event["delta_frame"]
+    #         tp_count1[i] = np.sum(score_label1[mask] == "TP")
+    #         fn_count1[i] = np.sum(score_label1[mask] == "FN")
+    #         tp_count2[i] = np.sum(score_label2[mask] == "TP")
+    #         fn_count2[i] = np.sum(score_label2[mask] == "FN")
 
-        if reversed:
-            score_label1, score_label2 = score_label2, score_label1
-            delta = -delta
+    #     # inverse for unit_id2
+    #     tp_count2 = tp_count2[::-1]
+    #     fn_count2 = fn_count2[::-1]
 
-        return score_label1, score_label2, delta
+    #     return tp_count1, fn_count1, tp_count2, fn_count2
 
-    def get_label_count_per_collision_bins(self, gt_unit_id1, gt_unit_id2, bins):
-        score_label1, score_label2, delta = self.get_label_for_collision(gt_unit_id1, gt_unit_id2)
+    # def compute_all_pair_collision_bins(self):
+    #     print('CollisionGTComparison.compute_all_pair_collision_bins')
+    #     d = int(self.collision_lag / 1000 * self.sampling_frequency)
+    #     bins = np.linspace(-d, d, self.nbins + 1)
+    #     self.bins = bins
 
-        tp_count1 = np.zeros(bins.size - 1)
-        fn_count1 = np.zeros(bins.size - 1)
-        tp_count2 = np.zeros(bins.size - 1)
-        fn_count2 = np.zeros(bins.size - 1)
+    #     unit_ids = self.sorting1.unit_ids
+    #     n = len(unit_ids)
 
-        for i in range(tp_count1.size):
-            l0, l1 = bins[i], bins[i + 1]
-            mask = (delta >= l0) & (delta < l1)
+    #     all_tp_count1 = []
+    #     all_fn_count1 = []
+    #     all_tp_count2 = []
+    #     all_fn_count2 = []
 
-            tp_count1[i] = np.sum(score_label1[mask] == "TP")
-            fn_count1[i] = np.sum(score_label1[mask] == "FN")
-            tp_count2[i] = np.sum(score_label2[mask] == "TP")
-            fn_count2[i] = np.sum(score_label2[mask] == "FN")
+    #     self.all_tp = np.zeros((n, n, self.nbins), dtype="int64")
+    #     self.all_fn = np.zeros((n, n, self.nbins), dtype="int64")
 
-        # inverse for unit_id2
-        tp_count2 = tp_count2[::-1]
-        fn_count2 = fn_count2[::-1]
+    #     for i in range(n):
+    #         print(i, n)
+    #         for j in range(i + 1, n):
+    #             u1 = unit_ids[i]
+    #             u2 = unit_ids[j]
 
-        return tp_count1, fn_count1, tp_count2, fn_count2
+    #             tp_count1, fn_count1, tp_count2, fn_count2 = self.get_label_count_per_collision_bins(u1, u2, bins)
+
+    #             self.all_tp[i, j, :] = tp_count1
+    #             self.all_tp[j, i, :] = tp_count2
+    #             self.all_fn[i, j, :] = fn_count1
+    #             self.all_fn[j, i, :] = fn_count2
 
     def compute_all_pair_collision_bins(self):
         d = int(self.collision_lag / 1000 * self.sampling_frequency)
         bins = np.linspace(-d, d, self.nbins + 1)
         self.bins = bins
 
-        unit_ids = self.sorting1.unit_ids
-        n = len(unit_ids)
+        collision_events = self.collision_events
+        labels_st1 = self._labels_st1
+        gt_unit_ids = self.sorting1.unit_ids
+
+        nbins = bins.size -1 
+        n = len(gt_unit_ids)
+        all_tp = np.zeros((n, n, nbins), dtype="int64")
+        all_fn = np.zeros((n, n, nbins), dtype="int64")
+
+        unit_ids1 = collision_events['unit_id1']
+        unit_indices1 = collision_events['unit_index1']
+        unit_ids2 = collision_events['unit_id2']
+        unit_indices2 = collision_events['unit_index2']
+
+        spike_indices1 = collision_events['index1']
+        spike_indices2 = collision_events['index2']
+        delta_frame = collision_events['delta_frame']
+        delta_frame
+        delta_bin = np.clip(np.floor((delta_frame - bins[0]) / (bins[1] - bins[0])), 0, nbins-1).astype('int64')
+        inv_delta_bin = np.clip(np.floor((-delta_frame - bins[0]) / (bins[1] - bins[0])), 0, nbins-1).astype('int64')
 
-        all_tp_count1 = []
-        all_fn_count1 = []
-        all_tp_count2 = []
-        all_fn_count2 = []
+        seg_index = 0
 
-        self.all_tp = np.zeros((n, n, self.nbins), dtype="int64")
-        self.all_fn = np.zeros((n, n, self.nbins), dtype="int64")
+        loop = range(len(unit_ids1))
+        if self.progress_bar:
+            loop = tqdm(loop, desc="collision by bin")
+        
+        for c in loop:
 
-        for i in range(n):
-            for j in range(i + 1, n):
-                u1 = unit_ids[i]
-                u2 = unit_ids[j]
+            score1 = labels_st1[unit_ids1[c]][seg_index][spike_indices1[c]]
+            score2 = labels_st1[unit_ids2[c]][seg_index][spike_indices2[c]]
 
-                tp_count1, fn_count1, tp_count2, fn_count2 = self.get_label_count_per_collision_bins(u1, u2, bins)
+            unit_index1 = unit_indices1[c]
+            unit_index2 = unit_indices2[c]
+
+            if score1 == "TP":
+                all_tp[unit_index1, unit_index2, delta_bin[c]] += 1
+            else:
+                all_fn[unit_index1, unit_index2, delta_bin[c]] += 1
+
+            if score2 == "TP":
+                all_tp[unit_index2, unit_index1, inv_delta_bin[c]] += 1
+            else:
+                all_fn[unit_index2, unit_index1, inv_delta_bin[c]] += 1
+
+        self.all_tp = all_tp
+        self.all_fn = all_fn
 
-                self.all_tp[i, j, :] = tp_count1
-                self.all_tp[j, i, :] = tp_count2
-                self.all_fn[i, j, :] = fn_count1
-                self.all_fn[j, i, :] = fn_count2
 
     def compute_collision_by_similarity(self, similarity_matrix, unit_ids=None, good_only=False, min_accuracy=0.9):
         if unit_ids is None:
@@ -172,77 +236,3 @@ def compute_collision_by_similarity(self, similarity_matrix, unit_ids=None, good
         pair_names = pair_names[order]
 
         return similarities, recall_scores, pair_names
-
-
-# This is removed at the moment.
-# We need to move this maybe one day in benchmark.
-# please do not delete this
-
-# class CollisionGTStudy(GroundTruthStudy):
-#     def run_comparisons(self, case_keys=None, exhaustive_gt=True, collision_lag=2.0, nbins=11, **kwargs):
-#         _kwargs = dict()
-#         _kwargs.update(kwargs)
-#         _kwargs["exhaustive_gt"] = exhaustive_gt
-#         _kwargs["collision_lag"] = collision_lag
-#         _kwargs["nbins"] = nbins
-#         GroundTruthStudy.run_comparisons(self, case_keys=None, comparison_class=CollisionGTComparison, **_kwargs)
-#         self.exhaustive_gt = exhaustive_gt
-#         self.collision_lag = collision_lag
-
-#     def get_lags(self, key):
-#         comp = self.comparisons[key]
-#         fs = comp.sorting1.get_sampling_frequency()
-#         lags = comp.bins / fs * 1000.0
-#         return lags
-
-#     def precompute_scores_by_similarities(self, case_keys=None, good_only=False, min_accuracy=0.9):
-#         import sklearn
-
-#         if case_keys is None:
-#             case_keys = self.cases.keys()
-
-#         self.all_similarities = {}
-#         self.all_recall_scores = {}
-#         self.good_only = good_only
-
-#         for key in case_keys:
-#             templates = self.get_templates(key)
-#             flat_templates = templates.reshape(templates.shape[0], -1)
-#             similarity = sklearn.metrics.pairwise.cosine_similarity(flat_templates)
-#             comp = self.comparisons[key]
-#             similarities, recall_scores, pair_names = comp.compute_collision_by_similarity(
-#                 similarity, good_only=good_only, min_accuracy=min_accuracy
-#             )
-#             self.all_similarities[key] = similarities
-#             self.all_recall_scores[key] = recall_scores
-
-#     def get_mean_over_similarity_range(self, similarity_range, key):
-#         idx = (self.all_similarities[key] >= similarity_range[0]) & (self.all_similarities[key] <= similarity_range[1])
-#         all_similarities = self.all_similarities[key][idx]
-#         all_recall_scores = self.all_recall_scores[key][idx]
-
-#         order = np.argsort(all_similarities)
-#         all_similarities = all_similarities[order]
-#         all_recall_scores = all_recall_scores[order, :]
-
-#         mean_recall_scores = np.nanmean(all_recall_scores, axis=0)
-
-#         return mean_recall_scores
-
-#     def get_lag_profile_over_similarity_bins(self, similarity_bins, key):
-#         all_similarities = self.all_similarities[key]
-#         all_recall_scores = self.all_recall_scores[key]
-
-#         order = np.argsort(all_similarities)
-#         all_similarities = all_similarities[order]
-#         all_recall_scores = all_recall_scores[order, :]
-
-#         result = {}
-
-#         for i in range(similarity_bins.size - 1):
-#             cmin, cmax = similarity_bins[i], similarity_bins[i + 1]
-#             amin, amax = np.searchsorted(all_similarities, [cmin, cmax])
-#             mean_recall_scores = np.nanmean(all_recall_scores[amin:amax], axis=0)
-#             result[(cmin, cmax)] = mean_recall_scores
-
-#         return result