Repository: SpikeInterface/spikeinterface
Pull request #540 (2022-03-04 16:50:43)

Add function to remove empty units


Patch:
diff --git a/spikeinterface/core/base.py b/spikeinterface/core/base.py
index 61b365e2d..8a3c25cab 100644
--- a/spikeinterface/core/base.py
+++ b/spikeinterface/core/base.py
@@ -62,7 +62,7 @@ class BaseExtractor:
             if self.get_num_segments() == 1:
                 return 0
             else:
-                raise ValueError()
+                raise ValueError("Multi-segment object. Provide 'segment_index'")
         else:
             return segment_index
 
diff --git a/spikeinterface/core/basesorting.py b/spikeinterface/core/basesorting.py
index 0963ae0e4..edb14bf38 100644
--- a/spikeinterface/core/basesorting.py
+++ b/spikeinterface/core/basesorting.py
@@ -129,12 +129,43 @@ class BaseSorting(BaseExtractor):
         return v
 
     def select_units(self, unit_ids, renamed_unit_ids=None):
+        """
+        Selects a subset of units
+
+        Parameters
+        ----------
+        unit_ids : numpy.array or list
+            List of unit ids to keep
+        renamed_unit_ids : numpy.array or list, optional
+            If given, the kept unit ids are renamed, by default None
+
+        Returns
+        -------
+        BaseSorting
+            Sorting object with selected units
+        """
         from spikeinterface import UnitsSelectionSorting
         sub_sorting = UnitsSelectionSorting(
             self, unit_ids, renamed_unit_ids=renamed_unit_ids)
         return sub_sorting
 
-    # add tests
+    def remove_empty_units(self):
+        """
+        Removes units with empty spike trains
+
+        Returns
+        -------
+        BaseSorting
+            Sorting object with non-empty units
+        """
+        units_to_keep = []
+        for segment_index in range(self.get_num_segments()):
+            for unit in self.get_unit_ids():
+                if len(self.get_unit_spike_train(unit, segment_index=segment_index)) > 0:
+                    units_to_keep.append(unit)
+        units_to_keep = np.unique(units_to_keep)
+        return self.select_units(units_to_keep)
+
     def frame_slice(self, start_frame, end_frame):
         from spikeinterface import FrameSliceSorting
         sub_sorting = FrameSliceSorting(


Test patch:
diff --git a/spikeinterface/core/testing_tools.py b/spikeinterface/core/testing_tools.py
index 73889e19e..f66821acf 100644
--- a/spikeinterface/core/testing_tools.py
+++ b/spikeinterface/core/testing_tools.py
@@ -34,20 +34,27 @@ def generate_sorting(
         num_units=5,
         sampling_frequency=30000.,  # in Hz
         durations=[10.325, 3.5],  #  in s for 2 segments
+        empty_units=None
 ):
     num_segments = len(durations)
     num_timepoints = [int(sampling_frequency * d) for d in durations]
 
     unit_ids = np.arange(num_units)
 
+    if empty_units is None:
+        empty_units = []
+
     units_dict_list = []
     for seg_index in range(num_segments):
         units_dict = {}
         for unit_id in unit_ids:
-            #  15 Hz for all units
-            n_spike = int(15. * durations[seg_index])
-            spike_times = np.sort(np.unique(np.random.randint(0, num_timepoints[seg_index], n_spike)))
-            units_dict[unit_id] = spike_times
+            if unit_id not in empty_units:
+                #  15 Hz for all units
+                n_spike = int(15. * durations[seg_index])
+                spike_times = np.sort(np.unique(np.random.randint(0, num_timepoints[seg_index], n_spike)))
+                units_dict[unit_id] = spike_times
+            else:
+                units_dict[unit_id] = np.array([], dtype=int)
         units_dict_list.append(units_dict)
     sorting = NumpySorting.from_dict(units_dict_list, sampling_frequency)
 
diff --git a/spikeinterface/core/tests/test_basesorting.py b/spikeinterface/core/tests/test_basesorting.py
index 047cfe67d..3e4ceb59b 100644
--- a/spikeinterface/core/tests/test_basesorting.py
+++ b/spikeinterface/core/tests/test_basesorting.py
@@ -10,7 +10,7 @@ import numpy as np
 from spikeinterface.core import NpzSortingExtractor, load_extractor
 from spikeinterface.core.base import BaseExtractor
 
-from spikeinterface.core.testing_tools import create_sorting_npz
+from spikeinterface.core.testing_tools import create_sorting_npz, generate_sorting
 
 if hasattr(pytest, "global_test_folder"):
     cache_folder = pytest.global_test_folder / "core"
@@ -65,6 +65,19 @@ def test_BaseSorting():
 
     spikes = sorting.to_spike_vector()
     # print(spikes)
+    
+    # select units
+    keep_units = [0, 1]
+    sorting_select = sorting.select_units(unit_ids=keep_units)
+    for unit in sorting_select.get_unit_ids():
+        assert unit in keep_units
+        
+    # remove empty units
+    empty_units = [1, 3]
+    sorting_empty = generate_sorting(empty_units=empty_units)
+    sorting_clean = sorting_empty.remove_empty_units()
+    for unit in sorting_clean.get_unit_ids():
+        assert unit not in empty_units
 
 
 if __name__ == '__main__':