Speed up FFT phase shifts for Neuropixels recordings

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

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()