[galenlynch]

amplitude_scalings is now 13 times faster on dense Neuropixels data, with the same results.

NP 2.0 recording, 519 units and radius_um=100 sparsity, on an i9-12900K. The 16-job times are from 60 s holding 311,699 spikes; the 1-job times are from 10 s, multiplied by 6.

On this recording 91 % of spikes collide, so most of the time goes to fit_collision. For each collision it builds a design matrix, mostly zeros, and passes it to scipy.optimize.nnls. Least squares needs only the matrix's Gram matrix and its product with the waveform. When templates are zero outside their sparsity, each entry of those is a sum over the channels two templates share. The new kernel in _amplitude_scalings_numba.py computes those sums, reusing each spike's terms across its collisions, and solves the k × k non-negative fit by Lawson–Hanson. It also fits the spikes that don't collide.

This builds on @JESUSROYETH's #4759, #4764, #4786 and #4797, and reuses the unit overlap matrix from #4797.

The kernel runs when numba is installed and every template is zero outside its sparsity mask. Otherwise, as with an explicit sparsity sparser than the templates, the numpy path runs unchanged. The kernel is a plain @njit, so n_jobs alone sets the parallelism.

The tests compare the two paths spike by spike to rtol=1e-9, on synthetic chunks with non-contiguous sparsity, negative amplitudes, a zero template, spikes at both chunk edges, gains of none, float32 and float64, collision handling on and off, and duplicate spikes. A further test runs the whole extension both ways, including the fallback. On 10 s of the real recording, the two paths give identical float32 output.

[JESUSROYETH]

Hi @galenlynch, thanks for the mention and for building on the previous changes!

I checked fa7d6cb locally and on an Ampere Altra VM. On a generated 64-unit/64-channel recording, the complete amplitude_scalings call went from 2.964 s to 0.145 s, median of five runs with one worker and a warm kernel. The float32 outputs were bit-identical. This is synthetic data, not a reproduction of your NP2 benchmark, but the improvement is clear.

Two cases showed different behaviour through the public API:

• With finite templates already computed, inserting a NaN at a colliding spike makes the numpy path raise ValueError. The numba path instead returns finite values, including zero for the affected spike. It would be useful to preserve the invalid-data check here.
• With return_in_uV=False, float16 and non-native-endian int16 recordings run on the numpy path but fail when the numba path is selected. The eligibility check currently covers sparsity, but not these dtypes.

There is also a numerical boundary worth checking. For two almost-identical templates, with a centred design-matrix condition number around 2.1e6, the numpy fit gives approximately [1.29945, 0.70055], while the kernel gives [2.0, 0.0] with a larger residual. This is a deliberately ill-conditioned numerical test, not evidence of incorrect neuron identification in real recordings. Would it make sense to fall back to the explicit fit when the small-pivot check fails?

[galenlynch]

Thanks for testing it on ARM, @JESUSROYETH. All three cases were real. I added two new commits fix them, and added regression tests for each case.

• NaN at a collision: in eb0c30f, a collision whose traces aren't finite goes to fit_collision, which raises the same ValueError as the numpy path.
• float16 and byte-swapped traces: in f041cfb, the node converts these to float32 and native byte order before the kernel, and both conversions are exact. Testing this turned up a bug on main: fit_collision built its template columns in the traces' dtype, so it truncated the templates for unscaled int16 or float16 traces. The same commit fixes it, which changes main's results for those traces only.
• Nearly identical templates: falling back was the right fix, but the kernel has to decide before the fit rather than when a solve fails. On your example the kernel stopped with one coefficient, because the second one's gradient fell below its tolerance, so no solve ever failed. The kernel now checks each collision's conditioning once, before the fit, and sends those that fail to fit_collision. Your case now matches the numpy path.

On 10 s of a Neuropixels recording, none of the 45,786 collisions fell back, the float32 output still matches the numpy path bit for bit, and the kernel time is unchanged at 170 ms per chunk.

[JESUSROYETH]

Thanks @galenlynch!

Retested the original cases on the updated version, all three are fixed and the 34 tests pass locally. LGTM from my side.

@alejoe91 @chrishalcrow, would be good to get your thoughts on this optimisation too.