From b84613b7811c67b5bed9ae46bdd0d83a3cd34e2d Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 09:28:43 +0100 Subject: [PATCH 1/9] Fix channel order bug --- src/spikeinterface/core/sortinganalyzer.py | 13 ++++++---- .../core/tests/test_sortinganalyzer.py | 24 ++++++++++++++++++- 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index c1b4d1f84f..2fd6ebb9be 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/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,8 +1805,7 @@ 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 @@ -1811,8 +1813,9 @@ def select_channels(self, channel_ids) -> "SortingAnalyzer": 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 diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 25aeb78c1a..3d8b0db2a7 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/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 From f21e73ff8973d2c079e794bdb0f32dbc6fcaacae Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 10:56:31 +0100 Subject: [PATCH 2/9] add waveforms and pc slicing --- .../core/analyzer_extension_core.py | 36 +++++++++++++++++++ .../postprocessing/principal_component.py | 36 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index dde1ba8620..4d7b97e527 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/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 ): diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index ad4bc17594..3e0861cc21 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/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 ): From cacbb314ff4d8c845350a7a93385726e6527626e Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 13:17:53 +0100 Subject: [PATCH 3/9] add and fix tests for selected analyzers --- .../core/analyzer_extension_core.py | 2 +- .../core/tests/test_sortinganalyzer.py | 75 ++++++++++++++++++- .../postprocessing/principal_component.py | 4 + 3 files changed, 79 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index 4d7b97e527..7ae5d614e2 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -598,7 +598,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(): diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 3d8b0db2a7..4c41e33eec 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -817,12 +817,85 @@ def test_select_channels(dataset): assert len(p) == len(keep_channel_ids) +def test_select_channels_sparse_waveforms_pca(dataset): + """ + Test that `_select_channels` selects the correct waveforms and 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 = 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", "templates"]) + + # Select channels, in a non-monotonic way + select_channel_ids = np.array(["3", "8", "7"]) + analyzer_seleted = 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(sorting_analyzer.channel_ids, analyzer_seleted.channel_ids_to_indices(analyzer_seleted.channel_ids)) + ) + + original_templates = sorting_analyzer.get_extension("templates").get_data() + selected_templates = analyzer_seleted.get_extension("templates").get_data() + + original_waveforms = sorting_analyzer.get_extension("waveforms") + selected_waveforms = analyzer_seleted.get_extension("waveforms") + + original_pca = sorting_analyzer.get_extension("principal_components") + selected_pca = analyzer_seleted.get_extension("principal_components") + + 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_seleted.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) + + 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: + + # 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] + 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_waveform = selected_waveforms_one_unit[:, :, channel_index_in_selected] + selected_unit_pca = selected_pca_one_unit[:, :, channel_index_in_selected] + + assert np.all(original_unit_waveform == selected_unit_waveform) + assert np.all(original_unit_pca == selected_unit_pca) + + def test_select_channels_independent(dataset): """ Test that `_select_channels` is independent of channel id order. """ recording, sorting = dataset - # Make a very sparse analyzer + # Make a sparse analyzer sorting_analyzer = create_sorting_analyzer( sorting, recording, format="memory", sparse=True, sparsity_kwargs={"method": "radius", "radius_um": 30} ) diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 3e0861cc21..6c5d5c304b 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -132,6 +132,10 @@ def _select_channels_extension_data(self, channel_ids): data = {"pca_projection": new_waveforms[:, :, :max_num_active_channels]} + for key, value in self.data.items(): + if key != "pca_projection": + data[key] = value + return data def _merge_extension_data( From 66396fba6ba40e49ea3941482b4a4fc7b29d65e8 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 13:19:34 +0100 Subject: [PATCH 4/9] respond to alessio --- src/spikeinterface/postprocessing/principal_component.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 6c5d5c304b..1a5059aae8 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -117,20 +117,20 @@ def _select_channels_extension_data(self, channel_ids): 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) + new_pcs = np.zeros_like(old_pcs) max_num_active_channels = 0 - for waveform_index, (waveform, spike) in enumerate(zip(old_pcs, random_spikes)): + for pc_index, (old_pc, 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] + new_pcs[pc_index, :, :size_of_new_mask] = old_pc[:, channel_slice] max_num_active_channels = max(max_num_active_channels, size_of_new_mask) - data = {"pca_projection": new_waveforms[:, :, :max_num_active_channels]} + data = {"pca_projection": new_pcs[:, :, :max_num_active_channels]} for key, value in self.data.items(): if key != "pca_projection": From c9e1caa88ea3014c0bcac647412cc4a8c7cbd7ba Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 13:20:03 +0100 Subject: [PATCH 5/9] oups in tests --- src/spikeinterface/core/tests/test_sortinganalyzer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 4c41e33eec..c72f90e34f 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -842,7 +842,7 @@ def test_select_channels_sparse_waveforms_pca(dataset): zip(sorting_analyzer.channel_ids, sorting_analyzer.channel_ids_to_indices(sorting_analyzer.channel_ids)) ) selected_id_index_map = dict( - zip(sorting_analyzer.channel_ids, analyzer_seleted.channel_ids_to_indices(analyzer_seleted.channel_ids)) + zip(analyzer_seleted.channel_ids, analyzer_seleted.channel_ids_to_indices(analyzer_seleted.channel_ids)) ) original_templates = sorting_analyzer.get_extension("templates").get_data() From 9502b828470a28ec3d00d3eb67980e4c2a9f3487 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 13:26:34 +0100 Subject: [PATCH 6/9] Move PCA test to postprocessing --- .../core/tests/test_sortinganalyzer.py | 15 ++---- .../tests/test_principal_component.py | 47 +++++++++++++++++++ 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index c72f90e34f..1c73fc7605 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -817,9 +817,9 @@ def test_select_channels(dataset): assert len(p) == len(keep_channel_ids) -def test_select_channels_sparse_waveforms_pca(dataset): +def test_select_channels_sparse_waveforms_templates(dataset): """ - Test that `_select_channels` selects the correct waveforms and principal components when the analyzer + 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 @@ -831,7 +831,7 @@ def test_select_channels_sparse_waveforms_pca(dataset): 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", "templates"]) + sorting_analyzer.compute(["random_spikes", "waveforms", "templates"]) # Select channels, in a non-monotonic way select_channel_ids = np.array(["3", "8", "7"]) @@ -851,9 +851,6 @@ def test_select_channels_sparse_waveforms_pca(dataset): original_waveforms = sorting_analyzer.get_extension("waveforms") selected_waveforms = analyzer_seleted.get_extension("waveforms") - original_pca = sorting_analyzer.get_extension("principal_components") - selected_pca = analyzer_seleted.get_extension("principal_components") - 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] @@ -862,9 +859,6 @@ def test_select_channels_sparse_waveforms_pca(dataset): original_waveforms_one_unit = original_waveforms.get_waveforms_one_unit(unit_id) selected_waveforms_one_unit = selected_waveforms.get_waveforms_one_unit(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: @@ -880,14 +874,11 @@ def test_select_channels_sparse_waveforms_pca(dataset): # 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] - 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_waveform = selected_waveforms_one_unit[:, :, channel_index_in_selected] - selected_unit_pca = selected_pca_one_unit[:, :, channel_index_in_selected] assert np.all(original_unit_waveform == selected_unit_waveform) - assert np.all(original_unit_pca == selected_unit_pca) def test_select_channels_independent(dataset): diff --git a/src/spikeinterface/postprocessing/tests/test_principal_component.py b/src/spikeinterface/postprocessing/tests/test_principal_component.py index 0e65bb2338..038483aaf8 100644 --- a/src/spikeinterface/postprocessing/tests/test_principal_component.py +++ b/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_seleted = sorting_analyzer._select_channels(channel_ids=select_channel_ids) + + # Prepare the data + original_pca = sorting_analyzer.get_extension("principal_components") + selected_pca = analyzer_seleted.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_seleted.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) From 821aeb61085f252861e2a35c4579dc0093cd2ed8 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 13:31:48 +0100 Subject: [PATCH 7/9] tweak loop when selecting --- src/spikeinterface/core/analyzer_extension_core.py | 4 ++-- src/spikeinterface/postprocessing/principal_component.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index 7ae5d614e2..c461a06019 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -262,6 +262,7 @@ 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 + # Compute how to slice the original sparsity to get the newly selected sparsity unit_sparsity_slices = {} for unit_id in unit_ids: channel_indices = [] @@ -278,9 +279,8 @@ def _select_channels_extension_data(self, channel_ids): new_waveforms = np.zeros_like(old_waveforms) max_num_active_channels = 0 - for waveform_index, (waveform, spike) in enumerate(zip(old_waveforms, random_spikes)): + for waveform_index, (waveform, unit_index) in enumerate(zip(old_waveforms, random_spikes["unit_index"])): - unit_index = spike["unit_index"] unit_id = unit_ids[unit_index] channel_slice = unit_sparsity_slices[unit_id] diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 1a5059aae8..257d02f95d 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -103,6 +103,7 @@ 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 + # Compute how to slice the original sparsity to get the newly selected sparsity unit_sparsity_slices = {} for unit_id in unit_ids: channel_indices = [] @@ -119,9 +120,8 @@ def _select_channels_extension_data(self, channel_ids): new_pcs = np.zeros_like(old_pcs) max_num_active_channels = 0 - for pc_index, (old_pc, spike) in enumerate(zip(old_pcs, random_spikes)): + for pc_index, (old_pc, unit_index) in enumerate(zip(old_pcs, random_spikes["unit_index"])): - unit_index = spike["unit_index"] unit_id = unit_ids[unit_index] channel_slice = unit_sparsity_slices[unit_id] From 3132f3baf3b848eeff181138c941a65446e290f9 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 24 Jul 2026 13:54:12 +0100 Subject: [PATCH 8/9] deal with zero channel case --- src/spikeinterface/core/analyzer_extension_core.py | 8 +++++--- src/spikeinterface/postprocessing/principal_component.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index c461a06019..ce416d45e2 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -284,10 +284,12 @@ def _select_channels_extension_data(self, channel_ids): 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] + if len(channel_slice) > 0: - max_num_active_channels = max(max_num_active_channels, size_of_new_mask) + 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]} diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 257d02f95d..496367919a 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -125,10 +125,12 @@ def _select_channels_extension_data(self, channel_ids): unit_id = unit_ids[unit_index] channel_slice = unit_sparsity_slices[unit_id] - size_of_new_mask = len(channel_slice) - new_pcs[pc_index, :, :size_of_new_mask] = old_pc[:, channel_slice] + if len(channel_slice) > 0: - max_num_active_channels = max(max_num_active_channels, size_of_new_mask) + size_of_new_mask = len(channel_slice) + new_pcs[pc_index, :, :size_of_new_mask] = old_pc[:, channel_slice] + + max_num_active_channels = max(max_num_active_channels, size_of_new_mask) data = {"pca_projection": new_pcs[:, :, :max_num_active_channels]} From 14d9a00f2dc2e0fce7f9a48ec8abe32e8d103b80 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Mon, 27 Jul 2026 10:35:41 +0100 Subject: [PATCH 9/9] respond to alessio review --- src/spikeinterface/core/sortinganalyzer.py | 2 +- src/spikeinterface/core/tests/test_sortinganalyzer.py | 10 +++++----- .../postprocessing/tests/test_principal_component.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 2fd6ebb9be..09993437f1 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1789,7 +1789,7 @@ def _select_channels(self, channel_ids) -> "SortingAnalyzer": 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] + 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) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 1c73fc7605..4ce383a093 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -835,26 +835,26 @@ def test_select_channels_sparse_waveforms_templates(dataset): # Select channels, in a non-monotonic way select_channel_ids = np.array(["3", "8", "7"]) - analyzer_seleted = sorting_analyzer._select_channels(channel_ids=select_channel_ids) + 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_seleted.channel_ids, analyzer_seleted.channel_ids_to_indices(analyzer_seleted.channel_ids)) + 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_seleted.get_extension("templates").get_data() + selected_templates = analyzer_selected.get_extension("templates").get_data() original_waveforms = sorting_analyzer.get_extension("waveforms") - selected_waveforms = analyzer_seleted.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_seleted.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) diff --git a/src/spikeinterface/postprocessing/tests/test_principal_component.py b/src/spikeinterface/postprocessing/tests/test_principal_component.py index 038483aaf8..879f4ff055 100644 --- a/src/spikeinterface/postprocessing/tests/test_principal_component.py +++ b/src/spikeinterface/postprocessing/tests/test_principal_component.py @@ -216,16 +216,16 @@ def test_select_channels_sparse_pca(): # Select channels, in a non-monotonic way select_channel_ids = np.array(["3", "8", "7"]) - analyzer_seleted = sorting_analyzer._select_channels(channel_ids=select_channel_ids) + 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_seleted.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_seleted.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)