matching handle segment_index

Author: samuelgarcia

src/spikeinterface/sortingcomponents/matching/circus.py

@@ -477,6 +477,8 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *ar
         spikes = spikes[:num_spikes]
         order = np.argsort(spikes["sample_index"])
         spikes = spikes[order]
+        
+        spikes["segment_index"] = segment_index
 
         return spikes
 
@@ -1105,6 +1107,8 @@ 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

@@ -78,6 +78,8 @@ 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

@@ -186,6 +186,8 @@ 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):
@@ -307,7 +309,7 @@ def _find_spikes_one_level(self, traces, level=0):
 
             spikes["cluster_index"][i] = cluster_index
             spikes["amplitude"][i] = amplitude
-
+        
         return spikes
 
 

src/spikeinterface/sortingcomponents/matching/wobble.py

@@ -450,7 +450,8 @@ 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