Fix channel order bug and make select_channels private

src/spikeinterface/core/analyzer_extension_core.py

@@ -257,6 +257,15 @@ def _select_units_extension_data(self, unit_ids):
 
         return new_data
 
+    def _select_channels_extension_data(self, channel_ids):
+
+        old_waveforms = self.data["waveforms"]
+
+        new_waveforms = _select_channels_sparse_data(self.sorting_analyzer, old_waveforms, channel_ids)
+        data = {"waveforms": new_waveforms}
+
+        return data
+
     def _merge_extension_data(
         self, merge_unit_groups, new_unit_ids, new_sorting_analyzer, keep_mask=None, verbose=False, **job_kwargs
     ):
@@ -562,7 +571,7 @@ def _select_units_extension_data(self, unit_ids):
         return new_data
 
     def _select_channels_extension_data(self, channel_ids):
-        keep_channel_indices = np.flatnonzero(np.isin(self.sorting_analyzer.channel_ids, channel_ids))
+        keep_channel_indices = [np.where(self.sorting_analyzer.channel_ids == id)[0][0] for id in channel_ids]
 
         new_data = {}
         for key, arr in self.data.items():
@@ -1727,3 +1736,43 @@ def _update_data_after_merge_or_split(old_analyzer, new_analyzer, old_arr, new_s
         raise NotImplementedError(
             "Only pandas DataFrame and numpy array are supported for merging and splitting extension data."
         )
+
+
+def _select_channels_sparse_data(old_sorting_analyzer, old_data, new_channel_ids):
+    """
+    When selecting channels from extensions which are sparse (waveforms, pcs),
+    this function remaps the sparsity for the underlying data.
+    """
+
+    unit_ids = old_sorting_analyzer.unit_ids
+    old_unit_id_to_channel_ids = old_sorting_analyzer.sparsity.unit_id_to_channel_ids
+
+    # Compute how to slice the original sparsity to get the new sparsity
+    unit_sparsity_slices = {}
+    for unit_index, unit_id in enumerate(unit_ids):
+        unit_sparsity_channel_indices = []
+        unit_channel_ids = old_unit_id_to_channel_ids[unit_id]
+
+        for channel_id in new_channel_ids:
+            if channel_id in unit_channel_ids:
+                idx = np.where(old_unit_id_to_channel_ids[unit_id] == channel_id)[0][0]
+                unit_sparsity_channel_indices.append(idx)
+        unit_sparsity_slices[unit_id] = np.array(unit_sparsity_channel_indices)
+
+    random_spikes = old_sorting_analyzer.get_extension("random_spikes").get_random_spikes()
+
+    new_data = np.zeros_like(old_data)
+    max_num_active_channels = 0
+    for data_index, (one_old_data, unit_index) in enumerate(zip(old_data, random_spikes["unit_index"])):
+
+        unit_id = unit_ids[unit_index]
+        channel_slice = unit_sparsity_slices[unit_id]
+
+        if len(channel_slice) > 0:
+
+            size_of_new_mask = len(channel_slice)
+            new_data[data_index, :, :size_of_new_mask] = one_old_data[:, channel_slice]
+
+            max_num_active_channels = max(max_num_active_channels, size_of_new_mask)
+
+    return new_data[:, :, :max_num_active_channels]

src/spikeinterface/core/sortinganalyzer.py

@@ -1404,7 +1404,7 @@ def split_by(self):
             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)
+            analyzer_split = analyzer_units._select_channels(channel_ids)
             split_analyzers[key] = analyzer_split
 
         return split_analyzers
@@ -1767,7 +1767,7 @@ def select_units(self, unit_ids, format="memory", folder=None) -> "SortingAnalyz
             folder = clean_zarr_folder_name(folder)
         return self._save_or_select_or_merge_or_split(format=format, folder=folder, unit_ids=unit_ids)
 
-    def select_channels(self, channel_ids) -> "SortingAnalyzer":
+    def _select_channels(self, channel_ids) -> "SortingAnalyzer":
         """
         This method is equivalent to `save_as()` but with a subset of channels.
         Filters channels by creating a new sorting analyzer object in a new folder.
@@ -1788,6 +1788,9 @@ def select_channels(self, channel_ids) -> "SortingAnalyzer":
         if not np.all(np.isin(channel_ids, self.channel_ids)):
             wrong_channel_ids = [ch for ch in channel_ids if ch not in self.channel_ids]
             raise ValueError(f"Some channel_ids are not in the current channel_ids: {wrong_channel_ids}")
+
+        select_channel_indices_in_old_recording = self.channel_ids_to_indices(channel_ids)
+
         if self.has_recording() or self.has_temporary_recording():
             new_recording = self.recording.select_channels(channel_ids)
             new_rec_attributes = None
@@ -1802,17 +1805,17 @@ def select_channels(self, channel_ids) -> "SortingAnalyzer":
                     values_arr = np.array(values)
                     if len(values_arr) == len(self.channel_ids):
                         # only slice properties that have the same length as channel_ids
-                        channel_indices = [np.where(self.channel_ids == id)[0][0] for id in channel_ids]
-                        new_properties[key] = values_arr[channel_indices]
+                        new_properties[key] = values_arr[select_channel_indices_in_old_recording]
                     else:
                         new_properties[key] = values_arr
                 new_rec_attributes["properties"] = new_properties
             if new_rec_attributes.get("probegroup") is not None:
                 slice_indices = self.channel_ids_to_indices(channel_ids)
                 new_probegroup = new_rec_attributes["probegroup"].get_slice(slice_indices)
                 new_rec_attributes["probegroup"] = new_probegroup
+
         if self.sparsity is not None:
-            sparsity_mask = self.sparsity.mask[:, np.isin(self.channel_ids, channel_ids)]
+            sparsity_mask = self.sparsity.mask[:, select_channel_indices_in_old_recording]
             new_sparsity = ChannelSparsity(sparsity_mask, self.unit_ids, np.array(channel_ids))
         else:
             new_sparsity = None

src/spikeinterface/core/tests/test_sortinganalyzer.py

@@ -798,7 +798,7 @@ def test_select_channels(dataset):
     sorting_analyzer.compute(["random_spikes", "templates", "noise_levels"])
     # select channels
     keep_channel_ids = recording.channel_ids[::2]
-    sorting_analyzer2 = sorting_analyzer.select_channels(channel_ids=keep_channel_ids)
+    sorting_analyzer2 = sorting_analyzer._select_channels(channel_ids=keep_channel_ids)
 
     assert np.array_equal(sorting_analyzer2.channel_ids, keep_channel_ids)
     assert np.array_equal(sorting_analyzer2.get_channel_locations(), recording.get_channel_locations(keep_channel_ids))
@@ -817,6 +817,92 @@ def test_select_channels(dataset):
         assert len(p) == len(keep_channel_ids)
 
 
+def test_select_channels_sparse_waveforms_templates(dataset):
+    """
+    Test that `_select_channels` selects the correct waveforms and templates when the analyzer
+    is sparse.
+
+    The actual code uses fancy indexing etc, so this test is designed to _not_ do this, and instead
+    just loop over all units and channels to check consistency.
+    """
+
+    recording, sorting = dataset
+    # Make a sparse analyzer
+    sorting_analyzer = create_sorting_analyzer(
+        sorting, recording, format="memory", sparse=True, sparsity_kwargs={"method": "radius", "radius_um": 30}
+    )
+    sorting_analyzer.compute(["random_spikes", "waveforms", "templates"])
+
+    # Select channels, in a non-monotonic way
+    select_channel_ids = np.array(["3", "8", "7"])
+    analyzer_selected = sorting_analyzer._select_channels(channel_ids=select_channel_ids)
+
+    # Prepare the data
+    original_id_index_map = dict(
+        zip(sorting_analyzer.channel_ids, sorting_analyzer.channel_ids_to_indices(sorting_analyzer.channel_ids))
+    )
+    selected_id_index_map = dict(
+        zip(analyzer_selected.channel_ids, analyzer_selected.channel_ids_to_indices(analyzer_selected.channel_ids))
+    )
+
+    original_templates = sorting_analyzer.get_extension("templates").get_data()
+    selected_templates = analyzer_selected.get_extension("templates").get_data()
+
+    original_waveforms = sorting_analyzer.get_extension("waveforms")
+    selected_waveforms = analyzer_selected.get_extension("waveforms")
+
+    for unit_index, unit_id in enumerate(sorting_analyzer.unit_ids):
+
+        original_units_to_channels = sorting_analyzer.sparsity.unit_id_to_channel_ids[unit_id]
+        selected_units_to_channels = analyzer_selected.sparsity.unit_id_to_channel_ids[unit_id]
+
+        original_waveforms_one_unit = original_waveforms.get_waveforms_one_unit(unit_id)
+        selected_waveforms_one_unit = selected_waveforms.get_waveforms_one_unit(unit_id)
+
+        for channel_id in select_channel_ids:
+            if channel_id in original_units_to_channels:
+
+                # Check templates, which are dense
+                original_channel_index = original_id_index_map[channel_id]
+                selected_channel_index = selected_id_index_map[channel_id]
+
+                original_channel = original_templates[unit_index, :, original_channel_index]
+                selected_channel = selected_templates[unit_index, :, selected_channel_index]
+
+                assert np.all(original_channel == selected_channel)
+
+                # Now check waveforms and PCs, which are sparse
+                channel_index_in_original = np.where(original_units_to_channels == channel_id)[0][0]
+                original_unit_waveform = original_waveforms_one_unit[:, :, channel_index_in_original]
+
+                channel_index_in_selected = np.where(selected_units_to_channels == channel_id)[0][0]
+                selected_unit_waveform = selected_waveforms_one_unit[:, :, channel_index_in_selected]
+
+                assert np.all(original_unit_waveform == selected_unit_waveform)
+
+
+def test_select_channels_independent(dataset):
+    """
+    Test that `_select_channels` is independent of channel id order.
+    """
+    recording, sorting = dataset
+    # Make a sparse analyzer
+    sorting_analyzer = create_sorting_analyzer(
+        sorting, recording, format="memory", sparse=True, sparsity_kwargs={"method": "radius", "radius_um": 30}
+    )
+
+    select_channel_ids = np.array(["3", "7", "8"])
+    sa_one = sorting_analyzer._select_channels(channel_ids=select_channel_ids)
+
+    # Make another analyzer with select_channel_ids ['7', '3', '8']
+    shuffle_order = np.array([1, 0, 2])
+    second_channel = select_channel_ids[shuffle_order]
+    sa_two = sorting_analyzer._select_channels(channel_ids=second_channel)
+
+    assert np.all(sa_one.sparsity.mask == sa_two.sparsity.mask[:, shuffle_order])
+    assert np.all(sa_one.get_channel_locations() == sa_two.get_channel_locations()[shuffle_order])
+
+
 def test_main_channel_from_templates_dense_recordingless(tmp_path):
     """When a dense analyzer has a `templates` extension but no attached recording, `get_main_channels`
     recovers each unit's main channel from the templates (its peak channel) and reports it consistently

src/spikeinterface/postprocessing/principal_component.py

@@ -12,7 +12,7 @@
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
 from spikeinterface.core.core_tools import slice_rows
 from spikeinterface.core.job_tools import TimeSeriesChunkExecutor, _shared_job_kwargs_doc, fix_job_kwargs
-from spikeinterface.core.analyzer_extension_core import _inplace_sparse_realign_waveforms
+from spikeinterface.core.analyzer_extension_core import _inplace_sparse_realign_waveforms, _select_channels_sparse_data
 
 _possible_modes = ["by_channel_local", "by_channel_global", "concatenated"]
 
@@ -98,6 +98,19 @@ def _select_units_extension_data(self, unit_ids):
                 new_data[k] = v
         return new_data
 
+    def _select_channels_extension_data(self, channel_ids):
+
+        old_pcs = self.data["pca_projection"]
+        new_pcs = _select_channels_sparse_data(self.sorting_analyzer, old_pcs, channel_ids)
+
+        data = {"pca_projection": new_pcs}
+
+        for key, value in self.data.items():
+            if key != "pca_projection":
+                data[key] = value
+
+        return data
+
     def _merge_extension_data(
         self, merge_unit_groups, new_unit_ids, new_sorting_analyzer, keep_mask=None, verbose=False, **job_kwargs
     ):

src/spikeinterface/postprocessing/tests/test_principal_component.py

@@ -3,6 +3,8 @@
 
 from spikeinterface.postprocessing import ComputePrincipalComponents
 from spikeinterface.postprocessing.tests.common_extension_tests import AnalyzerExtensionCommonTestSuite
+from spikeinterface.core.tests.test_sortinganalyzer import get_dataset
+from spikeinterface.core import create_sorting_analyzer
 
 
 class TestPrincipalComponentsExtension(AnalyzerExtensionCommonTestSuite):
@@ -196,6 +198,51 @@ def test_project_new(self):
         assert new_proj.shape[2] == ext_pca.data["pca_projection"].shape[2]
 
 
+def test_select_channels_sparse_pca():
+    """
+    Test that `_select_channels` selects the correct principal components when the analyzer
+    is sparse.
+
+    The actual code uses fancy indexing etc, so this test is designed to _not_ do this, and instead
+    just loop over all units and channels to check consistency.
+    """
+
+    recording, sorting = get_dataset()
+    # Make a sparse analyzer
+    sorting_analyzer = create_sorting_analyzer(
+        sorting, recording, format="memory", sparse=True, sparsity_kwargs={"method": "radius", "radius_um": 30}
+    )
+    sorting_analyzer.compute(["random_spikes", "waveforms", "principal_components"])
+
+    # Select channels, in a non-monotonic way
+    select_channel_ids = np.array(["3", "8", "7"])
+    analyzer_selected = sorting_analyzer._select_channels(channel_ids=select_channel_ids)
+
+    # Prepare the data
+    original_pca = sorting_analyzer.get_extension("principal_components")
+    selected_pca = analyzer_selected.get_extension("principal_components")
+
+    for unit_id in sorting_analyzer.unit_ids:
+
+        original_units_to_channels = sorting_analyzer.sparsity.unit_id_to_channel_ids[unit_id]
+        selected_units_to_channels = analyzer_selected.sparsity.unit_id_to_channel_ids[unit_id]
+
+        original_pca_one_unit, _ = original_pca.get_projections_one_unit(unit_id, sparse=True)
+        selected_pca_one_unit, _ = selected_pca.get_projections_one_unit(unit_id, sparse=True)
+
+        for channel_id in select_channel_ids:
+            if channel_id in original_units_to_channels:
+
+                # Now check waveforms and PCs, which are sparse
+                channel_index_in_original = np.where(original_units_to_channels == channel_id)[0][0]
+                original_unit_pca = original_pca_one_unit[:, :, channel_index_in_original]
+
+                channel_index_in_selected = np.where(selected_units_to_channels == channel_id)[0][0]
+                selected_unit_pca = selected_pca_one_unit[:, :, channel_index_in_selected]
+
+                assert np.all(original_unit_pca == selected_unit_pca)
+
+
 if __name__ == "__main__":
     test = TestPrincipalComponentsExtension()
     test.test_get_projections(sparse=True)