Skip to content

Speed up soft merges with vectorized spike remapping - #4848

Open
JESUSROYETH wants to merge 1 commit into
SpikeInterface:mainfrom
JESUSROYETH:radar/perf-vectorized-soft-merge-remap
Open

JESUSROYETH wants to merge 1 commit into
SpikeInterface:mainfrom
JESUSROYETH:radar/perf-vectorized-soft-merge-remap

Conversation

@JESUSROYETH

Copy link
Copy Markdown
Contributor

apply_merges_to_sorting() currently builds absolute spike indices for every unit and segment, then loops over them to rewrite unit_index. It does this even when censor_ms=None, but that path only needs a direct old-to-new unit mapping.

This PR replaces the per-unit index building with a compact lookup table applied to the spike vector in one NumPy operation. The old index-building path stays for censored merges, where it is still needed.

The public signature and return_extra=True behavior are unchanged. The keep_mask docstring is corrected too: it is only returned when return_extra=True, and it is always a boolean mask (all True when censor_ms is None). It was never actually None on that path, so the old docstring was simply wrong.

I measured the complete public SortingAnalyzer.merge_units() call, including propagation of eight loaded extensions, on a one-hour generated analyzer with 300 units, 384 channels and 5,397,781 spikes. Five paired runs after one warm-up ran on a fresh GCP c2-standard-8:

Baseline median (range) Candidate median (range) Change
Complete soft merge 1409.96 ms (1395.17-1430.00) 1192.22 ms (1185.14-1214.27) -15.44% (217.75 ms)

Every pair improved by 12.97-15.91%, and every baseline/candidate run had the same SHA-256 across the merged spike vector and all deterministic extension outputs. The public 10-second MEARec fixture gave identical outputs too, and a 200-case randomized sweep covered multi-segment, multi-group and censored merges.

A new focused test also covers the return_extra=True, censor_ms=None combination, so the unchanged keep_mask contract for that branch now has a test behind it, not just the docstring. The full sorting-tools test file, plus the memory, binary_folder and zarr analyzer cases, all pass (34 tests). Black and the repository style checks are clean.

Related to #4310.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant