oups : spikes in margin

Author: samuelgarcia

src/spikeinterface/sortingcomponents/matching/base.py

@@ -27,7 +27,18 @@ def get_dtype(self):
     def get_trace_margin(self):
         raise NotImplementedError  
 
-    def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
+    def compute(self, traces, start_frame, end_frame, segment_index, max_margin):
+        spikes = self.compute_matching(traces, start_frame, end_frame, segment_index)
+        spikes["segment_index"] = segment_index
+
+        margin = self.get_trace_margin()
+        if margin > 0:
+            keep = (spikes["sample_index"] >= margin) & (spikes["sample_index"] < (traces.shape[0] - margin))
+            spikes = spikes[keep]
+
+        return spikes
+
+    def compute_matching(self, traces, start_frame, end_frame, segment_index):
         raise NotImplementedError
 
     def get_extra_outputs(self):

src/spikeinterface/sortingcomponents/matching/circus.py

@@ -274,7 +274,7 @@ def get_extra_outputs(self):
     def get_trace_margin(self):
         return self.margin    
 
-    def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
+    def compute_matching(self, traces, start_frame, end_frame, segment_index):
         import scipy.spatial
         import scipy
 
@@ -478,8 +478,6 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *ar
         order = np.argsort(spikes["sample_index"])
         spikes = spikes[order]
 
-        spikes["segment_index"] = segment_index
-
         return spikes
 
 
@@ -1024,7 +1022,7 @@ def get_trace_margin(self):
         return self.margin
 
 
-    def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
+    def compute_matching(self, traces, start_frame, end_frame, segment_index):
 
         neighbor_window = self.num_samples - 1
 
@@ -1107,8 +1105,6 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *ar
         order = np.argsort(spikes["sample_index"])
         spikes = spikes[order]
 
-        spikes["segment_index"] = segment_index
-        
         return spikes
 
 

src/spikeinterface/sortingcomponents/matching/naive.py

@@ -51,7 +51,7 @@ def get_trace_margin(self):
         return self.margin
 
 
-    def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
+    def compute_matching(self, traces, start_frame, end_frame, segment_index):
 
         if self.margin > 0:
             peak_traces = traces[self.margin:-self.margin, :]
@@ -78,8 +78,6 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *ar
             spikes["cluster_index"][i] = cluster_index
             spikes["amplitude"][i] = 0.0
 
-        spikes["segment_index"] = segment_index
-        
         return spikes
 
 

src/spikeinterface/sortingcomponents/matching/tdc.py

@@ -160,7 +160,7 @@ def __init__(self, recording, return_output=True, parents=None,
     def get_trace_margin(self):
         return self.margin
 
-    def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
+    def compute_matching(self, traces, start_frame, end_frame, segment_index):
         traces = traces.copy()
 
         all_spikes = []
@@ -186,8 +186,6 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *ar
         else:
             all_spikes = np.zeros(0, dtype=_base_matching_dtype)
 
-        all_spikes["segment_index"] = segment_index
-
         return all_spikes
 
     def _find_spikes_one_level(self, traces, level=0):

src/spikeinterface/sortingcomponents/matching/wobble.py

@@ -383,7 +383,7 @@ def __init__(self, recording, return_output=True, parents=None,
     def get_trace_margin(self):
         return self.margin
 
-    def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
+    def compute_matching(self, traces, start_frame, end_frame, segment_index):
 
         # Unpack method_kwargs
         # nbefore, nafter = method_kwargs["nbefore"], method_kwargs["nafter"]
@@ -450,8 +450,7 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *ar
         spikes["cluster_index"] = spike_train[:, 1]
         spikes["channel_index"] = channel_inds
         spikes["amplitude"] = amplitudes
-        spikes["segment_index"] = segment_index
-        
+
         return spikes
 
     # TODO: Replace this method with equivalent from spikeinterface

src/spikeinterface/sortingcomponents/tests/test_template_matching.py

@@ -77,7 +77,7 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
     #     comp = si.compare_sorter_to_ground_truth(gt_sorting, sorting)
     #     si.plot_agreement_matrix(comp, ax=ax)
     #     ax.set_title(method)
-    #     plt.show()
+        plt.show()
 
 
 if __name__ == "__main__":