Expose agreement_method ark in basepairsortercomparison

Author: alejoe91

src/spikeinterface/comparison/collision.py

@@ -15,11 +15,11 @@ class CollisionGTComparison(GroundTruthComparison):
 
     Parameters
     ----------
-    gt_sorting : SortingExtractor
+    gt_sorting : BaseSorting
         The first sorting for the comparison
     collision_lag : float, default 2.0
         Collision lag in ms.
-    tested_sorting : SortingExtractor
+    tested_sorting : BaseSorting
         The second sorting for the comparison
     nbins : int, default : 11
         Number of collision bins

src/spikeinterface/comparison/comparisontools.py

@@ -72,7 +72,7 @@ def do_count_event(sorting):
 
     Parameters
     ----------
-    sorting : SortingExtractor
+    sorting : BaseSorting
         A sorting extractor
 
     Returns
@@ -345,9 +345,9 @@ def make_agreement_scores(
 
     Parameters
     ----------
-    sorting1 : SortingExtractor
+    sorting1 : BaseSorting
         The first sorting extractor
-    sorting2 : SortingExtractor
+    sorting2 : BaseSorting
         The second sorting extractor
     delta_frames : int
         Number of frames to consider spikes coincident
@@ -549,9 +549,9 @@ def do_score_labels(sorting1, sorting2, delta_frames, unit_map12, label_misclass
 
     Parameters
     ----------
-    sorting1 : SortingExtractor instance
+    sorting1 : BaseSorting
         The ground truth sorting
-    sorting2 : SortingExtractor instance
+    sorting2 : BaseSorting
         The tested sorting
     delta_frames : int
         Number of frames to consider spikes coincident
@@ -902,7 +902,7 @@ def make_collision_events(sorting, delta):
 
     Parameters
     ----------
-    sorting : SortingExtractor
+    sorting : BaseSorting
         The sorting extractor object for counting collision events
     delta : int
         Number of frames for considering collision events
@@ -1297,9 +1297,10 @@ def compute_distance_matrix(sorting1: BaseSorting, sorting2: BaseSorting, delta_
         squared_norm_1_vector[:, np.newaxis] + squared_norm_2_vector[np.newaxis, :] - 2 * dot_product_matrix
     )
 
-    distance_metrix = np.sqrt(squared_distance_matrix)
+    distance_matrix = np.sqrt(squared_distance_matrix)
+    agreement_matrix = 1 / ((distance_matrix**2 / dot_product_matrix) + 1)
 
-    return distance_metrix
+    return distance_matrix, agreement_matrix
 
 
 def calculate_generalized_comparison_metrics(

src/spikeinterface/comparison/correlogram.py

@@ -17,9 +17,9 @@ class CorrelogramGTComparison(GroundTruthComparison):
 
     Parameters
     ----------
-    gt_sorting : SortingExtractor
+    gt_sorting : BaseSorting
         The first sorting for the comparison
-    tested_sorting : SortingExtractor
+    tested_sorting : BaseSorting
         The second sorting for the comparison
     bin_ms : float, default: 1.0
         Size of bin for correlograms

src/spikeinterface/comparison/paircomparisons.py

@@ -2,7 +2,8 @@
 
 import numpy as np
 
-from spikeinterface.core.core_tools import define_function_from_class
+from ..core import BaseSorting
+from ..core.core_tools import define_function_from_class
 from .basecomparison import BasePairComparison, MixinSpikeTrainComparison, MixinTemplateComparison
 from .comparisontools import (
     do_count_event,
@@ -12,6 +13,7 @@
     do_confusion_matrix,
     do_count_score,
     compute_performance,
+    compute_distance_matrix,
 )
 from ..postprocessing import compute_template_similarity_by_pair
 
@@ -23,16 +25,17 @@ class BasePairSorterComparison(BasePairComparison, MixinSpikeTrainComparison):
 
     def __init__(
         self,
-        sorting1,
-        sorting2,
-        sorting1_name=None,
-        sorting2_name=None,
-        delta_time=0.4,
-        match_score=0.5,
-        chance_score=0.1,
-        ensure_symmetry=False,
-        n_jobs=1,
-        verbose=False,
+        sorting1: BaseSorting,
+        sorting2: BaseSorting,
+        sorting1_name: str | None = None,
+        sorting2_name: str | None = None,
+        delta_time: float = 0.4,
+        match_score: float = 0.5,
+        chance_score: float = 0.1,
+        ensure_symmetry: bool = False,
+        agreement_method: str = "from_count",
+        n_jobs: int = 1,
+        verbose: bool = False,
     ):
         if sorting1_name is None:
             sorting1_name = "sorting1"
@@ -59,6 +62,7 @@ def __init__(
         self.unit2_ids = self.sorting2.get_unit_ids()
 
         self.ensure_symmetry = ensure_symmetry
+        self.agreement_method = agreement_method
 
         self._do_agreement()
         self._do_matching()
@@ -85,18 +89,30 @@ def _do_agreement(self):
 
         # common to GroundTruthComparison and SymmetricSortingComparison
         # spike count for each spike train
-        self.event_counts1 = do_count_event(self.sorting1)
-        self.event_counts2 = do_count_event(self.sorting2)
+        if self.agreement_method == "from_count":
+            self.event_counts1 = do_count_event(self.sorting1)
+            self.event_counts2 = do_count_event(self.sorting2)
 
-        # matrix of  event match count for each pair
-        self.match_event_count = make_match_count_matrix(
-            self.sorting1, self.sorting2, self.delta_frames, ensure_symmetry=self.ensure_symmetry
-        )
+            # matrix of  event match count for each pair
+            self.match_event_count = make_match_count_matrix(
+                self.sorting1, self.sorting2, self.delta_frames, ensure_symmetry=self.ensure_symmetry
+            )
 
-        # agreement matrix score for each pair
-        self.agreement_scores = make_agreement_scores_from_count(
-            self.match_event_count, self.event_counts1, self.event_counts2
-        )
+            # agreement matrix score for each pair
+            self.agreement_scores = make_agreement_scores_from_count(
+                self.match_event_count, self.event_counts1, self.event_counts2
+            )
+        elif self.agreement_method == "distance_function":
+            import pandas as pd
+
+            _, agreement_matrix = compute_distance_matrix(self.sorting1, self.sorting2, self.delta_frames)
+
+            self.agreement_scores = pd.DataFrame(
+                agreement_matrix, index=self.sorting1.unit_ids, columns=self.sorting2.unit_ids
+            )
+
+        else:
+            raise ValueError("agreement_method must be 'from_count' or 'distance_matrix'")
 
 
 class SymmetricSortingComparison(BasePairSorterComparison):
@@ -112,9 +128,9 @@ class SymmetricSortingComparison(BasePairSorterComparison):
 
     Parameters
     ----------
-    sorting1 : SortingExtractor
+    sorting1 : BaseSorting
         The first sorting for the comparison
-    sorting2 : SortingExtractor
+    sorting2 : BaseSorting
         The second sorting for the comparison
     sorting1_name : str, default: None
         The name of sorter 1
@@ -126,6 +142,9 @@ class SymmetricSortingComparison(BasePairSorterComparison):
         Minimum agreement score to match units
     chance_score : float, default: 0.1
         Minimum agreement score to for a possible match
+    agreement_method : "from_count" | "distance_function", default: "from_count"
+        The method to compute agreement scores. The "from_count" method computes agreement scores from spike counts.
+        The "distance_function" method computes agreement scores from spike time distance functions.
     n_jobs : int, default: -1
         Number of cores to use in parallel. Uses all available if -1
     verbose : bool, default: False
@@ -139,15 +158,16 @@ class SymmetricSortingComparison(BasePairSorterComparison):
 
     def __init__(
         self,
-        sorting1,
-        sorting2,
-        sorting1_name=None,
-        sorting2_name=None,
-        delta_time=0.4,
-        match_score=0.5,
-        chance_score=0.1,
-        n_jobs=-1,
-        verbose=False,
+        sorting1: BaseSorting,
+        sorting2: BaseSorting,
+        sorting1_name: str | None = None,
+        sorting2_name: str | None = None,
+        delta_time: float = 0.4,
+        match_score: float = 0.5,
+        chance_score: float = 0.1,
+        agreement_method: str = "from_count",
+        n_jobs: int = -1,
+        verbose: bool = False,
     ):
         BasePairSorterComparison.__init__(
             self,
@@ -159,6 +179,7 @@ def __init__(
             match_score=match_score,
             chance_score=chance_score,
             ensure_symmetry=True,
+            agreement_method=agreement_method,
             n_jobs=n_jobs,
             verbose=verbose,
         )
@@ -167,10 +188,13 @@ def get_matching(self):
         return self.hungarian_match_12, self.hungarian_match_21
 
     def get_matching_event_count(self, unit1, unit2):
-        if (unit1 is not None) and (unit2 is not None):
-            return self.match_event_count.at[unit1, unit2]
+        if self.agreement_method == "from_count":
+            if (unit1 is not None) and (unit2 is not None):
+                return self.match_event_count.at[unit1, unit2]
+            else:
+                raise Exception("get_matching_event_count: unit1 and unit2 must not be None.")
         else:
-            raise Exception("get_matching_event_count: unit1 and unit2 must not be None.")
+            raise Exception("get_matching_event_count is valid only if agreement_method='from_count'")
 
     def get_best_unit_match1(self, unit1):
         return self.best_match_12[unit1]
@@ -215,9 +239,9 @@ class GroundTruthComparison(BasePairSorterComparison):
 
     Parameters
     ----------
-    gt_sorting : SortingExtractor
+    gt_sorting : BaseSorting
         The first sorting for the comparison
-    tested_sorting : SortingExtractor
+    tested_sorting : BaseSorting
         The second sorting for the comparison
     gt_name : str, default: None
         The name of sorter 1
@@ -243,6 +267,9 @@ class GroundTruthComparison(BasePairSorterComparison):
         For instance, MEArec simulated dataset have exhaustive_gt=True
     match_mode : "hungarian" | "best", default: "hungarian"
         The method to match units
+    agreement_method : "from_count" | "distance_function", default: "from_count"
+        The method to compute agreement scores. The "from_count" method computes agreement scores from spike counts.
+        The "distance_function" method computes agreement scores from spike time distance functions.
     n_jobs : int, default: -1
         Number of cores to use in parallel. Uses all available if -1
     compute_labels : bool, default: False
@@ -260,22 +287,23 @@ class GroundTruthComparison(BasePairSorterComparison):
 
     def __init__(
         self,
-        gt_sorting,
-        tested_sorting,
-        gt_name=None,
-        tested_name=None,
-        delta_time=0.4,
-        match_score=0.5,
-        well_detected_score=0.8,
-        redundant_score=0.2,
-        overmerged_score=0.2,
-        chance_score=0.1,
-        exhaustive_gt=False,
-        n_jobs=-1,
-        match_mode="hungarian",
-        compute_labels=False,
-        compute_misclassifications=False,
-        verbose=False,
+        gt_sorting: BaseSorting,
+        tested_sorting: BaseSorting,
+        gt_name: str | None = None,
+        tested_name: str | None = None,
+        delta_time: float = 0.4,
+        match_score: float = 0.5,
+        well_detected_score: float = 0.8,
+        redundant_score: float = 0.2,
+        overmerged_score: float = 0.2,
+        chance_score: float = 0.1,
+        exhaustive_gt: bool = False,
+        agreement_method: str = "from_count",
+        n_jobs: int = -1,
+        match_mode: str = "hungarian",
+        compute_labels: bool = False,
+        compute_misclassifications: bool = False,
+        verbose: bool = False,
     ):
         import pandas as pd
 
@@ -293,6 +321,7 @@ def __init__(
             match_score=match_score,
             chance_score=chance_score,
             ensure_symmetry=False,
+            agreement_method=agreement_method,
             n_jobs=n_jobs,
             verbose=verbose,
         )

src/spikeinterface/comparison/tests/test_symmetricsortingcomparison.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from spikeinterface.core import generate_sorting
+from spikeinterface.core import generate_sorting, aggregate_units
 from spikeinterface.extractors import NumpySorting
 from spikeinterface.comparison import compare_two_sorters
 
@@ -24,9 +24,13 @@ def test_compare_two_sorters():
         ],
         [0, 0, 5],
     )
-    sc = compare_two_sorters(sorting1, sorting2)
+    sc_from_counts = compare_two_sorters(sorting1, sorting2, agreement_method="from_count")
+    sc_from_distance = compare_two_sorters(sorting1, sorting2, agreement_method="distance_function")
 
-    print(sc.agreement_scores)
+    np.testing.assert_array_equal(
+        sc_from_counts.hungarian_match_12.to_numpy(),
+        sc_from_distance.hungarian_match_12.to_numpy(),
+    )
 
 
 def test_compare_multi_segment():
@@ -37,6 +41,24 @@ def test_compare_multi_segment():
     assert np.allclose(np.diag(cmp_multi.agreement_scores), np.ones(len(sort.unit_ids)))
 
 
+def test_agreements():
+    """
+    Test that the agreement scores are the same when using from_count and distance_matrix
+    """
+    sorting1 = generate_sorting(num_units=100)
+    sorting_extra = generate_sorting(num_units=50)
+    sorting2 = aggregate_units([sorting1, sorting_extra])
+    sorting2 = sorting2.select_units(unit_ids=sorting2.unit_ids[np.random.permutation(len(sorting2.unit_ids))])
+
+    sc_from_counts = compare_two_sorters(sorting1, sorting2, agreement_method="from_count")
+    sc_from_distance = compare_two_sorters(sorting1, sorting2, agreement_method="distance_function")
+
+    np.testing.assert_array_equal(
+        sc_from_counts.hungarian_match_12.to_numpy(),
+        sc_from_distance.hungarian_match_12.to_numpy(),
+    )
+
+
 if __name__ == "__main__":
     test_compare_two_sorters()
     test_compare_multi_segment()