Fix channel order bug and make select_channels private

src/spikeinterface/core/analyzer_extension_core.py

@@ -257,6 +257,42 @@ def _select_units_extension_data(self, unit_ids):
 
         return new_data
 
+    def _select_channels_extension_data(self, channel_ids):
+
+        unit_ids = self.sorting_analyzer.unit_ids
+        old_unit_id_to_channel_ids = self.sorting_analyzer.sparsity.unit_id_to_channel_ids
+
+        unit_sparsity_slices = {}
+        for unit_id in unit_ids:
+            channel_indices = []
+            unit_channel_ids = old_unit_id_to_channel_ids[unit_id]
+
+            for channel_id in channel_ids:
+                if channel_id in unit_channel_ids:
+                    idx = np.where(old_unit_id_to_channel_ids[unit_id] == channel_id)[0][0]
+                    channel_indices.append(idx)
+            unit_sparsity_slices[unit_id] = np.array(channel_indices)
+
+        old_waveforms = self.data["waveforms"]
+        random_spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes()
+
+        new_waveforms = np.zeros_like(old_waveforms)
+        max_num_active_channels = 0
+        for waveform_index, (waveform, spike) in enumerate(zip(old_waveforms, random_spikes)):
+
+            unit_index = spike["unit_index"]
+            unit_id = unit_ids[unit_index]
+            channel_slice = unit_sparsity_slices[unit_id]
+
+            size_of_new_mask = len(channel_slice)
+            new_waveforms[waveform_index, :, :size_of_new_mask] = waveform[:, channel_slice]
+
+            max_num_active_channels = max(max_num_active_channels, size_of_new_mask)
+
+        data = {"waveforms": new_waveforms[:, :, :max_num_active_channels]}
+
+        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/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 = [np.where(self.channel_ids == id)[0][0] for id in 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,28 @@ def test_select_channels(dataset):
         assert len(p) == len(keep_channel_ids)
 
 
+def test_select_channels_independent(dataset):
+    """
+    Test that `_select_channels` is independent of channel id order.
+    """
+    recording, sorting = dataset
+    # Make a very 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

@@ -98,6 +98,42 @@ def _select_units_extension_data(self, unit_ids):
                 new_data[k] = v
         return new_data
 
+    def _select_channels_extension_data(self, channel_ids):
+
+        unit_ids = self.sorting_analyzer.unit_ids
+        old_unit_id_to_channel_ids = self.sorting_analyzer.sparsity.unit_id_to_channel_ids
+
+        unit_sparsity_slices = {}
+        for unit_id in unit_ids:
+            channel_indices = []
+            unit_channel_ids = old_unit_id_to_channel_ids[unit_id]
+
+            for channel_id in channel_ids:
+                if channel_id in unit_channel_ids:
+                    idx = np.where(old_unit_id_to_channel_ids[unit_id] == channel_id)[0][0]
+                    channel_indices.append(idx)
+            unit_sparsity_slices[unit_id] = np.array(channel_indices)
+
+        old_pcs = self.data["pca_projection"]
+        random_spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes()
+
+        new_waveforms = np.zeros_like(old_pcs)
+        max_num_active_channels = 0
+        for waveform_index, (waveform, spike) in enumerate(zip(old_pcs, random_spikes)):
+
+            unit_index = spike["unit_index"]
+            unit_id = unit_ids[unit_index]
+            channel_slice = unit_sparsity_slices[unit_id]
+
+            size_of_new_mask = len(channel_slice)
+            new_waveforms[waveform_index, :, :size_of_new_mask] = waveform[:, channel_slice]
+
+            max_num_active_channels = max(max_num_active_channels, size_of_new_mask)
+
+        data = {"pca_projection": new_waveforms[:, :, :max_num_active_channels]}
+
+        return data
+
     def _merge_extension_data(
         self, merge_unit_groups, new_unit_ids, new_sorting_analyzer, keep_mask=None, verbose=False, **job_kwargs
     ):