Make soft merges/split faster

[chrishalcrow]

I've been implementing "live" curations in si-gui (SpikeInterface/spikeinterface-gui#224), and to make it more usable it would be great to speed up soft merges/splits. So this is an issue to try and direct some optimisation work! Just very quick and dirty: I benchmarked a single two-unit merge on a typical 1-hour NP2 recording. Here are the times to re-compute each extension, with metrics further split up. It's not all extensions, just the ones I use in gui.

Times for each extension re-compute in seconds:

template_similarity,3.515722459065730
firing_rate,1.814500082982700
synchrony,1.3482307499507400
rp_violation,0.54650704190135
waveforms,0.1620308329584080
amplitude_cutoff,0.10835533309727900
templates,0.1060247499262910
correlograms,0.061571332975290700
amplitude_median,0.049069040920585400
spike_locations,0.02714691695291550
snr,0.016673249891027800
num_spikes,0.0160850411048159
amplitude_cv,0.012816166039556300
sliding_rp_violation,0.011563500040210800
unit_locations,0.010269250022247400
spike_amplitudes,0.004084207932464780
firing_range,0.0012063339818269000
repolarization_slope,0.000768749974668026
presence_ratio,0.0004059579223394390
recovery_slope,0.00016670895274728500
isi_violation,0.00012345798313617700
random_spikes,8.01659189164639E-05
sd_ratio,4.77919820696116E-05
half_width,1.93340238183737E-05
noise_levels,1.19999749585986E-05
peak_to_valley,1.12500274553895E-05
peak_trough_ratio,5.79189509153366E-06

(total = 7 seconds. Note: this doesn't scale linearly, two merges take much less time than 2 times one merge)

Good news: the slow ones seem reasonable to optimize! For a soft template (a linear combination of old templates) template similarity can be computed from already-computed info. Firing rate and rp_violation are slow because we re-compute for all units (which we don't need to do). synchrony should be numbafy-able.

Now a single split!

template_similarity,4.082042583962900
correlograms,2.585401250049470
firing_rate,1.890601999941280
synchrony,1.315767249907370
templates,0.44488545798230900
rp_violation,0.43966795899905300
waveforms,0.12408262502867700
amplitude_cutoff,0.12074912490788800
amplitude_median,0.04578441695775840
amplitude_cv,0.022253792034462100
snr,0.016999625018797800
num_spikes,0.011724042007699600
sliding_rp_violation,0.011248083901591600
unit_locations,0.005210084025748070
firing_range,0.000748499995097518
presence_ratio,0.0006416250253096220
repolarization_slope,0.00044879200868308500
recovery_slope,0.0002276250161230560
isi_violation,0.00018608407117426400
random_spikes,0.00010933401063084600
sd_ratio,4.31250082328916E-05
half_width,1.81249342858791E-05
peak_to_valley,1.36250164359808E-05
noise_levels,1.11659755930305E-05
spike_amplitudes,8.66595655679703E-06
peak_trough_ratio,7.12496694177389E-06
spike_locations,5.91692514717579E-06

Harder to optimize splits: we do need to re-compute correlograms for a unit (pain), we do need to recompute the similarity between the new templates (pain). Maybe the entire template_similarity computation can be optimized better? Maybe could be numba-fied or sparsified in some way?

@yger @samuelgarcia please share any ideas! I'm happy to try and compete to make things faster. Seems like a fun project for new nerds on the project (@tayheau ;) ).

Let's post here if we decide to work on something. I'll work on easy optimizations in the metrics first.

[yger]

Hi ! I've been working exactly on that too! #4308 should speed up quite a lot the computations of the quality metrics. I had no time yet to also dig into template metrics, but this is worth doing, since most of them were done before the spike_vector

[yger]

Regarding the optimization for template_similarity, I don't think we can do much more. This is already using numba, parallel, and making use of sparsification. #4305 should speed up drastically the computations of correlograms with the use of prange in numba, at the cost of a slight numerical instability. However, this is not activate by default (you would need to turn on the option fast_mode=True). Let's see how this would be propagated in the GUI, but I think this is worth considering

[yger]

I've also now made rp_violation aware of the unit_ids, such that it is only computed for the needed ones. I'll see what to do with synchrony, the last one, not so simple to optimize

[yger]

But let me know how the new PR behaves, I'll be curious to see if there is a proper speedup also in he GUI. Currently, my benchmark use case is for 1000 units and 1000s, but I guess this should also be noticeable with your way of measuring times. Thanks !

[chrishalcrow]

Thanks @yger - amazing!! I'll check it out. For template similarity, we can significantly speed up soft merges. If $t_i$ are templates and $\alpha t_1 + \beta t_2$ is the newly merged template we can use

$$ \langle \alpha t_1 + \beta t_2 , t_3 \rangle = \alpha \langle t_1, t_3 \rangle + \beta \langle t_2, t_3 \rangle $$

since all the similarities are true metrics. Right?

[yger]

not sure, I think this is more complicated because the metrics are normalized and also given the sparsity mask used for the computations (the union, which depend on every pair):

we are computing in fact something

$$ \langle t_i, t_j \rangle = \frac{\langle t_i, t_j \rangle}{|t_i| |t_j}| $$

So not sure we can sum like that. Plus, the template_similarity metrics is not aware of number of spikes (but maybe this could be changed)

[chrishalcrow]

I don't think the sparsity is a problem. I think formally the sparsity is a projection map, which you can compose a metric function with and still get a metric function.

The normalization is annoying, though! Hmmm. Lemme think...

[yger]

Actually, this would not be so much of a big deal to save the energies of the templates, such that we could multiply by the ones of the pruned templates, and then update values accordingly. I guess this is doable, no? As said, the extension simply needs to know the amount of spikes in every templates, and such energies

[chrishalcrow]

Another idea is to estimate it. The obvious hunch is that if your new template is $\alpha t_i + \beta t_j$ then the template similarity of the joint template with another template $t_k$ is $T_{new, k} = \alpha T_{ik} + \beta T_{jk}$.

I had some fun and proved that this is true, provided $|t_i - t_j|$ is small. I think that this is a reasonable assumption, since we're merging them, and we should just do this estimate for soft merges.

show_template_similarity.pdf