Extensions are for all spikes

SpikeInterface · Jul 1, 2024 · 3211df7 · 3211df7
1 parent a3cc09b
commit 3211df7
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py
@@ -123,10 +123,9 @@ def _merge_extension_data(
             new_data["collision_mask"] = self.data["collision_mask"]
 
         if kept_indices is not None:
-            valid = kept_indices[self.sorting_analyzer.get_extension("random_spikes")._get_data()]
-            new_data["amplitude_scalings"] = new_data["amplitude_scalings"][valid]
+            new_data["amplitude_scalings"] = new_data["amplitude_scalings"][kept_indices]
             if self.params["handle_collisions"]:
-                new_data["collision_mask"] = new_data["collision_mask"][valid]
+                new_data["collision_mask"] = new_data["collision_mask"][kept_indices]
 
         return new_data
 

diff --git a/src/spikeinterface/postprocessing/spike_amplitudes.py b/src/spikeinterface/postprocessing/spike_amplitudes.py
@@ -88,8 +88,7 @@ def _merge_extension_data(
         new_data["amplitudes"] = self.data["amplitudes"]
 
         if kept_indices is not None:
-            valid = kept_indices[self.sorting_analyzer.get_extension("random_spikes")._get_data()]
-            new_data["amplitudes"] = new_data["amplitudes"][valid]
+            new_data["amplitudes"] = new_data["amplitudes"][kept_indices]
 
         return new_data
 

diff --git a/src/spikeinterface/postprocessing/spike_locations.py b/src/spikeinterface/postprocessing/spike_locations.py
@@ -97,8 +97,7 @@ def _merge_extension_data(
     ):
         new_spike_locations = self.data["spike_locations"]
         if kept_indices is not None:
-            valid = kept_indices[self.sorting_analyzer.get_extension("random_spikes")._get_data()]
-            new_spike_locations = new_spike_locations[valid]
+            new_spike_locations = new_spike_locations[kept_indices]
 
         ### In theory here, we should recompute the locations since the peak positions
         ### in a merged could be different. Should be discussed