From a3ead81655dfad247f5f93d8762d9237f07a58cf Mon Sep 17 00:00:00 2001 From: Sebastien Date: Mon, 8 Jul 2024 09:55:14 +0200 Subject: [PATCH] Some more optimizations, to reduce memory after merges --- src/spikeinterface/core/analyzer_extension_core.py | 4 ++-- src/spikeinterface/postprocessing/principal_component.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index d2dc771500..98e9e225dd 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -255,6 +255,7 @@ def _merge_extension_data( if sparsity is not None: + max_num_chans = np.sum(new_sorting_analyzer.sparsity.mask, 1).max() new_data["waveforms"] = waveforms.copy() for to_be_merged, unit_id in zip(units_to_merge, new_unit_ids): new_channel_ids = sparsity.unit_id_to_channel_ids[unit_id] @@ -263,8 +264,7 @@ def _merge_extension_data( ) num_chans = new_waveforms.shape[2] new_data["waveforms"][spike_indices, :, :num_chans] = new_waveforms - new_data["waveforms"][spike_indices, :, num_chans:] = 0 - + new_data["waveforms"] = new_data["waveforms"][:, :, :max_num_chans] else: new_data["waveforms"] = waveforms diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index e357ce9df5..fe6e842453 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -118,6 +118,7 @@ def _merge_extension_data( if sparsity is not None: + max_num_chans = np.sum(new_sorting_analyzer.sparsity.mask, 1).max() new_data["pca_projection"] = pca_projections.copy() for to_be_merge, unit_id in zip(units_to_merge, new_unit_ids): new_channel_ids = sparsity.unit_id_to_channel_ids[unit_id] @@ -126,7 +127,7 @@ def _merge_extension_data( ) num_chans = new_projections.shape[2] new_data["pca_projection"][spike_indices, :, :num_chans] = new_projections - new_data["pca_projection"][spike_indices, :, num_chans:] = 0 + new_data["pca_projection"] = new_data["pca_projection"][:, :, :max_num_chans] else: new_data["pca_projection"] = pca_projections