Faster exact median for common_reference

[galenlynch]

common_reference(operator="median") calls np.median(traces, axis=1) on every chunk. np.median partitions each row with branches that depend on the data, and on 384 channels it mispredicts about one branch per value.

This PR adds row_median, a numba kernel that returns exactly np.median's values and dtype:

• int16 with at least 4 channels, which is what common_reference gets in the default chain: each row is counted into a histogram around the previous row's median. Neighbouring samples have nearly the same median, so most rows take one pass; a row whose median falls outside the histogram is sorted instead.
• float32 with at least 64 channels: counting passes narrow in on the middle values, and the few left are sorted. Rows holding NaN or inf take an exact path.

Other inputs, and installs without numba, keep np.median.

Speed. 1 s of 384 highpassed channels, column-major as filter steps return them, one thread, i9-12900K:

	main	this PR
median, int16	72 ms	11 ms
median, float32	101 ms	38 ms
common_reference, int16	94 ms	33 ms
common_reference, float32	106 ms	43 ms
common_reference(reference="local"), int16	640 ms	379 ms

The median of 1 s of highpassed traces against the number of channels in each median, as reference="local" and groups produce.

Mispredicts per row drop from 354 to 4 for int16, and from 392 to 47 for float32, measured with hardware counters.

One case is slower: float32 holding only whole numbers, such as unfiltered int16 cast to float32, takes 108 ms against 97.

Tests. row_median equals np.median on int8 to uint16 and float32, including ties, saturated channels, a median jumping outside the histogram, column-major and non-contiguous input, and NaN and inf. common_reference output is unchanged for int16 and float32.

Update: float rows now use np.sort, as #4836 does, and the float32 kernel described below has been removed (see this comment). The int16 kernel is unchanged.

[chrishalcrow]

These functions look very hard to understand and maintain. There's an alternative CMR PR open at the moment (#4836) , which looks simpler. It hasn't been jit-ed, but looks like it could be.

[galenlynch]

The open pr is another sorting approach to finding the median, which will have a lot of branches. The histogram approach for integers and counting approach for float should have fewer branches than any sorting approach, and therefore be faster on modern hardware. I can try to JIT the other approach and directly compare if you’d like. My guess is it would still be slower. Happy to add comments to the code here if ease of understanding the code is a problem.

[galenlynch]

Thanks for the pointer to #4836. Its approach is simpler, and for float32 it's also faster than my kernel was (29 against 39 ms), so this PR now sorts float rows the way #4836 does and the float kernel is gone (beec896). The integer kernel that remains has comments on each step (ff73ac5).

[JoeZiminski]

Thanks @galenlynch this is very cool, however my 2 cents is as a library I do not think we should be in the business of rewriting core mathematical functions. This also applies to #4836. In general it makes the library harder to maintain and the risk of bugs greatly increases. Unless the speed up is absolutely mindboggling it is safer to rely on these standard library implementations wherever we can.

[galenlynch]

Thanks for the reviews, @JoeZiminski and @chrishalcrow, both are fair points. I've cut this PR down to one integer kernel, a 55-line numba function with comments, about the size of the correlogram kernels in postprocessing/correlograms.py. For floats this PR now uses #4836's approach, which is faster than what I had. The kernel that remains makes the median 5.9 times faster on int16, the default dtype. Once the other steps of common preprocessing chains are fast, that matters. With #4563 and #4856, a chain of phase_shift → bandpass_filter → common_reference, each with its default parameters, takes 7.9 minutes of one core per hour of neuropixels 2.0 recording, and 11.0 minutes with np.median.

One core, 384 channels at 30 kHz, int16 data in RAM. "main" is upstream main; #4859 is the fused prototype.

Here is the timing of the median alone, on what common_reference receives from phase_shift → bandpass_filter: 1 s of 384 channels, on one thread of an i9-12900K.

median of each row	float32	int16
np.median (main)	100 ms	70 ms
np.sort, as in #4836	29 ms	182 ms
counting kernels, as first posted here	39 ms	12 ms

The common_reference step as a whole, which also subtracts the median and casts back to int16, goes from 91 to 33 ms, 2.8 times faster, as in the figure.

For float32, np.sort beats my counting kernel, so this PR now sorts float rows, and my float kernel, which was most of the code, is gone.

For int16, np.sort is 2.6 times slower than np.median, because numpy vectorises 16-bit sorts only on CPUs with AVX-512 (not mine). The kernel I've kept is much faster on ephys data, but not general enough to contribute to numpy. It counts each value into a histogram centred on the previous time point's median, which avoids the hard-to-predict comparisons in np.median's partition. That suits filtered traces, whose median moves little between samples. Where the median jumps, rows fall back to sorting and are slower than np.median.

The tests compare it with np.median value for value on every integer dtype it accepts, from 1 to 384 channels, including the full value range, ties, saturated channels, a median jumping outside the histogram, and reversed, strided and column-major input. Like the correlograms, it falls back to numpy, here np.median, when numba isn't installed.

I think this one kernel earns its place, since it is small, exact, on the default path, and 5.9 times faster than np.median, but you know the maintenance costs better than I do. If you'd rather not carry its maintenance then I'll close this and leave float32 to #4836.