Merge pull request #3456 from samuelgarcia/fix_select_spikes

fix random_spikes_selection()
SpikeInterface · Oct 2, 2024 · 10e90db · 10e90db
2 parents dc1496f + 036691b
commit 10e90db
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py
@@ -197,17 +197,23 @@ def random_spikes_selection(
         cum_sizes = np.cumsum([0] + [s.size for s in spikes])
 
         # this fast when numba
-        spike_indices = spike_vector_to_indices(spikes, sorting.unit_ids)
+        spike_indices = spike_vector_to_indices(spikes, sorting.unit_ids, absolute_index=False)
 
         random_spikes_indices = []
         for unit_index, unit_id in enumerate(sorting.unit_ids):
             all_unit_indices = []
             for segment_index in range(sorting.get_num_segments()):
-                inds_in_seg = spike_indices[segment_index][unit_id] + cum_sizes[segment_index]
+                # this is local index
+                inds_in_seg = spike_indices[segment_index][unit_id]
                 if margin_size is not None:
-                    inds_in_seg = inds_in_seg[inds_in_seg >= margin_size]
-                    inds_in_seg = inds_in_seg[inds_in_seg < (num_samples[segment_index] - margin_size)]
-                all_unit_indices.append(inds_in_seg)
+                    local_spikes = spikes[segment_index][inds_in_seg]
+                    mask = (local_spikes["sample_index"] >= margin_size) & (
+                        local_spikes["sample_index"] < (num_samples[segment_index] - margin_size)
+                    )
+                    inds_in_seg = inds_in_seg[mask]
+                # go back to absolut index
+                inds_in_seg_abs = inds_in_seg + cum_sizes[segment_index]
+                all_unit_indices.append(inds_in_seg_abs)
             all_unit_indices = np.concatenate(all_unit_indices)
             selected_unit_indices = rng.choice(
                 all_unit_indices, size=min(max_spikes_per_unit, all_unit_indices.size), replace=False, shuffle=False

diff --git a/src/spikeinterface/core/tests/test_sorting_tools.py b/src/spikeinterface/core/tests/test_sorting_tools.py
@@ -162,8 +162,8 @@ def test_generate_unit_ids_for_merge_group():
 if __name__ == "__main__":
     # test_spike_vector_to_spike_trains()
     # test_spike_vector_to_indices()
-    # test_random_spikes_selection()
+    test_random_spikes_selection()
 
-    test_apply_merges_to_sorting()
-    test_get_ids_after_merging()
+    # test_apply_merges_to_sorting()
+    # test_get_ids_after_merging()
     # test_generate_unit_ids_for_merge_group()