diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 3000d0c143..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 @@ -145,24 +142,21 @@ 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) + # 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") - # 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..23718f6ac6 100644 --- a/src/spikeinterface/preprocessing/tests/test_phase_shift.py +++ b/src/spikeinterface/preprocessing/tests/test_phase_shift.py @@ -93,5 +93,24 @@ def test_phase_shift(): # ~ plt.show() +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) + + 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__": test_phase_shift()