Add SortingAnalyzer.split() function for aggregated analyzers (#4660)

Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com>

Author: alejoe91

src/spikeinterface/core/sortinganalyzer.py

@@ -1280,6 +1280,37 @@ def get_main_channels(self, outputs: Literal["index", "id"] = "index", with_dict
         else:
             return main_chans
 
+    def is_aggregated(self):
+        """
+        Returns True if the SortingAnalyzer is aggregated, False otherwise.
+        """
+        return (
+            "aggregation_key" in self.get_sorting_property_keys()
+            and "aggregation_key" in self.get_recording_property_keys()
+        )
+
+    def split_by(self):
+        """
+        Returns a dictionary of SortingAnalyzer objects, split by the aggregation_key.
+        The keys of the dictionary are the unique values of the aggregation_key, and the values
+        are the SortingAnalyzer objects corresponding to each unique value.
+        """
+        if not self.is_aggregated():
+            raise ValueError("SortingAnalyzer is not aggregated")
+
+        units_aggregation_key = self.get_sorting_property("aggregation_key")
+        channel_aggregation_key = self.get_recording_property("aggregation_key")
+        unique_keys = np.unique(units_aggregation_key)
+        split_analyzers = {}
+        for key in unique_keys:
+            unit_ids = self.unit_ids[units_aggregation_key == key]
+            channel_ids = self.channel_ids[channel_aggregation_key == key]
+            analyzer_units = self.select_units(unit_ids)
+            analyzer_split = analyzer_units.select_channels(channel_ids)
+            split_analyzers[key] = analyzer_split
+
+        return split_analyzers
+
     def are_units_mergeable(
         self,
         merge_unit_groups: list[str | int],

src/spikeinterface/core/tests/test_sortinganalyzer.py

@@ -277,6 +277,12 @@ def test_create_by_dict():
         assert not np.any(sparsity_mask[i][other_group_channel_indices])
         assert sparsity_mask[i, main_channel_index]
 
+    # test split aggregated
+    analyzers_split = analyzer_with_sparsity.split_by()
+    for group, analyzer in analyzers_split.items():
+        assert np.all(analyzer.get_sorting_property("aggregation_key") == group)
+        assert np.all(analyzer.recording.get_property("aggregation_key") == group)
+
 
 def test_load_without_runtime_info(tmp_path, dataset):
     import zarr