Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Drift Correction issues with multi-shank probe #686

Open
Selmaan opened this issue May 7, 2024 · 27 comments
Open

Drift Correction issues with multi-shank probe #686

Selmaan opened this issue May 7, 2024 · 27 comments
Assignees
@marius10p

Comments

@Selmaan
Copy link

Selmaan commented May 7, 2024

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
@jacobpennington
Copy link
Collaborator

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
Copy link
Collaborator

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

@Selmaan
Copy link
Author

Selmaan commented May 10, 2024 •
edited
Loading

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...

@marius10p
Copy link
Contributor

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.

@Selmaan
Copy link
Author

Selmaan commented May 10, 2024 via email

@Selmaan
Copy link
Author

Selmaan commented May 10, 2024

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

@jacobpennington
Copy link
Collaborator

@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
Copy link
Collaborator

@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
Copy link
Collaborator

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.

@Selmaan
Copy link
Author

Selmaan commented Jun 3, 2024 via email

@jacobpennington
Copy link
Collaborator 

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`.

@Selmaan
Copy link
Author

Selmaan commented Jun 4, 2024

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

@jacobpennington
Copy link
Collaborator

Smoothing is applied prior to generating those plots.

@Selmaan
Copy link
Author

Selmaan commented Jun 4, 2024

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

@marius10p
Copy link
Contributor

@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.

@Selmaan
Copy link
Author

Selmaan commented Jun 6, 2024

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

@marius10p
Copy link
Contributor

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?

@Selmaan
Copy link
Author

Selmaan commented Jun 7, 2024

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!

@jacobpennington
Copy link
Collaborator

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).

@Selmaan
Copy link
Author

Selmaan commented Jul 1, 2024 via email

@jacobpennington
Copy link
Collaborator

jacobpennington commented Jul 1, 2024 •
edited
Loading

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

@Selmaan
Copy link
Author

Selmaan commented Jul 1, 2024 via email

@jacobpennington
Copy link
Collaborator

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
Copy link
Collaborator

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.

@Selmaan
Copy link
Author

Selmaan commented Jul 3, 2024 via email

@cgallimore25
Copy link

" 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

@marius10p
Copy link
Contributor

Sure, here's two that I know of:

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

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

@marius10p marius10p

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
4 participants
@marius10p @Selmaan @jacobpennington @cgallimore25