Running preprocessing steps as one pass over each chunk #4859

Each preprocessing step currently computes a whole chunk before the next step starts. I prototyped phase_shift → bandpass → median common_reference as a single pass over small tiles of each chunk. Its output matches the steps run one after another, and it runs 1.4 times faster on 1 thread and 1.5 times faster on 8, with a third less peak RAM.

Before opening any PRs, I'd like your views on three design questions.

1 s chunks of a 384-channel Neuropixels recording read from RAM, on an i9-12900K. "FIR steps" are phase_shift(method="fir") (#4563), bandpass_filter(method="fir") (#4856) and common_reference with the median from #4854, run one after another. On 8 threads each step splits its work across the threads, as proposed in #4857.

MS/s counts channel-samples, so one probe in real time is 11.5 MS/s. Best of 7; peak RAM is above the loaded chunk.

The single pass removes the phase shift's pass and the memory traffic between steps

• As @oliche suggested on #4563, each channel's phase-shift kernel is convolved into the bandpass kernel, which grows from 503 to 518 taps. The phase shift then costs about 1% of the filter.
• The bandpass runs on tiles small enough to stay in L2, instead of streaming the whole chunk through memory once per step.

The median gains little, since compute rather than memory limits it: on 1 thread it takes 35 ms per 1 s chunk fused and 39 ms as a separate step.

On 10 s of the recording, saved in 1 s chunks, the fused output differs from the separate steps by 6.2e-6 ADC counts rms, a millionth of the signal's 6.5 counts rms, and by at most 1.3e-4. The error is float32 rounding: it is the same within 32 samples of a chunk or tile edge as elsewhere, so the seams leave no artefacts.

Proposal

A user would opt in by calling fuse on the end of a chain:

recording = phase_shift(raw, method="fir")
recording = bandpass_filter(recording, method="fir")
recording = common_reference(recording)
recording = fuse(recording)  # one FusedFIRReferenceRecording on top of raw

fuse walks back from the last step. If the steps form a chain it knows, it replaces them with one fused step on the same parent; otherwise it returns the recording unchanged. The fused step keeps the last step's dtype and metadata, so code downstream sees the same recording.

The fused step's get_traces reads the chunk with the filter's margin, then works through tiles of about 1,600 samples by a slab of channels:

1. Filter the tile with the combined phase-shift and bandpass kernel, in one FFT and one inverse.
2. Once every slab of a time block is filtered, subtract the median or average across channels from those rows.

The arithmetic is in float32, converted to the last step's dtype once, at the end.

The prototype recognises one chain by its classes: an optional phase_shift(method="fir"), then fir_filter or a highpass_filter or bandpass_filter with method="fir", then an optional global common_reference. To fuse other chains without writing a fused step for each, steps could declare how they read their input:

A step that declares nothing is never fused and works as it does now.

Threads

The fused step splits its tiles across the threads budgeted by #4857. That budget alone makes the separate FIR steps 5.1 times faster on 8 threads, and fusion adds another 1.5 times.

Questions

1. Should fuse() be called explicitly, or applied automatically when a chain matches?
2. Is converting the dtype once, at the end, acceptable? An int16 chain currently rounds to int16 after every step, so its output would change by that rounding noise, about 0.3 ADC counts rms per step.
3. Should a fused recording save as its original chain, so saved files and provenance stay the same?

The prototype is on branch proto/fused-v2, on top of #4563, #4853, #4854, #4855 and #4856.