From 5e025057150963dee5179837d159e5a6ce620381 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Tue, 3 Nov 2026 10:33:19 -0600 Subject: [PATCH 1/3] Shift subsets of channels at a time in phase_shift to reduce peak memory --- .../preprocessing/phase_shift.py | 33 ++++++++++++------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index a6beb78e1b..89060119c9 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -6,6 +6,7 @@ from spikeinterface.core.time_series_tools import _taper_margins from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment +from .filter import NUM_CHANNELS_PER_BLOCK class PhaseShiftRecording(BasePreprocessor): @@ -87,25 +88,35 @@ def get_traces(self, start_frame, end_frame, channel_indices): if channel_indices is None: channel_indices = slice(None) - # The in-place taper requires an owned buffer, including for float64 parents. traces_chunk, left_margin, right_margin = get_chunk_with_margin( self.parent_recording_segment, start_frame, end_frame, channel_indices, self.margin, - dtype=self.tmp_dtype, add_zeros=True, - copy=True, ) - _taper_margins(traces_chunk, self._margin_window) - traces_shift = apply_frequency_shift(traces_chunk, self.sample_shifts[channel_indices], axis=0) - - traces_shift = traces_shift[left_margin:-right_margin, :] - if self.tmp_dtype is not None: - if np.issubdtype(self.dtype, np.integer): - traces_shift = traces_shift.round() - traces_shift = traces_shift.astype(self.dtype) + sample_shifts = self.sample_shifts[channel_indices] + work_dtype = self.tmp_dtype if self.tmp_dtype is not None else traces_chunk.dtype + + # Channels are shifted independently, so shifting in blocks gives the same result with a + # much lower peak memory (the float64 copy and its complex FFT scale with the block) + num_samples = traces_chunk.shape[0] - left_margin - right_margin + num_channels = traces_chunk.shape[1] + traces_shift = np.empty((num_samples, num_channels), dtype=self.dtype) + + for channel_index in range(0, num_channels, NUM_CHANNELS_PER_BLOCK): + channel_slice = slice(channel_index, channel_index + NUM_CHANNELS_PER_BLOCK) + # astype copies, so the in-place taper does not touch the parent + block = traces_chunk[:, channel_slice].astype(work_dtype) + _taper_margins(block, self._margin_window) + block_shift = apply_frequency_shift(block, sample_shifts[channel_slice], axis=0) + block_shift = block_shift[left_margin:-right_margin, :] + + if self.tmp_dtype is not None and np.issubdtype(self.dtype, np.integer): + block_shift = block_shift.round() + + traces_shift[:, channel_slice] = block_shift return traces_shift From 665a10e00442ae844b3c7c6c480d35e54d923280 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Mon, 9 Nov 2026 11:53:58 -0600 Subject: [PATCH 2/3] alessio review --- src/spikeinterface/preprocessing/phase_shift.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 11d5acd8eb..574dacd2ce 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -140,6 +140,8 @@ def get_traces(self, start_frame, end_frame, channel_indices): return self._get_traces_fft(start_frame, end_frame, channel_indices) def _get_traces_fft(self, start_frame, end_frame, channel_indices): + from scipy.fft import irfft, next_fast_len, rfft + traces_chunk, left_margin, right_margin = get_chunk_with_margin( self.parent_recording_segment, start_frame, @@ -156,14 +158,24 @@ def _get_traces_fft(self, start_frame, end_frame, channel_indices): num_channels = traces_chunk.shape[1] traces_shift = np.empty((num_samples, num_channels), dtype=self.dtype) + # Neuropixels channels reuse a small set of ADC sampling delays, so each rotation is computed once per + # chunk and every block looks up the columns of its channels + fft_length = next_fast_len(traces_chunk.shape[0], real=True) + angular_frequencies = 2 * np.pi * np.fft.rfftfreq(fft_length) + unique_shifts, shift_indices = np.unique(sample_shifts, return_inverse=True) + unique_rotations = np.exp(-1j * (angular_frequencies[:, np.newaxis] * unique_shifts[np.newaxis, :])) + for channel_index in range(0, num_channels, NUM_CHANNELS_PER_BLOCK): channel_slice = slice(channel_index, channel_index + NUM_CHANNELS_PER_BLOCK) # The shift computes in float64 for every parent and output dtype. astype copies, so the # in-place taper does not touch the parent block = traces_chunk[:, channel_slice].astype(np.float64) _taper_margins(block, self._margin_window) - block_shift = apply_frequency_shift(block, sample_shifts[channel_slice], axis=0) - traces_shift[:, channel_slice] = _cast_owned_traces(block_shift[left_margin:-right_margin, :], self.dtype) + spectrum = rfft(block, n=fft_length, axis=0, overwrite_x=True) + spectrum *= unique_rotations[:, shift_indices[channel_slice]] + block_shift = irfft(spectrum, n=fft_length, axis=0, overwrite_x=True) + block_shift = block_shift[left_margin : left_margin + num_samples, :] + traces_shift[:, channel_slice] = _cast_owned_traces(block_shift, self.dtype) return traces_shift From 32ae6b5653700cb324bcdac0ac8e5fa68eda7dbe Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Wed, 11 Nov 2026 09:47:35 -0600 Subject: [PATCH 3/3] Clarify block casting before phase-shift tapering --- src/spikeinterface/preprocessing/phase_shift.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 574dacd2ce..a9a20aaecb 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -167,8 +167,9 @@ def _get_traces_fft(self, start_frame, end_frame, channel_indices): for channel_index in range(0, num_channels, NUM_CHANNELS_PER_BLOCK): channel_slice = slice(channel_index, channel_index + NUM_CHANNELS_PER_BLOCK) - # The shift computes in float64 for every parent and output dtype. astype copies, so the - # in-place taper does not touch the parent + # Tapering needs floating point before the FFT to preserve fractional values. + # Keep float64 precision, but cast only this block to limit peak memory. + # astype copies, so tapering in place does not modify the parent. block = traces_chunk[:, channel_slice].astype(np.float64) _taper_margins(block, self._margin_window) spectrum = rfft(block, n=fft_length, axis=0, overwrite_x=True)