Repository: SpikeInterface/spikeinterface
Pull request #2700 (2024-04-10 19:42:24)

Seepdup random_spikes_selection() and put "all" method in it.
@yger : for you


Patch:
diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py
index 247517fc9..2aae5a85b 100644
--- a/src/spikeinterface/core/analyzer_extension_core.py
+++ b/src/spikeinterface/core/analyzer_extension_core.py
@@ -52,15 +52,11 @@ class ComputeRandomSpikes(AnalyzerExtension):
     def _run(
         self,
     ):
-        if self.params["method"] == "all":
-            all_spikes_indices = np.arange(len(self.sorting_analyzer.sorting.to_spike_vector()))
-            self.data["random_spikes_indices"] = all_spikes_indices
-        else:
-            self.data["random_spikes_indices"] = random_spikes_selection(
-                self.sorting_analyzer.sorting,
-                num_samples=self.sorting_analyzer.rec_attributes["num_samples"],
-                **self.params,
-            )
+        self.data["random_spikes_indices"] = random_spikes_selection(
+            self.sorting_analyzer.sorting,
+            num_samples=self.sorting_analyzer.rec_attributes["num_samples"],
+            **self.params,
+        )
 
     def _set_params(self, method="uniform", max_spikes_per_unit=500, margin_size=None, seed=None):
         params = dict(method=method, max_spikes_per_unit=max_spikes_per_unit, margin_size=margin_size, seed=seed)
diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py
index 8b829496c..fbecb95e8 100644
--- a/src/spikeinterface/core/sorting_tools.py
+++ b/src/spikeinterface/core/sorting_tools.py
@@ -160,7 +160,7 @@ def random_spikes_selection(
         The number of samples per segment.
         Can be retrieved from recording with
         num_samples = [recording.get_num_samples(seg_index) for seg_index in range(recording.get_num_segments())]
-    method: "uniform", default: "uniform"
+    method: "uniform"  | "all", default: "uniform"
         The method to use. Only "uniform" is implemented for now
     max_spikes_per_unit: int, default: 500
         The number of spikes per units
@@ -174,38 +174,38 @@ def random_spikes_selection(
     random_spikes_indices: np.array
         Selected spike indices coresponding to the sorting spike vector.
     """
-    rng = np.random.default_rng(seed=seed)
-    spikes = sorting.to_spike_vector()
 
-    random_spikes_indices = []
-    for unit_index, unit_id in enumerate(sorting.unit_ids):
-        all_unit_indices = np.flatnonzero(unit_index == spikes["unit_index"])
+    if method == "uniform":
+        rng = np.random.default_rng(seed=seed)
 
-        if method == "uniform":
+        spikes = sorting.to_spike_vector(concatenated=False)
+        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)
+
+        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]
+                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)
+            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
             )
-        else:
-            raise ValueError(f"random_spikes_selection wrong method {method}, currently only 'uniform' can be used.")
-
-        if margin_size is not None:
-            assert num_samples is not None
-            margin_size = int(margin_size)
-            keep = np.ones(selected_unit_indices.size, dtype=bool)
-            # left margin
-            keep[selected_unit_indices < margin_size] = False
-            # right margin
-            for segment_index in range(sorting.get_num_segments()):
-                remove_mask = np.flatnonzero(
-                    (spikes[selected_unit_indices]["segment_index"] == segment_index)
-                    & (spikes[selected_unit_indices]["sample_index"] >= (num_samples[segment_index] - margin_size))
-                )
-                keep[remove_mask] = False
-            selected_unit_indices = selected_unit_indices[keep]
+            random_spikes_indices.append(selected_unit_indices)
 
-        random_spikes_indices.append(selected_unit_indices)
+        random_spikes_indices = np.concatenate(random_spikes_indices)
+        random_spikes_indices = np.sort(random_spikes_indices)
 
-    random_spikes_indices = np.concatenate(random_spikes_indices)
-    random_spikes_indices = np.sort(random_spikes_indices)
+    elif method == "all":
+        spikes = sorting.to_spike_vector()
+        random_spikes_indices = np.arange(spikes.size)
+    else:
+        raise ValueError(f"random_spikes_selection(): method must be 'all' or 'uniform'")
 
     return random_spikes_indices


Test patch:
diff --git a/src/spikeinterface/core/tests/test_sorting_tools.py b/src/spikeinterface/core/tests/test_sorting_tools.py
index fb184d335..1aefeeb06 100644
--- a/src/spikeinterface/core/tests/test_sorting_tools.py
+++ b/src/spikeinterface/core/tests/test_sorting_tools.py
@@ -42,7 +42,7 @@ def test_spike_vector_to_indices():
 
 def test_random_spikes_selection():
     recording, sorting = generate_ground_truth_recording(
-        durations=[30.0],
+        durations=[20.0, 10.0],
         sampling_frequency=16000.0,
         num_channels=10,
         num_units=5,
@@ -56,6 +56,7 @@ def test_random_spikes_selection():
     random_spikes_indices = random_spikes_selection(
         sorting, num_samples, method="uniform", max_spikes_per_unit=max_spikes_per_unit, margin_size=None, seed=2205
     )
+    random_spikes_indices1 = random_spikes_indices
     spikes = sorting.to_spike_vector()
     some_spikes = spikes[random_spikes_indices]
     for unit_index, unit_id in enumerate(sorting.unit_ids):
@@ -69,8 +70,12 @@ def test_random_spikes_selection():
     # in that case the number is not garanty so it can be a bit less
     assert random_spikes_indices.size >= (0.9 * sorting.unit_ids.size * max_spikes_per_unit)
 
+    # all
+    random_spikes_indices = random_spikes_selection(sorting, num_samples, method="all")
+    assert random_spikes_indices.size == spikes.size
+
 
 if __name__ == "__main__":
     # test_spike_vector_to_spike_trains()
-    test_spike_vector_to_indices()
-    # test_random_spikes_selection()
+    # test_spike_vector_to_indices()
+    test_random_spikes_selection()