Speed up FFT phase shifts for Neuropixels recordings

[JESUSROYETH]

What this changes

Neuropixels channels reuse a small set of ADC sampling delays, but the current FFT path still computes one phase vector per channel. The real NP 2.0 fixture used here has 16 unique delays across 384 channels.

This computes one rotation per unique delay and gathers it back to the channel layout. The optimised path is limited to inputs with at least 128 samples, 32 channels and 4x average delay reuse. Smaller or less repetitive inputs keep the original path.

The API and numerical expression do not change.

Benchmark

BinaryRecordingExtractor -> phase_shift -> common_reference(global median) -> write_binary_recording, 384 channels, 10 seconds, n_jobs=1, one-second chunks. Five alternating runs after one warm-up:

CPU	before	after	paired speedup	peak RSS
Intel Ice Lake	9.392 s	7.742 s	17.48%	526.2 -> 479.1 MB
AMD Milan	5.681 s	4.310 s	23.88%	525.8 -> 478.9 MB
ARM Altra	7.687 s	6.205 s	19.60%	524.9 -> 477.7 MB

The five runs on each CPU produced identical output. The real Open Ephys NP2 fast-path fixture is also byte-identical.

This complements #4563 by speeding up the exact FFT method. If FIR becomes the default, this still helps users who keep the exact path. Cross-chunk caching is left for a separate change.

[JoeZiminski]

Hi @JESUSROYETH thanks for this, really appreciate the speed-up work you are doing in general on SI!

For this one, could you run this on a very long recording e.g. 1h? I am curious of how the speed up scales. The new addition although faster does make the code harder to read (and this is already quite a fiddly part of the codebase) so would be good to ensure the speedup is definitely worth it.

[JESUSROYETH]

@JoeZiminski Thanks for pointing this out

I ran it on two isolated VMs. Each run processed a logical 1-hour, 384-channel recording as 3,600 one-second chunks with n_jobs=8. The 1-hour input repeats a deterministic 10-second fixture, so this isolates the preprocessing compute and does not include disk-write time.

Across six measured paired runs after warm-up, the baseline took a median 675.7 s and this branch 575.8 s. The median paired speedup was 15.08% (range 11.41–18.25%), which saves about 100 seconds per 1-hour recording. Peak RSS also decreased by a median 9.31%. The sampled checksum across every chunk matched in all six runs; the shorter physical-file benchmark still gives byte-identical output.

So the speedup does hold at 1 hour, although with eight workers it is a little lower than the 17.5% measured with one worker on the 10-second pipeline. I think the gain is still worth the extra branch, but the readability point is fair and I can simplify the comments if you prefer ..

[JoeZiminski]

Hey @JESUSROYETH so sorry, I completely misread the diff the first time and though the top changes were to the codebase, and missed all the thoroughly documented stuff below! Will read now

[JoeZiminski]

Thanks @JESUSROYETH sorry I haven't reviewed the entire thing yet but will take a look at the tests next week. Yes the code itself is easier to follow than anticipated apologies. However I would cut back the comments, in places are a little hard to understand (e.g. Gathering regresses at 2x average reuse on small channel counts).

In general I would focus comments on additional context the reader may need to understand what is happening rather than justifying benchmarking compared to old (now unobservable code). Here IMO the comments are actually too extensive and make it more difficult to follow the code itself.

This preprocessing step is (I think exclusively?) used for neuropixels which all have high channel counts and 12 or 16 reuse. Therefore, I would just force the fast path and not worry about a slight slowdown for the (I think impossible) slow path. This would basically get the diff on production code down to a few lines + a quick 1-line comment on reuse of unique vectors for speed. LMKWYT!

[JESUSROYETH]

Thanks @JoeZiminski, that makes sense. I can simplify the production code to always use the repeated-delay path and reduce the comments to one short note. Also, feel free to edit the PR directly if you prefer another shape, I can run the validation again afterwards or change the direction based on what you think ..

[h-mayorquin]

I agree with most of the comments of @JoeZiminski. I can confirm that this is only used for Neuropixels: every extractor that sets the inter_sample_shift property (spikeglx, openephys, cbin_ibl, ibl and spikegadgets) is reading Neuropixels probes. So we can always take the fast path and drop the thresholds. In fact, maybe the docstring of phase_shift should say so.

Regarding the testing, what we should have is a test for correctness that does not depend on the method. Something like sampling a known signal at the delayed times of each channel and checking that phase_shift brings all the channels back to the signal without delay:

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)  # Neuropixels-like: few delays, many channels

    def signal(t):
        return np.sin(2 * np.pi * 2.5 * t) + np.sin(2 * np.pi * 8.5 * t)

    # each channel samples the same signal inter_sample_shift samples late
    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)  # the recording borders have edge effects
    expected = np.broadcast_to(signal(times)[:, np.newaxis], traces.shape)
    np.testing.assert_allclose(shifted_traces[interior], expected[interior], atol=1e-3)

On main the error in the interior is 3e-5 against 5e-2 without the shift, so it fails clearly if the shift is wrong. Right now we only check that chunked and non-chunked processing agree, which would pass with a wrong shift too. I think this test can replace the new ones in this PR.

[JESUSROYETH]

Thanks @h-mayorquin, updated. The repeated-delay path is now always used, the docstring clarifies this is for Neuropixels recordings, and the tests now check that the original signal is recovered.