Improve stability of velocity fits in template metrics

[ecobost]

In template metrics, velocity_above and velocity_below estimate how fast an spike moves along the axon/dendrites by fitting a line where x is distance (template_channel_location - location of soma) and y is time (position of the peak at each template_channel compared to the soma expressed in msecs). A small slope means the spike takes a long time to move through the probe (so mm/msecs), a higher slope means the spike moves fast. Currently when a peak position is the same (or very close) along the template channels, VelocityFits() tries to fit a straight vertical line; which gives an infinite slope (in practice, the fitting fails with NaN because X is ill-conditioned). This is pretty common.

This PR switches the fitting to regress peak_ms (in y) onto channel_distances (in x) and then take 1/slope to obtain the same velocities as before. This is more stable and also justified in the original source "Because the time difference between the trough of adjacent sites could be 0, to avoid infinite numbers, we calculated the inverse of velocity (ms/mm) instead by fitting a regression line to the time of waveform trough at different sites against the distance of the sites relative to soma" (Jia et al, 2020). I also centered the data before fitting the regressor. Centering helps stability and avoids having to learn an intercept which gives the model a bit more robustness.

Here's the velocities calculated with the current method and the new method (proposed in this PR):

Predicted velocities are consistent for lower velocities (center of the plot) and more stable at higher velocities. It is also able to make predictions for a lot more cases (this is a 3h neuropixels recording and the proportions of NaNs drops from 0.37 to 0.26). I also manually checked some results (where there were nans before) and they look sensible.

Overall, I think it's just a sensible change for added stability and better predictions.

[chrishalcrow]

This looks great, I'll have a close look tomorrow :)

[chrishalcrow]

Looks great!

Confirmed that it gives similar results on real data, is fast, and gives fewer nan results. Win, win win. Made a tiny comment which would marginally make the error messages less annoying.

@ecobost out of interest, since you've been looking at this score: I've only really got NP2 data. Do you know if this algorithm generalizes to large square array probes? Or is it designed only for long ones?

[ecobost]

Hi @chrishalcrow, thanks for having a look at this

I found out that sklearn's TheilSenRegressor requires fit_intercept=True (otherwise it fits lines without intercept y=wx with a single datapoint rather than two datapoints, pretty odd behavior and not what we want as the entire dataset might have intercept 0 but any particular pair of datapoints does not) so I changed that back in this commit 0630977

I then realized that because this regressor is designed to work on multi-dimensional input, it actually returns an estimate of the median (the spatial median across the parameter space) rather than the median slope (closed-form solution when num_features=1, our case); that is inexact and more time-consuming (as it has to solve a small optimization problem). So I just implemented the algorithm myself (it is pretty straightforward and hopefully easy to understand), it returns almost the same results as sklearn but it's 5-7x faster (and exact).

Overall, this should produce better more stable fits. Here's the same plot as above (with velocity computed as in the current code in x axis vs proposed 1/inv_velocity), kind of just looks the same

re: your question. I work with NHP 1.0 probes so a unit template is essentially a single column with channels 20-microns apart, which is the canonical use case for this method. Assumptions are that the neuron runs parallel to your depth dimension (y by default, i.e., along the length of the probe for neuropixels) and that channels are nearby in x (code has a column range to limit the x breadth if need be). One could compute a 2-d propagation speed tracking how the spike moves in x and y (think of having a 2-d heatmap of spike delays) but it will require some extra thought

[chrishalcrow]

Hello, my results are fairly different for the new implementation (sometimes a x2 difference near velocity of 1000: figure below is velocity_above), and I went down a rabbit hole this morning...

I think your PR fixes something more fundamental than just stability. The Theil-Sen regressor is robust to outliers in y, not in x. For the velocity fits, the distances come from channel locations and are very easy to compute while the peak times are more noisy (since they are to do with peak locations on multiple templates: these can be hard to find when you're further from the signal). Since the peak times are noisier, they should be treated like the independent variable in the regression. This PR makes this change.

So: I would say this PR makes quite a big change in the metric computation but it fixes an error in the previous implementation.