From 382467d0eed12364c1e1e0a2390362b6535b2d77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Wed, 23 Sep 2026 09:39:05 -0300 Subject: [PATCH 1/2] Speed up FFT phase shifts for Neuropixels recordings --- .../preprocessing/phase_shift.py | 45 ++++-- .../preprocessing/tests/test_phase_shift.py | 130 ++++++++++++++++-- 2 files changed, 153 insertions(+), 22 deletions(-) diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 3000d0c143..1f0eb37279 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -145,24 +145,45 @@ def apply_frequency_shift(signal, shift_samples, axis=0): from scipy.fft import rfft, irfft signal_length = signal.shape[axis] - num_channels = shift_samples.size - fourier_signal_size = signal_length // 2 + 1 - frequency_domain_signal = rfft(signal, n=signal_length, axis=axis, overwrite_x=True) - fourier_signal_size = frequency_domain_signal.shape[0] if axis == 0: - frequency_grid = np.empty(shape=(fourier_signal_size, num_channels)) - # Note that np.fft.rfttfreq handles both even and odd signal lengths - frequency_grid[:, :] = 2 * np.pi * np.fft.rfftfreq(signal_length)[:, np.newaxis] - shifts = np.multiply(frequency_grid, shift_samples[np.newaxis, :], out=frequency_grid) + angular_frequencies = 2 * np.pi * np.fft.rfftfreq(signal_length) + # The crossover was measured at 32 and 384 channels, and at signal lengths from 4 to 384 samples. + # np.unique's overhead is fixed per call but the exp() savings it buys scale with the signal length + # (one rotation per frequency bin), so a short chunk can lose even at the channel/reuse ratio that + # wins for a realistically sized one. 128 samples clears the measured worst case (32 channels, exact + # 4x reuse) with margin; skip the fast path below that instead of adding fixed overhead to short chunks. + take_fast_path = False + if signal_length >= 128 and shift_samples.size >= 32: + unique_shifts, shift_indices = np.unique(shift_samples, return_inverse=True) + # Gathering regresses at 2x average reuse on small channel counts, but wins consistently from 4x + # across measured 32- and 384-channel inputs. Inverse indices handle skewed reuse and non-contiguous + # channel groups without a per-channel Python loop. + take_fast_path = unique_shifts.size * 4 <= shift_samples.size + + if take_fast_path: + # Neuropixels channels share a small number of ADC sampling delays (12 for NP 1.0). The original + # path always rounds the angle to a float64 buffer right after the multiply, independently of + # input dtype; downcast the angle product here (not the raw shift beforehand) so a wider shift + # dtype rounds at the same point instead of losing precision one multiply earlier. + unique_angles = (angular_frequencies[:, np.newaxis] * unique_shifts[np.newaxis, :]).astype( + np.float64, copy=False + ) + unique_rotations = np.exp(-1j * unique_angles) + rotations = unique_rotations[:, shift_indices] + # scipy.fft preserves a float32 input as complex64, while rotations above is complex128. The ufunc + # promotes the multiply to complex128 here, matching the original path's float64-precision output. + phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) + else: + frequency_grid = np.empty(shape=frequency_domain_signal.shape) + frequency_grid[:, :] = angular_frequencies[:, np.newaxis] + shifts = np.multiply(frequency_grid, shift_samples[np.newaxis, :], out=frequency_grid) + rotations = np.exp(-1j * shifts) + phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) else: raise NotImplementedError("Axis != 0 is not implemented yet") - # Rotate the signal in the frequency domain - rotations = np.exp(-1j * shifts) - phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) - # Inverse FFT to get the translated signal shifted_signal = irfft(phase_shifted_signal, n=signal_length, axis=axis, overwrite_x=True) return shifted_signal diff --git a/src/spikeinterface/preprocessing/tests/test_phase_shift.py b/src/spikeinterface/preprocessing/tests/test_phase_shift.py index 6bf7d50dd4..f0e9bd231e 100644 --- a/src/spikeinterface/preprocessing/tests/test_phase_shift.py +++ b/src/spikeinterface/preprocessing/tests/test_phase_shift.py @@ -2,9 +2,10 @@ from spikeinterface import NumpyRecording from spikeinterface.preprocessing import phase_shift +from spikeinterface.preprocessing.phase_shift import apply_frequency_shift -def create_shifted_channel(): +def create_shifted_channel(inter_sample_shift=None): duration = 5.0 sr_h = 10000.0 times_h = np.arange(0, duration, 1 / sr_h) @@ -16,15 +17,11 @@ def create_shifted_channel(): ratio = 10 sr = sr_h / ratio - times = times_h[0::ratio] - delay_sample = 4 - sig0 = sig_h[0::ratio] - sig1 = sig_h[delay_sample::ratio] - - inter_sample_shift = [0.0, delay_sample / ratio] - - traces = np.stack([sig0, sig1], axis=1) - traces *= 1000 + if inter_sample_shift is None: + inter_sample_shift = [0.0, 0.4] + channels = [sig_h[int(round(shift * ratio)) :: ratio] for shift in inter_sample_shift] + min_length = min(channel.size for channel in channels) + traces = np.stack([channel[:min_length] for channel in channels], axis=1) * 1000 return traces, sr, inter_sample_shift @@ -93,5 +90,118 @@ def test_phase_shift(): # ~ plt.show() +def test_phase_shift_repeated_delays(): + inter_sample_shift = np.tile(np.arange(8) / 10, 8) + traces, sampling_frequency, _ = create_shifted_channel(inter_sample_shift) + recording = NumpyRecording([traces], sampling_frequency) + recording.set_property("inter_sample_shift", inter_sample_shift) + + recording2 = phase_shift(recording, margin_ms=40.0) + traces2 = recording2.get_traces() + recording3 = recording2.save(format="memory", chunk_size=1000, n_jobs=1, progress_bar=False) + traces3 = recording3.get_traces() + + rms = np.sqrt(np.mean(traces**2)) + error_mean = np.sqrt(np.mean((traces2 - traces3) ** 2)) + error_max = np.sqrt(np.max((traces2 - traces3) ** 2)) + assert error_mean / rms < 0.001 + assert error_max / rms < 0.02 + + # Reading from recording2 calls apply_frequency_shift again. These non-contiguous blocks retain eight + # delays across 32 channels, so the subset also enters the gather path at the exact 4x boundary. + subset_indices = np.r_[0:16, 32:48] + subset_ids = recording2.channel_ids[subset_indices] + traces_slice = recording2.get_traces(channel_ids=subset_ids) + np.testing.assert_allclose(traces_slice, traces2[:, subset_indices], rtol=1e-12, atol=1e-12) + + +def test_apply_frequency_shift_repeated_and_unique_shifts(): + rng = np.random.default_rng(0) + boundary_shifts = np.tile(np.arange(8) / 8, 4) + skewed_shifts = np.concatenate([np.full(25, 0.1), np.arange(1, 8) / 9]) + rng.shuffle(skewed_shifts) + _, skewed_counts = np.unique(skewed_shifts, return_counts=True) + assert np.unique(boundary_shifts).size * 4 == boundary_shifts.size + assert skewed_counts.min() == 1 and skewed_counts.max() == 25 + + cases = [] + # 100/101 sit below the 128-sample fast-path eligibility floor (fallback path); 128/129 sit at/above it + # (gather path), so both branches are covered for repeated, unique, exact-4x, and skewed shifts. + for num_samples in (100, 101, 128, 129): + signal = rng.standard_normal((num_samples, 32)) + shift_cases = (np.tile(np.arange(4) / 4, 8), np.arange(32) / 32, boundary_shifts, skewed_shifts) + for shift_samples in shift_cases: + frequencies = 2 * np.pi * np.fft.rfftfreq(signal.shape[0])[:, np.newaxis] + expected = np.fft.irfft( + np.fft.rfft(signal, axis=0) * np.exp(-1j * frequencies * shift_samples[np.newaxis, :]), + n=signal.shape[0], + axis=0, + ) + cases.append((signal, shift_samples, expected)) + + for signal, shift_samples, expected in cases: + actual = apply_frequency_shift(signal.copy(), shift_samples) + np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) + + +def test_apply_frequency_shift_float32_input_matches_float64_oracle(): + rng = np.random.default_rng(0) + # 128 samples clears the fast-path signal-length floor so repeated_shifts below exercises the + # repeated-shift branch this oracle targets, not the fallback. + signal32 = rng.standard_normal((128, 32)).astype(np.float32) + repeated_shifts = np.tile(np.arange(4) / 4, 8) + unique_shifts = np.arange(32) / 32 + frequencies = 2 * np.pi * np.fft.rfftfreq(signal32.shape[0])[:, np.newaxis] + + for shift_samples in (repeated_shifts, unique_shifts): + # Independent oracle: upcast the signal to float64 before the FFT, rather than after, as the + # repeated-shift branch does. The two are not bit-identical (the float32 FFT already rounds at + # float32 precision) but must agree well past that noise floor. + expected = np.fft.irfft( + np.fft.rfft(signal32.astype(np.float64), axis=0) * np.exp(-1j * frequencies * shift_samples[np.newaxis, :]), + n=signal32.shape[0], + axis=0, + ) + actual = apply_frequency_shift(signal32.copy(), shift_samples) + assert actual.dtype == np.float64 + np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5) + + +def test_apply_frequency_shift_output_dtype_independent_of_shift_dtype(): + rng = np.random.default_rng(0) + # 128 samples clears the fast-path signal-length floor so repeated_shifts below exercises the + # repeated-shift branch this test targets, not the fallback. + signal = rng.standard_normal((128, 32)) + + repeated_shifts = np.tile(np.array([0.13, 0.27, 0.41, 0.59]), 8) + unique_shifts = np.linspace(0.013, 0.917, 32) + for shift_samples in (repeated_shifts, unique_shifts): + reference = apply_frequency_shift(signal.copy(), shift_samples) + widened = apply_frequency_shift(signal.copy(), shift_samples.astype(np.longdouble)) + assert widened.dtype == reference.dtype == np.float64 + np.testing.assert_allclose(widened, reference, rtol=1e-12, atol=1e-12) + + +def test_apply_frequency_shift_repeated_path_rounds_after_multiplying_not_before(): + # Parse directly into longdouble: casting up a float64 cannot recover precision it never had. + wide_shift = np.longdouble("1000000000000.0001") + # 128 samples clears the fast-path signal-length floor so this exercises the repeated-shift branch. + num_samples, num_channels = 128, 32 + rng = np.random.default_rng(0) + signal = rng.standard_normal((num_samples, num_channels)) + shift_samples = np.full(num_channels, wide_shift, dtype=np.longdouble) + + # Independent oracle: multiply at the shift's native precision, then round the angle to float64 + # at the same point as the original vectorized path's float64 output buffer. + angular_frequencies = 2 * np.pi * np.fft.rfftfreq(num_samples) + angles = (angular_frequencies[:, np.newaxis] * shift_samples[np.newaxis, :]).astype(np.float64) + rotations = np.exp(-1j * angles) + expected = np.fft.irfft(np.fft.rfft(signal, axis=0) * rotations, n=num_samples, axis=0) + + actual = apply_frequency_shift(signal.copy(), shift_samples) + assert actual.dtype == np.float64 + np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) + + if __name__ == "__main__": test_phase_shift() From bb6d42956f1741aaf49879bf0b600d6ef256aa36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Sun, 4 Oct 2026 21:19:13 -0300 Subject: [PATCH 2/2] Address phase-shift review feedback --- .../preprocessing/phase_shift.py | 45 ++---- .../preprocessing/tests/test_phase_shift.py | 139 +++--------------- 2 files changed, 33 insertions(+), 151 deletions(-) diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 1f0eb37279..c01a9acf59 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -9,10 +9,7 @@ class PhaseShiftRecording(BasePreprocessor): """ - This apply a phase shift to a recording to cancel the small sampling - delay across for some recording system. - - This is particularly relevant for neuropixel recording. + Apply phase shifts to compensate for the staggered ADC sampling times in Neuropixels recordings. This code is inspired from from IBL lib. https://github.com/int-brain-lab/ibllib/blob/master/ibllib/dsp/fourier.py @@ -149,38 +146,14 @@ def apply_frequency_shift(signal, shift_samples, axis=0): if axis == 0: angular_frequencies = 2 * np.pi * np.fft.rfftfreq(signal_length) - # The crossover was measured at 32 and 384 channels, and at signal lengths from 4 to 384 samples. - # np.unique's overhead is fixed per call but the exp() savings it buys scale with the signal length - # (one rotation per frequency bin), so a short chunk can lose even at the channel/reuse ratio that - # wins for a realistically sized one. 128 samples clears the measured worst case (32 channels, exact - # 4x reuse) with margin; skip the fast path below that instead of adding fixed overhead to short chunks. - take_fast_path = False - if signal_length >= 128 and shift_samples.size >= 32: - unique_shifts, shift_indices = np.unique(shift_samples, return_inverse=True) - # Gathering regresses at 2x average reuse on small channel counts, but wins consistently from 4x - # across measured 32- and 384-channel inputs. Inverse indices handle skewed reuse and non-contiguous - # channel groups without a per-channel Python loop. - take_fast_path = unique_shifts.size * 4 <= shift_samples.size - - if take_fast_path: - # Neuropixels channels share a small number of ADC sampling delays (12 for NP 1.0). The original - # path always rounds the angle to a float64 buffer right after the multiply, independently of - # input dtype; downcast the angle product here (not the raw shift beforehand) so a wider shift - # dtype rounds at the same point instead of losing precision one multiply earlier. - unique_angles = (angular_frequencies[:, np.newaxis] * unique_shifts[np.newaxis, :]).astype( - np.float64, copy=False - ) - unique_rotations = np.exp(-1j * unique_angles) - rotations = unique_rotations[:, shift_indices] - # scipy.fft preserves a float32 input as complex64, while rotations above is complex128. The ufunc - # promotes the multiply to complex128 here, matching the original path's float64-precision output. - phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) - else: - frequency_grid = np.empty(shape=frequency_domain_signal.shape) - frequency_grid[:, :] = angular_frequencies[:, np.newaxis] - shifts = np.multiply(frequency_grid, shift_samples[np.newaxis, :], out=frequency_grid) - rotations = np.exp(-1j * shifts) - phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) + # Neuropixels channels reuse a small set of ADC sampling delays, so compute each rotation once. + unique_shifts, shift_indices = np.unique(shift_samples, return_inverse=True) + unique_angles = (angular_frequencies[:, np.newaxis] * unique_shifts[np.newaxis, :]).astype( + np.float64, copy=False + ) + unique_rotations = np.exp(-1j * unique_angles) + rotations = unique_rotations[:, shift_indices] + phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) else: raise NotImplementedError("Axis != 0 is not implemented yet") diff --git a/src/spikeinterface/preprocessing/tests/test_phase_shift.py b/src/spikeinterface/preprocessing/tests/test_phase_shift.py index f0e9bd231e..23718f6ac6 100644 --- a/src/spikeinterface/preprocessing/tests/test_phase_shift.py +++ b/src/spikeinterface/preprocessing/tests/test_phase_shift.py @@ -2,10 +2,9 @@ from spikeinterface import NumpyRecording from spikeinterface.preprocessing import phase_shift -from spikeinterface.preprocessing.phase_shift import apply_frequency_shift -def create_shifted_channel(inter_sample_shift=None): +def create_shifted_channel(): duration = 5.0 sr_h = 10000.0 times_h = np.arange(0, duration, 1 / sr_h) @@ -17,11 +16,15 @@ def create_shifted_channel(inter_sample_shift=None): ratio = 10 sr = sr_h / ratio - if inter_sample_shift is None: - inter_sample_shift = [0.0, 0.4] - channels = [sig_h[int(round(shift * ratio)) :: ratio] for shift in inter_sample_shift] - min_length = min(channel.size for channel in channels) - traces = np.stack([channel[:min_length] for channel in channels], axis=1) * 1000 + times = times_h[0::ratio] + delay_sample = 4 + sig0 = sig_h[0::ratio] + sig1 = sig_h[delay_sample::ratio] + + inter_sample_shift = [0.0, delay_sample / ratio] + + traces = np.stack([sig0, sig1], axis=1) + traces *= 1000 return traces, sr, inter_sample_shift @@ -90,117 +93,23 @@ def test_phase_shift(): # ~ plt.show() -def test_phase_shift_repeated_delays(): - inter_sample_shift = np.tile(np.arange(8) / 10, 8) - traces, sampling_frequency, _ = create_shifted_channel(inter_sample_shift) +def test_phase_shift_recovers_delayed_signal(): + sampling_frequency = 1000.0 + times = np.arange(5000) / sampling_frequency + inter_sample_shift = np.tile([0.0, 0.25, 0.5, 0.75], 8) + + def signal(time): + return np.sin(2 * np.pi * 2.5 * time) + np.sin(2 * np.pi * 8.5 * time) + + traces = signal(times[:, np.newaxis] + inter_sample_shift / sampling_frequency) recording = NumpyRecording([traces], sampling_frequency) recording.set_property("inter_sample_shift", inter_sample_shift) - recording2 = phase_shift(recording, margin_ms=40.0) - traces2 = recording2.get_traces() - recording3 = recording2.save(format="memory", chunk_size=1000, n_jobs=1, progress_bar=False) - traces3 = recording3.get_traces() - - rms = np.sqrt(np.mean(traces**2)) - error_mean = np.sqrt(np.mean((traces2 - traces3) ** 2)) - error_max = np.sqrt(np.max((traces2 - traces3) ** 2)) - assert error_mean / rms < 0.001 - assert error_max / rms < 0.02 - - # Reading from recording2 calls apply_frequency_shift again. These non-contiguous blocks retain eight - # delays across 32 channels, so the subset also enters the gather path at the exact 4x boundary. - subset_indices = np.r_[0:16, 32:48] - subset_ids = recording2.channel_ids[subset_indices] - traces_slice = recording2.get_traces(channel_ids=subset_ids) - np.testing.assert_allclose(traces_slice, traces2[:, subset_indices], rtol=1e-12, atol=1e-12) - - -def test_apply_frequency_shift_repeated_and_unique_shifts(): - rng = np.random.default_rng(0) - boundary_shifts = np.tile(np.arange(8) / 8, 4) - skewed_shifts = np.concatenate([np.full(25, 0.1), np.arange(1, 8) / 9]) - rng.shuffle(skewed_shifts) - _, skewed_counts = np.unique(skewed_shifts, return_counts=True) - assert np.unique(boundary_shifts).size * 4 == boundary_shifts.size - assert skewed_counts.min() == 1 and skewed_counts.max() == 25 - - cases = [] - # 100/101 sit below the 128-sample fast-path eligibility floor (fallback path); 128/129 sit at/above it - # (gather path), so both branches are covered for repeated, unique, exact-4x, and skewed shifts. - for num_samples in (100, 101, 128, 129): - signal = rng.standard_normal((num_samples, 32)) - shift_cases = (np.tile(np.arange(4) / 4, 8), np.arange(32) / 32, boundary_shifts, skewed_shifts) - for shift_samples in shift_cases: - frequencies = 2 * np.pi * np.fft.rfftfreq(signal.shape[0])[:, np.newaxis] - expected = np.fft.irfft( - np.fft.rfft(signal, axis=0) * np.exp(-1j * frequencies * shift_samples[np.newaxis, :]), - n=signal.shape[0], - axis=0, - ) - cases.append((signal, shift_samples, expected)) - - for signal, shift_samples, expected in cases: - actual = apply_frequency_shift(signal.copy(), shift_samples) - np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) - - -def test_apply_frequency_shift_float32_input_matches_float64_oracle(): - rng = np.random.default_rng(0) - # 128 samples clears the fast-path signal-length floor so repeated_shifts below exercises the - # repeated-shift branch this oracle targets, not the fallback. - signal32 = rng.standard_normal((128, 32)).astype(np.float32) - repeated_shifts = np.tile(np.arange(4) / 4, 8) - unique_shifts = np.arange(32) / 32 - frequencies = 2 * np.pi * np.fft.rfftfreq(signal32.shape[0])[:, np.newaxis] - - for shift_samples in (repeated_shifts, unique_shifts): - # Independent oracle: upcast the signal to float64 before the FFT, rather than after, as the - # repeated-shift branch does. The two are not bit-identical (the float32 FFT already rounds at - # float32 precision) but must agree well past that noise floor. - expected = np.fft.irfft( - np.fft.rfft(signal32.astype(np.float64), axis=0) * np.exp(-1j * frequencies * shift_samples[np.newaxis, :]), - n=signal32.shape[0], - axis=0, - ) - actual = apply_frequency_shift(signal32.copy(), shift_samples) - assert actual.dtype == np.float64 - np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5) - - -def test_apply_frequency_shift_output_dtype_independent_of_shift_dtype(): - rng = np.random.default_rng(0) - # 128 samples clears the fast-path signal-length floor so repeated_shifts below exercises the - # repeated-shift branch this test targets, not the fallback. - signal = rng.standard_normal((128, 32)) - - repeated_shifts = np.tile(np.array([0.13, 0.27, 0.41, 0.59]), 8) - unique_shifts = np.linspace(0.013, 0.917, 32) - for shift_samples in (repeated_shifts, unique_shifts): - reference = apply_frequency_shift(signal.copy(), shift_samples) - widened = apply_frequency_shift(signal.copy(), shift_samples.astype(np.longdouble)) - assert widened.dtype == reference.dtype == np.float64 - np.testing.assert_allclose(widened, reference, rtol=1e-12, atol=1e-12) - - -def test_apply_frequency_shift_repeated_path_rounds_after_multiplying_not_before(): - # Parse directly into longdouble: casting up a float64 cannot recover precision it never had. - wide_shift = np.longdouble("1000000000000.0001") - # 128 samples clears the fast-path signal-length floor so this exercises the repeated-shift branch. - num_samples, num_channels = 128, 32 - rng = np.random.default_rng(0) - signal = rng.standard_normal((num_samples, num_channels)) - shift_samples = np.full(num_channels, wide_shift, dtype=np.longdouble) - - # Independent oracle: multiply at the shift's native precision, then round the angle to float64 - # at the same point as the original vectorized path's float64 output buffer. - angular_frequencies = 2 * np.pi * np.fft.rfftfreq(num_samples) - angles = (angular_frequencies[:, np.newaxis] * shift_samples[np.newaxis, :]).astype(np.float64) - rotations = np.exp(-1j * angles) - expected = np.fft.irfft(np.fft.rfft(signal, axis=0) * rotations, n=num_samples, axis=0) - - actual = apply_frequency_shift(signal.copy(), shift_samples) - assert actual.dtype == np.float64 - np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) + shifted_traces = phase_shift(recording).get_traces() + + interior = slice(100, -100) + expected = np.broadcast_to(signal(times)[:, np.newaxis], traces.shape) + np.testing.assert_allclose(shifted_traces[interior], expected[interior], atol=1e-3) if __name__ == "__main__":