Add split_by method to BaseSorting

Author: alejoe91

src/spikeinterface/core/baserecordingsnippets.py

@@ -514,7 +514,7 @@ def split_by(self, property="group", outputs="dict"):
             recordings = {}
         for value in np.unique(values).tolist():
             (inds,) = np.nonzero(values == value)
-            new_channel_ids = self.get_channel_ids()[inds]
+            new_channel_ids = self.channel_ids[inds]
             subrec = self.select_channels(new_channel_ids)
             subrec.set_annotation("split_by_property", value=property)
             if outputs == "list":

src/spikeinterface/core/basesorting.py

@@ -626,6 +626,47 @@ def time_slice(self, start_time: float | None, end_time: float | None) -> BaseSo
 
         return self.frame_slice(start_frame=start_frame, end_frame=end_frame)
 
+    def split_by(self, property="group", outputs="dict"):
+        """
+        Splits object based on a certain property (e.g. "group")
+
+        Parameters
+        ----------
+        property : str, default: "group"
+            The property to use to split the object, default: "group"
+        outputs : "dict" | "list", default: "dict"
+            Whether to return a dict or a list
+
+        Returns
+        -------
+        dict or list
+            A dict or list with grouped objects based on property
+
+        Raises
+        ------
+        ValueError
+            Raised when property is not present
+        """
+        assert outputs in ("list", "dict")
+        values = self.get_property(property)
+        if values is None:
+            raise ValueError(f"property {property} is not set")
+
+        if outputs == "list":
+            sortings = []
+        elif outputs == "dict":
+            sortings = {}
+        for value in np.unique(values).tolist():
+            (inds,) = np.nonzero(values == value)
+            new_channel_ids = self.unit_ids[inds]
+            subsort = self.select_units(new_channel_ids)
+            subsort.set_annotation("split_by_property", value=property)
+            if outputs == "list":
+                sortings.append(subsort)
+            elif outputs == "dict":
+                sortings[value] = subsort
+        return sortings
+
     def time_to_sample_index(self, time, segment_index=0):
         """
         Transform time in seconds into sample index

src/spikeinterface/core/tests/test_sortinganalyzer.py

@@ -215,10 +215,9 @@ def test_create_by_dict():
     sort.set_property(key="group", values=[2, 2, 2, 1, 2, 2, 2, 1, 2, 1])
 
     unit_ids = sort.unit_ids
-
-    split_sort = {key: sort.select_units(unit_ids=unit_ids[sort.get_property("group") == key]) for key in [1, 2]}
-
-    analyzer = create_sorting_analyzer(split_sort, rec.split_by("group"))
+    split_sort = sort.split_by("group")
+    split_rec = rec.split_by("group")
+    analyzer = create_sorting_analyzer(split_sort, split_rec)
     analyzer_unit_ids = analyzer.unit_ids
 
     assert set(analyzer.unit_ids) == set(sort.unit_ids)