From 57ad0cf9db435821cef50231a4e0fd9795098658 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 25 Sep 2026 18:10:11 +0200 Subject: [PATCH] perf: avoid duplication of PCA array with reshape --- spikeinterface_gui/controller.py | 42 +++++++++++++++++++++++------ spikeinterface_gui/ndscatterview.py | 8 +++--- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/spikeinterface_gui/controller.py b/spikeinterface_gui/controller.py index c765f60..9263efa 100644 --- a/spikeinterface_gui/controller.py +++ b/spikeinterface_gui/controller.py @@ -891,16 +891,42 @@ def get_units_table(self): return self.units_table def get_all_pcs(self): - if self._pc_projections is None and self.pc_ext is not None: - self._pc_projections, self._pc_indices = self.pc_ext.get_some_projections( - channel_ids=self.analyzer.channel_ids, - unit_ids=self.analyzer.unit_ids - ) - - return self._pc_indices, self._pc_projections - else: + self._pc_indices, self._pc_projections = self._get_dense_pcs_channel_major() + if self._pc_projections is None: return None, None + return self._pc_indices, self._pc_projections + + def _get_dense_pcs_channel_major(self): + """Dense PCs of the random spikes as (num_spikes, num_channels, num_components), C-contiguous. + + Channel-major so that NDScatterView can flatten it with a free reshape + (column = channel_index * num_components + component) instead of holding a + second full copy. Built directly from the sparse projections, so the peak is + one dense array + the sparse one. + """ + sparsity = self.analyzer.sparsity + if sparsity is None: + pcs, unit_indices = self.pc_ext.get_some_projections( + channel_ids=self.analyzer.channel_ids, unit_ids=self.analyzer.unit_ids + ) + return unit_indices, np.ascontiguousarray(pcs.swapaxes(1, 2)) + + # read all rows at once: per-unit reads on a lazy zarr array re-fetch the same chunks + sparse_pcs = np.asarray(self.pc_ext.data["pca_projection"]) + num_spikes, num_components, _ = sparse_pcs.shape + unit_indices = np.asarray(self.analyzer.get_extension("random_spikes").get_random_spikes()["unit_index"]) + + dense = np.zeros((num_spikes, self.num_channels, num_components), dtype=sparse_pcs.dtype) + order = np.argsort(unit_indices, kind="stable") + bounds = np.searchsorted(unit_indices[order], np.arange(len(self.analyzer.unit_ids) + 1)) + for unit_index, unit_id in enumerate(self.analyzer.unit_ids): + rows = order[bounds[unit_index]:bounds[unit_index + 1]] + if rows.size == 0: + continue + chan_inds = sparsity.unit_id_to_channel_indices[unit_id] + dense[rows[:, None], chan_inds[None, :], :] = sparse_pcs[rows, :, : chan_inds.size].swapaxes(1, 2) + return unit_indices.copy(), dense def get_sparsity_mask(self): if self.external_sparsity is not None: diff --git a/spikeinterface_gui/ndscatterview.py b/spikeinterface_gui/ndscatterview.py index 77864c3..8d63c09 100644 --- a/spikeinterface_gui/ndscatterview.py +++ b/spikeinterface_gui/ndscatterview.py @@ -26,7 +26,9 @@ def __init__(self, controller=None, parent=None, backend="qt"): assert controller.has_extension('principal_components') self.pc_unit_index, self.pc_data = controller.get_all_pcs() - self.data = self.pc_data.swapaxes(1,2).reshape(self.pc_data.shape[0], -1) + # pc_data is (num_spikes, num_channels, num_components) and C-contiguous, so this is a view + # column = channel_index * num_components + component (see update_selected_components) + self.data = self.pc_data.reshape(self.pc_data.shape[0], -1) self.random_spikes_indices = controller.random_spikes_indices if self.data.shape[1] == 1: @@ -262,7 +264,7 @@ def update_selected_components(self): if self._best_mode: self.selected_comp[:] = True return - n_pc_per_chan = self.pc_data.shape[1] + n_pc_per_chan = self.pc_data.shape[2] n = min(self.settings['num_pc_per_channel'], n_pc_per_chan) self.selected_comp[:] = False for i in range(n): @@ -367,7 +369,7 @@ def _qt_make_layout(self): # self.graphicsview2.setMaximumSize(200, 200) - self.settings.param('num_pc_per_channel').setLimits((1, self.pc_data.shape[1])) + self.settings.param('num_pc_per_channel').setLimits((1, self.pc_data.shape[2])) # the color vector is precomputed # spike_colors = self.controller.get_spike_colors(self.pc_unit_index)