Kilosort icon indicating copy to clipboard operation
Kilosort copied to clipboard

Drift Correction issues with multi-shank probe

Open Selmaan opened this issue 9 months ago • 27 comments

Describe the issue:

I am using kilosort 4 with a 2-shank 64ch probe. Most of my recordings have essentially no drift and the algorithm has been working beautifully, thanks for the work put into this! However I am concerned there is a bug with how depth is handled for multi-shank probes for the drift correction step, which screws up the drift correction module.

Below is the channel map displayed in the ks4 gui. Note that the probe goes from 0-330um depth. image image

this is the spike position across probe colored by cluster, showing that spikes are well localized on the probe and match the channel map dimensions image

However when I look at the spike map used in drift correction, I see that spike depths range twice what they should, i.e. from 0 to ~600 image

Furthermore, even though there is a clear single instability in the recording corresponding to a shift of <20um, the drift correction algorithm has a noisy and unstable estimate which vastly overcompensates, flickering btw 0 and 60ums of drift early in the recording: image

(this is using 180,000 batch size with a 30,000 sampling rate, to account for fewer spikes on a 64ch probe than a neuropixels)

Reproduce the bug:

No response

Error message:

No response

Version information:

Following install instructions, using python 3.9, pytorch-cuda 11.8, kilosort v4.0.6

Context for the issue:

No response

Experiment information:

No response

Selmaan avatar May 07 '24 18:05 Selmaan

The issue with the drift scatter plot is a bug in the plot code, I'll have that fixed soon. Thanks for catching that! As for the noisy drift estimate, I would take that as a sign that you should not use drift correction for that probe (i.e. set nblocks = 0). ~64 channels (on a single shank) is the minimum we recommend to get reliable drift estimates.

jacobpennington avatar May 07 '24 23:05 jacobpennington

Another suggestion that might help, is you could try continuing to increase the batch size until you no longer see that noisy outlier.

jacobpennington avatar May 10 '24 18:05 jacobpennington

Is there any way to enforce a smoother drift estimate (besides using massive batch sizes, I'm already using much larger batches than default)? I am working in an unusual model species and brain area where we have many more neurons per unit volume than typical for mouse cortex (usually 100-200 units per recording), and also much less electrode drift, and I know from previous experience this kind of data is often sufficient to get reliable drift estimates (at least assuming rigidity, which seems valid to me on visual examination). But occasionally there are units which essentially turn on or off throughout the recording and disrupt motion estimates. I think you can see some of this in the spike amplitude plot I shared? I often end up turning drift correction off entirely, because estimated drift is usually +-10um or so, unless occasionally I get these massive errors. But it would be fantastic to impose some kind of smoother or robustness or prior on the drift correction...

Selmaan avatar May 10 '24 22:05 Selmaan

There is already a temporal smoothing step in the drift correction, which smooths the "Z" cross-correlation curves before taking the max (which is the drift estimate). This effectively should work like a prior. We could make the smoothing constant a parameter (@jacobpennington).

However, Jacob and I discussed this case today and we think there might be more going on. It looks like there is a bit of a large physical readjustment happening around 3000 sec. There are many neurons that are either only active before or only active after the re-adjustment. This could be due to physical modifications in tissue or distortions beyond what we can fix with drift correction. In the extreme case, it would correspond to neurons dying off or getting silenced by the big physical movement.

If you increase the batch size and it still does not help, then I think my hypothesis above might be true. In that case, the next best thing you could is split data around such critical points, spike sort separately, and use one of the packages that match neurons over days to make the matchings, since that matching is likely to be highly nonlinear.

marius10p avatar May 10 '24 22:05 marius10p

Let me get back to you with a better example, I agree with your point about the plot I shared in the first post and I have better examples to demonstrate the issue I mentioned in the most recent post. But the short version is I would definitely explore adjusting the smoothing if you can expose that to the user in some way.

On Fri, May 10, 2024 at 6:31 PM Marius Pachitariu @.***> wrote:

There is already a temporal smoothing step in the drift correction, which smooths the "Z" cross-correlation curves before taking the max (which is the drift estimate). This effectively should work like a prior. We could make the smoothing constant a parameter @.*** https://github.com/jacobpennington).

However, Jacob and I discussed this case today and we think there might be more going on. It looks like there is a bit of a large physical readjustment happening around 3000 sec. There are many neurons that are either only active before or only active after the re-adjustment. This could be due to physical modifications in tissue or distortions beyond what we can fix with drift correction. In the extreme case, it would correspond to neurons dying off or getting silenced by the big physical movement.

If you increase the batch size and it still does not help, then I think my hypothesis above might be true. In that case, the next best thing you could is split data around such critical points, spike sort separately, and use one of the packages that match neurons over days to make the matchings, since that matching is likely to be highly nonlinear.

— Reply to this email directly, view it on GitHub https://github.com/MouseLand/Kilosort/issues/686#issuecomment-2105345322, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABJNQGWQLMQW3QJ42RFZHXDZBVDCTAVCNFSM6AAAAABHLRKYRWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMBVGM2DKMZSGI . You are receiving this because you authored the thread.Message ID: @.***>

Selmaan avatar May 10 '24 23:05 Selmaan

To be clear, it sounds like the original bug I mentioned is just a plotting issue and will be resolved, so we can close this issue if you'd prefer. But on the topic of smoothing out drift correction estimates, this is an example of a spike position plot for a cleaner recording. You can see some high amplitude neurons simply turn on or off, and reappear at the same location after turning off, and this is not particularly synchronized across neurons. Overall the recording is pretty stable, although there is a bit of drift particularly early in the recording. My experience is that this kind of data sometimes catastrophically fails the drift correct step, and so I usually turn drift correction off entirely, although that's clearly not optimal image

same data, slightly different plot: image

Selmaan avatar May 10 '24 23:05 Selmaan

@marius10p is sig_interp not the smoothness parameter you're referring to?

    'sig_interp': {
        'gui_name': 'sig_interp', 'type': float, 'min': 0, 'max': np.inf,
        'exclude': [0], 'default': 20, 'step': 'preprocessing',
        'description':
            """
            For drift correction, sigma for interpolation (spatial standard
            deviation). Approximate smoothness scale in units of microns.
            """
    },

jacobpennington avatar May 15 '24 14:05 jacobpennington

@Selmaan I'm going to add the smoothness parameter Marius mentioned, it's not the one I mentioned above. In the meantime, if you want to try it with your data and see if it helps you can change the values on line 139 in kilosort.datashift.py:

    # smooth the dot-product matrices across blocks, batches and across vertical offsets
    dcs = gaussian_filter(dcs, [0.5, 0.5, 0.5])

The 0.5s are the sigma for x,y, and time, respectively.

jacobpennington avatar May 21 '24 18:05 jacobpennington

This parameter has been added in the latest version as drift_smoothing. I'll close this as resolved for now, but if you run into any more problems please let us know.

jacobpennington avatar May 28 '24 18:05 jacobpennington

great, thanks! I'm just back from a break, updated kilosort and now seeing the options. Just one q, what are the units for these parameters?

On Tue, May 28, 2024 at 2:30 PM Jacob Pennington @.***> wrote:

Closed #686 https://github.com/MouseLand/Kilosort/issues/686 as completed.

— Reply to this email directly, view it on GitHub https://github.com/MouseLand/Kilosort/issues/686#event-12961718745, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABJNQGW6MTNB6N7DGBWKAIDZETEK5AVCNFSM6AAAAABHLRKYRWVHI2DSMVQWIX3LMV45UABCJFZXG5LFIV3GK3TUJZXXI2LGNFRWC5DJN5XDWMJSHE3DCNZRHA3TINI . You are receiving this because you were mentioned.Message ID: @.***>

Selmaan avatar Jun 03 '24 15:06 Selmaan

Amount of gaussian smoothing to apply to the spatiotemporal drift
estimation, for x,y,time axes in units of registration blocks
(for x,y axes) and batch size (for time axis). The x,y smoothing has
no effect for `nblocks = 1`.

jacobpennington avatar Jun 04 '24 00:06 jacobpennington

Ah I see. Just FYI, I ran this with t=5-block smoothing for the following example session with a pretty stable recording. Oddly, the resulting drift trace appears to have unsmoothed noise...is smoothing applied upstream or downstream of this plot? image image

Selmaan avatar Jun 04 '24 19:06 Selmaan

Smoothing is applied prior to generating those plots.

jacobpennington avatar Jun 04 '24 20:06 jacobpennington

ok. I'll look into this more then trying a couple different settings...my expectation was that there would be less 'spikiness' in the plots

Selmaan avatar Jun 04 '24 22:06 Selmaan

@Selmaan that is a very stable drift trace, and the apparent "noise" is just due to hitting the integer floor of the estimation process which is about half a micron for NP. You won't have drift related problems with this data (and you could even turn drift correction off). This is very different from the first example you posted.

marius10p avatar Jun 04 '24 22:06 marius10p

Yes it is, and I understand now about the discretization causing apparent "noise", so it was probably a bad place to start. Here are results for a session with noticeable drift, and a high temporal smoothing param (10). It is possibly oversmoothed and I could slightly reduce the smoothing param, but it's definitely tracking the coarse drift. There appears to be some kind of edge effect on the last bin? image image

Selmaan avatar Jun 06 '24 21:06 Selmaan

Does that correspond to batches with no spikes, or very few spikes? Do you wait for the probe to settle at the beginning of the recording?

marius10p avatar Jun 07 '24 01:06 marius10p

Sorry, I'm not sure the best way to determine that, although it doesn't seem to me like these times are missing spikes based on a very coarse judgment in the drift scatter. More generally this level of drift is unusual, I experienced it often recording from one bird for unclear reasons, but doing reasonable things (e.g. advance probe past target recording depth, retract, and leave to settle before recording) leads to much less drift in most of my recordings, like the more stable traces I've shared. Overall I'm usually able to manage drift mostly on the data collection side which is of course preferable, and usually drift correction works well. So here I was just focusing on a particularly challenging case, and also trying to ensure that the new parameter you added for smoothing drift correction is working as expected. On that last point, I'm still a bit unsure, or at least results are not matching my intuition.

For example, here is a drift map for the initial session I shared in this thread, (which we all agreed is problematic for a couple reasons, and now looks different with the correct depth plotting for multi-shank probes). image

The only thing I wanted to highlight is that I seem to be getting quite similar estimated drift traces regardless of the smoothing parameter. So for example here is the estimated drift trace when massively increasing the temporal smoothing (7s batches, and [0.5, 0.5, 60.0] drift smoothing) image and zoom in image

all that being said, the original issue of depth for multishank probes being misplotted is definitely solved, and it is rare for my data to have this level of drift issue, so I am ok closing this if you think drift correction is working fine and this is just a data issue!

Selmaan avatar Jun 07 '24 16:06 Selmaan

I'll check this on some of our high drift data as well to be sure, but FYI the documentation for the parameter was incorrect. The first axis is for a correlation we're computing (don't recommend changing this value), the second is for time (in units of batches), and the third is for the vertical axis on the probe (units of registration blocks).

jacobpennington avatar Jun 25 '24 18:06 jacobpennington

Ah this might explain why it was behaving differently than I expected! I’ll retry with corrected parameters.

On Tue, Jun 25, 2024 at 2:34 PM Jacob Pennington @.***> wrote:

I'll check this on some of our high drift data as well to be sure, but FYI the documentation for the parameter was incorrect. The first axis is for a correlation we're computing (don't recommend changing this value), the second is for time (in units of batches), and the third is for the vertical axis on the probe (units of registration blocks).

— Reply to this email directly, view it on GitHub https://github.com/MouseLand/Kilosort/issues/686#issuecomment-2189699024, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABJNQGRTHN5T45DBAHGJLNTZJGZZPAVCNFSM6AAAAABHLRKYRWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCOBZGY4TSMBSGQ . You are receiving this because you were mentioned.Message ID: @.***>

Selmaan avatar Jul 01 '24 12:07 Selmaan

This is what I got when testing on a high drift simulation for different values for the second and third (time in batches, y in registration blocks) positions. Let me know if anything is unclear.

If you want additional details, the drift_smoothing parameter is just used as the sigma argument for scipy.ndimage.gaussian_filter.

smoothing_test

jacobpennington avatar Jul 01 '24 15:07 jacobpennington

I am a little unclear actually. The registration block smoothing makes complete sense to me. But the temporal smoothing by batch is a bit confusing, it looks here like the drift traces are more coarsely discretized? I expected it to look smoother and/or to dampen fluctuations as the temporal smoothing gets very high. So maybe I'm not getting what the parameter really does?

On Mon, Jul 1, 2024 at 11:01 AM Jacob Pennington @.***> wrote:

This is what I got when testing on a high drift simulation for different values for the second and third (time in batches, y in registration blocks) positions. Let me know if anything is unclear.

smoothing_test.png (view on web) https://github.com/MouseLand/Kilosort/assets/35777376/0eb10111-b0e0-4b5d-aa08-9cb0f458c862

— Reply to this email directly, view it on GitHub https://github.com/MouseLand/Kilosort/issues/686#issuecomment-2200400573, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABJNQGVPVCHXQUBBYYJRLUTZKFVLFAVCNFSM6AAAAABHLRKYRWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMBQGQYDANJXGM . You are receiving this because you were mentioned.Message ID: @.***>

Selmaan avatar Jul 01 '24 15:07 Selmaan

You're right, that is counter-intuitive. Sorry, I was thinking about it backwards. I did some more testing and zoomed in a bit for a single block. It looks like increasing the time smoothing does smooth out the trace as expected, but also results in outliers for some points which is why those zoomed out traces look less smooth overall. I'll look into this some more to figure out why that's happening. image

jacobpennington avatar Jul 02 '24 17:07 jacobpennington

Okay, we did identify a bug related to this that we'll have to work on. Thanks again for your patience and for helping point that out.

jacobpennington avatar Jul 03 '24 18:07 jacobpennington

no problem, thanks for working on this feature!

On Wed, Jul 3, 2024 at 2:06 PM Jacob Pennington @.***> wrote:

Okay, we did identify a bug related to this that we'll have to work on. Thanks again for your patience and for helping point that out.

— Reply to this email directly, view it on GitHub https://github.com/MouseLand/Kilosort/issues/686#issuecomment-2206917800, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABJNQGSHV5S66PELNGGSLF3ZKQ4RHAVCNFSM6AAAAABHLRKYRWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMBWHEYTOOBQGA . You are receiving this because you were mentioned.Message ID: @.***>

Selmaan avatar Jul 03 '24 19:07 Selmaan

" In that case, the next best thing you could is split data around such critical points, spike sort separately, and use one of the packages that match neurons over days to make the matchings, since that matching is likely to be highly nonlinear."

@marius10p I'm finding myself in a situation where it may be best for me to sort 2 experimental runs separately, as they were done on a low channel count probe (16 channels, 50um inter-contact distance) and separated by a 30-min wait period (pre- and post- drug application). These multielectrode specs make drift correction inaccurate. Would you mind pointing me to one of these packages you mention that matches neurons over days?

I'm trying to brainstorm a good way to match the same neuron across the 2 recordings which has drifted 1 channel above/below its original position in the baseline recording. It helps that in my recordings, individual neurons generally drift no more than 1 channel (making knowledge of its channel index very helpful), though a quantitative way would be best. Thanks

cgallimore25 avatar Jul 25 '24 21:07 cgallimore25

Sure, here's two that I know of:

https://github.com/EnnyvanBeest/UnitMatch https://github.com/janelia-TDHarrisLab/Yuan-Neuron_Tracking

marius10p avatar Jul 25 '24 22:07 marius10p