Skip to content

Avoid dense waveform allocation when exporting PC features #4789

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JESUSROYETH wants to merge 1 commit into
base: main
Choose a base branch
from
+4 −6
Open

Avoid dense waveform allocation when exporting PC features #4789

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions src/spikeinterface/postprocessing/principal_component.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -649,16 +649,14 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):

nsamples = nbefore + nafter

# Extract all waveforms in the chunk at once
# Prepare waveform sample offsets for all spikes in the chunk
spike_times_in_chunk = spike_times[i0:i1]
# Offset spike times to be relative to the start of the traces buffer
spike_times_offset = spike_times_in_chunk - start - nbefore
spike_indices = np.arange(i0, i1)

# Build waveform array: (n_spikes, nsamples, n_channels)
# Use fancy indexing to extract all snippets at once
sample_indices = spike_times_offset[:, None] + np.arange(nsamples)[None, :] # (n_spikes, nsamples)
all_wfs = traces[sample_indices] # (n_spikes, nsamples, n_channels)
# This adds a sample axis using strides and does not allocate waveform data.
waveform_view = np.lib.stride_tricks.sliding_window_view(traces, nsamples, axis=0)

# Vectorized PCA: batch by channel across all spikes in the chunk.
# For each unique channel, find all spikes that use it (via their unit's
Expand All @@ -683,7 +681,7 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):
global_idxs = spike_indices[all_local_idxs]

# Batch waveforms for this channel: (n_spikes, nsamples)
wfs_batch = all_wfs[all_local_idxs, :, chan_ind]
wfs_batch = waveform_view[spike_times_offset[all_local_idxs], chan_ind]

if wfs_batch.size == 0:
continue
Expand Down
Loading