Correct and speed up multi-segment synchrony metrics

[JESUSROYETH]

Description

Synchrony metrics can credit the wrong units when a sorting has multiple segments. Equal sample indices from different segments get grouped together, and the code then slices the spikes as if those occurrences were contiguous. That can mark solitary units as synchronous, and it can also miss units that belong to a real synchronous event.

The spike vector is already ordered by segment and sample, so this change just detects the adjacent events and counts the unit pairs with NumPy instead. Single-segment results, duplicate-unit behaviour, and the ordering through the supported periods path stay the same, and it adds no dependency. This is the synchrony bottleneck discussed in #4310, and it follows the ordering contract from #4606.

Benchmark

I measured this through compute_synchrony_metrics, one thread, median of five fresh-process runs after warm-up:

Workload	Before	After	Speedup
20M spikes, 400 units; Xeon 8481C	5.777 s (5.772–5.954)	0.654 s (0.648–0.655)	8.84x
#4310 proxy: 5M spikes, 1,000 units, 1,000 s; Arm Altra	2.647 s	0.250 s	10.58x

Peak RSS on the 20M-spike workload dropped from 1051.7 to 315.7 MiB, and the output hashes matched in both workloads.

Validation

I ran the quality-metrics module: 32 tests pass. The patch also matches an independent oracle across 400 randomised multi-segment cases and the public MEArec fixture. Black is clean, and I tested it on Linux.

[chrishalcrow]

Hi @JESUSROYETH this is great!!! Thank you!

I had a quick play and think the new implementation is numba jit compliant with a few changes (which I've made as suggestions in the review). Would it be possible to add these and jit the function following what is done with other metrics (e.g. use numba is HAVE_NUMBA is True). This case is slightly different because the function is the same if a user has numba or not.

[alejoe91]

Thanks @JESUSROYETH

I'm curious now! What is the numba speedup????

[JESUSROYETH]

Thanks! @alejoe91 @chrishalcrow
I'm rerunning the end-to-end benchmark on the same 20M-spike workload, including both cold and warm Numba paths. I'll come back with the results and update the PR body.

[chrishalcrow]

Very neat!

[JESUSROYETH]

Interesting result @chrishalcrow @alejoe91: on the same 20M-spike flow, warm Numba was 3.3% slower (0.668 vs 0.649 s), while a fresh process was 10.4x slower due to compilation. Five-run medians, with identical output hashes. The function is already vectorized and dominated by np.unique; Numba uses sort-and-deduplicate there, while NumPy uses its faster hash path. Should I remove the JIT wrapper, or would you prefer to validate it on your workload first?

[chrishalcrow]

Interesting! I get a 2x speedup from numpy to numba on my data (macbook, 30hr 1-shank NP2 recording). @alejoe91 do you want to try on some of your real data?
I'm not worried about the compilation time, because the main issue is when doing repeated merges/splits.

[JESUSROYETH]

yeah, that makes sense @chrishalcrow. If this runs several times during merges/splits, then the compilation time is not really important. Also, the 2x on real NP2 data is more useful than my synthetic test.

[alejoe91]

Let me test on a couple of datasets on my end!

[alejoe91]

Small dataset: 142 units / ~2M spikes

compute("quality_metrics", metric_names=["synchrony"], save=False)
main: 3.12 s ± 106 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
this PR (numba): 3.52 s ± 187 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
this PR (numpy): 3.06 s ± 82.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

compute_synchrony_metrics(analyzer)
main: 170 ms ± 2.59 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
this PR (numba): 39 ms ± 1.47 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
this PR (numpy): 36.3 ms ± 458 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Medium dataset: 325 units / ~15M spikes

compute("quality_metrics", metric_names=["synchrony"], save=False)
main: 10.5 s ± 512 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
this PR (numba): 9.58 s ± 879 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
this PR (numpy): 8.45 s ± 194 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

compute_synchrony_metrics(analyzer)
main: 3.16 s ± 75.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
this PR (numba): 381 ms ± 13.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
this PR (numpy): 1.01 s ± 17.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Large dataset: 314 units / ~150M spikes

compute("quality_metrics", metric_names=["synchrony"], save=False)
main: 24.98s
this PR (numba): 15s
this PR (numpy): 17s

compute_synchrony_metrics(analyzer)
main: 13.68s
this PR (numba): 2.93s
this PR (numpy): 5.76s

I ran the synchrony metric before the timeit to get it compiled. I would say to keep numba as default?

[JESUSROYETH]

Yes, I think keeping Numba as default makes sense for the large repeated merge/split case. There is a regression on the small dataset, but the NumPy fallback is still available when Numba is not installed. Thanks for testing the different sizes