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.

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

Introduce extractor for Kilosort's temp_wh.dat #1954

Closed

Conversation

JoeZiminski
Copy link
Collaborator

This PR is currently in-progress. The main is to load Kilosort's temp_wh.dat file for Kilosort 1-3.

For kilosort 2.5 and 3, Phy uses the temp_wh.dat file, which (as far as I can tell) always contains all available channels.

For kilosort 1 and 2, the temp_wh.dat file is not used for Phy (although waveforms, etc. have previoulsy been calculated from it prior to saving in .npy files as it is the kilosort-preprocessed data). Kilosort 1 and 2 will remove channels for inactivity prior to sorting, and saves the used channels in channel_map.npy.

The temp_wh.dat is zero-padded at the end to match size of N * NT. The attached powerpoint indicates the method kilosort 2.5 and 3 uses for batching during preprocessing. The method is slightly different for 1, 2 but as far as I can tell results in the same output.

The main thing the extractor needs to do is handle the missing channels. The recording also contains zero-padding, but I'm hoping these can just be left there and will not effect waveforms as all spiketimes should be outside of the zero-padded region.

kilosort_preprocessing_batching.pptx

self.sorter_output_path = output_path / "sorter_output"
# TODO: store opts e.g. ntb, Nbatch etc here.

params = runpy.run_path(self.sorter_output_path / "params.py")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice! I didn't know about it!

Comment on lines 32 to 45
original_recording = load_extractor(output_path / "spikeinterface_recording.json", base_folder=output_path)
original_channel_ids = original_recording.get_channel_ids()

if original_recording.has_scaled():
gain_to_uV = original_recording.get_property("gain_to_uV")[
channel_indices
] # TODO: check this assumption - does KS change the scale / offset? can check by performing no processing...
offset_to_uV = original_recording.get_property("offset_to_uV")[channel_indices]
else:
gain_to_uV = None
offset_to_uV = None

self.original_recording_num_samples = original_recording.get_num_samples()
new_channel_ids = original_channel_ids[channel_indices] # TODO: check whether this will erroneously re-order
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this should be optional. In case someone runs KS ourside of SI, this class should still be able to load it and infer the channels etc. somewhere else, no?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Like the chanMap.mat ;) (if available)!

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @alejoe91 thanks for this! sorry was at a conference last week.

That's a good point, it should not be dependent on having previously run sorting with spikeinterface. The passed filepath could be checked and if it contains spikeinterface_recording.json + sorter_output, I think it is safe to assume it is SI-run sorting. Otherwise, we can check if it has some key files (temp_wh.dat, channel_map.npy) and assume it is the sorter output.

I guess the channel information, positions can be read from channel_map.npy and channel_positions.npy. In this case, the channel IDs could just be the channel index as a string? And I don't think the uV scaling can be recovered.

Otherwise, if it is loaded from an SI recording, we can recover the uV scaling and the channel ids can be called by their proper channel ID (as currently done). Otherwise these two methods (loading with SI metadata vs. loading from kilosort output) should be equivilent.

Does that make sense / sound reasonable?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perfectly reasonable! That's exactly what I was suggesting ;)

@alejoe91 alejoe91 added the extractors Related to extractors module label Sep 4, 2023
Comment on lines 156 to 158
# if channel_map.ndim == 2: # kilosort > 2
# does kilosort > 2 store shanks differently? channel_indices = channel_map.ravel()

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JoeZiminski ,

Kilosort 2 no longer uses the shank definition (MouseLand/Kilosort#155), so this is definitely true for Kilosort2. Based on my search of the code base of Kilosort 3 it seems like kcoords (shanks) are still not being used: (https://github.com/search?q=repo%3AMouseLand%2FKilosort%20kcoords&type=code) (ie the code is for loading the kcoords and putting them in rez, but then they aren't actually used.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually not used in KS2.5 either.
MouseLand/Kilosort#262

@JoeZiminski
Copy link
Collaborator Author

JoeZiminski commented Sep 21, 2023

Moving some discussion from #1908 here. That

  1. The raw file looks much nicer in KS than the temp_wh.dat for KS2, but by default KS uses temp_wh.dat in KS2.5 and KS3.
  2. In general the waveforms look worse after whitening. Generally KS / Phy un-whitens the spikes on the fly for presentation / template analysis (see below).

In general this raises some important issues:

TLDR: I wonder if it is worth automatically unwhitening the temp_wh.dat as part of this extractor. This allows waveforms to be extracted from drift-corrected data but maintains waveform shape.

Thanks @zm711 and @alejoe91 for bringing up this very important point. I was operating under the misapprehension that somehow regions of data including spikes were not whitened (in fact, it is just that these regions are not used in estimation of the covariance matrix for whitening in KS).

I wonder if the changes between KS2 and KS2.5, KS3 are related to this issue and this discussion. I can confirm the full whitening matrix is saved for KS2 but only a scaling diagonal matrix for more recent versions. In phy, templates are un-whitened on the fly for display (as far as I can tell, although exactly when this is on / off I am not entirely sure) e.g. here. I do not think this happens for the waveforms. I think for KS2.5 and KS3 that the displayed waveforms are incorrect, because the output inv whitening matrix KS outputs is diagonal. Also the issues and this discussion above indicates that KS unwhitens the spike for computing the final template versions (as the matlab results variable rez.Wraw contain the unwhitened templates).

The primary reason for #1954 was to ensure quality metrics are computed on data that is correct with respect to the kilosort output (e.g. spike times, channels on which spikes are strongest). However, if the case of the whitened data, it does not make sense to use this for quality metrics if waveform data must be un-whitened before interpretation. In this case, I can see that may makes more sense to use the spikeinterface pre-processed data for quality metrics calculation (assuming the user has applied bandpass filtering in spikeinterface and not kilosort).

However, if drift correction is applied, then I have two concerns, 1) is main one relevant for this discussion:

1) That the waveforms extracted from the spikeinterface pre-processed data but drift-correction is applied in kilosort, then template calculation will be incorrect if the spikeinterface-preprocessed data is used. This is because all waveforms for a unit are averaged across channels, but if the spikes drifting across different channels, then the average will be incorrect. Maybe I have made an oversight though in how the templates are computed on the SI side.
2) That applying the whitening matrix calculated prior to drift correction will be incorrect after drift correction because the cross-channel covariance will be different after the kriging interpolation applied during channel shifting. This is more a discussion for another day.

This means we have a bit of a problem, if I understand correctly. The temp_wh.dat should be used for quality metric when drift-correction is used in Kilosort (which, I think, is true in the majority of cases). However, using this file natively may lead to sub-optimal results because everything is white. Thus, it may be required for this extractor to un-whiten the data prior to waveform extraction (I hope, the inv. whitening matrix can be found somewhere in KS2.5 and KS3).

@zm711
Copy link
Collaborator

zm711 commented Sep 21, 2023

@JoeZiminski

Just a few more comments to keep in mind. Using KS in SI leads to the generation of both the whitening and inverse whitening matrices being generated, but currently the export_to_phy does not provide a whitening matrix. So if someone does kilosort and then SI and then does export to phy the whitening matrix that is left in the folder (if not deleted) may not apply correctly to the data any more.

If phy is not given a whitening matrix it just uses np.eye to create a diagonal matrix of 1s. And thus it uses the same matrix to do all of its unwhitening. Since SI export_to_phy does not supply a whitening users that do use that method will not actually (maybe because the file management going from KS to SI is a bit tricky with all the delete options) be seeing the correct unwhitened data in Phy if using temp_wh.dat.

And as you said the whitening story in KS2.5 and KS3 is a bit confusing since they were trying to do some diagonal matrix (I forget the issue you linked).

Based on all this discussion it seems to me that coming from SpikeInterface it may be best to just use the recording.dat for visualization unless we can supply an appropriate whitening matrix/unwhitening matrix for Phy to use. If a user just wants to use KS inside of SpikeInterface then it is likely not a problem, but as soon as people start doing pre/post processing etc in SpikeInterface I think the temp_wh.dat will no longer be synced up.

What I don't know at all is how the drift correction fits into this. I suppose we could do the drift correction ['kilosort-like'] (I forget the actual argument) to try to approximate the correct drift, but I think that runs into your (2) above which we can save for another time.

Sorry for the stream of consciousness, but wanted to jot down my ideas (as inchoate as they are), before I forgot.

@JoeZiminski
Copy link
Collaborator Author

Hi @zm711 thanks a lot, not at all that's very useful. I have not used export_to_phy, what are the kind of cases it is used it? I always thought it was to convert non-kilosort sorting outputs for use in Phy, but did not realise you would use it if you had run KS already (because KS output is already phy-compatible). Is this because there are some postprocessing that can be done in SI and then exported back to Phy format?

I agree I think if the whitening matrix is not available, then this temp-wh.dat should not be used. I am leaning towards having the get_traces() method of this extractor apply the whitening matrix by default. That way, waveforms will be computed similar to that in Phy, although I fear the scaling is going to be a bit of a nightmare especially across kilosort versions.

@zm711
Copy link
Collaborator

zm711 commented Sep 21, 2023

I think you're right that it is allows any SI output to be manually curated since as far as I know Phy is still the de facto standard for curation. So my guess would be that the majority of users are using it because they are not using KS/SC. But with the target of modularity in spikeinterface it is possible at some point someone would only use part of the KS workflow and thus would not generate all the npy files; or that someone wants to do some automatic merges/splits or test qcmetrics and then wants to verify their results manually. So despite running KS they then need to export their post-processed/automatically curated data into a Phy format, in which case they may rely on the export_to_phy.

That being said I do not have numbers of users doing that. So maybe this is an edge case that doesn't need to be worried about. But I just want to keep it on our radar since it is a possible complication of extracting on a temp_wh that may no longer reflect data if it is heavily (post)processed in spikeinterface.

although I fear the scaling is going to be a bit of a nightmare especially across kilosort versions.

Agreed. I think that is going to be brutal. But I'm happy to test the extractor (locally) when it's ready just to see if I see anything crazy.

@JoeZiminski
Copy link
Collaborator Author

JoeZiminski commented Sep 21, 2023

Thanks @zm711 that makes sense, I think it is nice to cater for as many edge cases as possible. As I understand it in the case of KS sorting, reading the temp_wh.dat after any export_to_phy should not be a problem as export_to_phy is only concerned with operations at a higher level than the underlying preprocessed data. e.g. if new unit ids, template amplitudes or principal components are computed from the preprocessed data and written to new .npy files, that is okay an should not affect the operation of this extractor. As long as export_to_phy is not changing the preprocessed data in any way (e.g. the binary itself or scaling information such as the whiten matrix npy), then I believe there should be no conflict.

That's a good point about using part of the KS workflow. The cases I can think of are:

  1. No SI preprocessing - preprocessing is KS, drift correction in KS+ sorting in KS
    • this is best case scenario for this extractor. Because the waveforms based on the SI recording will be completely wrong but using temp_wh.dat will be correct.
  2. SI preprocessing - preprocessing is KS, drift correction in KS+ sorting in KS
    • Similar to above, the main problem been the units are determined on drift-corrected data but waveforms computed on non-drift corrected data.
  3. preprocessing in SI, skip preprocessing in KS, drift correction in KS+ sorting in KS
    • This is a cool case, in this case the temp_wh.dat is not written but KS does overwrite the data during drift correction, I assume this is done on the recording.dat. This would be nice as it is an easy way to load the drift-corrected data for computing waveforms (as can just use SI to re-load the recording.dat, that has been drift corrected by KS) . In this case, this KilosortTempWhExtractor could not be used.
  4. preprocessing in SI, skip preprocessing in KS, skip direct correction in KS + sorting in KS
    • This is nice, as temp_wh.dat is not written so the normal spikeinterface pipeline could be used without issue.
  5. Any of the above, but changing the sorter outputs using export_to_phy
    • I believe as above, as long as export_to_phy is changing only sorting results and not data preprocessing this should have no additional effect.

Agreed. I think that is going to be brutal. But I'm happy to test the extractor (locally) when it's ready just to see if I see anything crazy.

Thanks for this!

@JoeZiminski
Copy link
Collaborator Author

JoeZiminski commented Sep 21, 2023

Just to get a handle on the actually effect of calculating templates based on drift vs. noncorrected drift data I tested. First preprocessing run in spikeinterface, then KS is run with skip_kilosort_preprocessing=True but do_correction=True in kilosort2.5. This means the recording.dat is written with KS drift correction. Code for creation here, maybe there is though some option I have not considered that is also having an effect.

Then the waveforms can be calculated from the preprocessed data, or preprocessed + drift correction. Some example units from a 10 minute recording in the attached PDF. In general the spikes are sharper after drift correction, as you'd expect,

dirft_vs_nodrift.pdf

@zm711
Copy link
Collaborator

zm711 commented Sep 22, 2023

Thanks! I'll take a look at it this weekend.

@JoeZiminski
Copy link
Collaborator Author

Just a reminder to self, one benefit of this will would be to allow visualisation of the KS data prior to sorting, this might shed light on some of the strange KS indexing errors (one possibility is the input data is completely scrambled in some way by the preprocessing).

@JoeZiminski
Copy link
Collaborator Author

JoeZiminski commented Apr 16, 2024

Hey @zm711 @alejoe91 further to the discussion in #2650 I think it would make sense to close this PR. Realistically I doubt I will ever complete it as I wouldn't be confident it could be robust across KS versions. Also KS versions < 4 seem to be basically unsupported now.

Effort would probably be better placed in recommending users run preprocessing in spikeinterface and skip KS preprocessing in order to inspect preprocessing outputs. Also if really necessary it is possible to hack the KS code to skip whitening in order to assess the KS motion correction directly. But I think handling the un-whitening and various offset corrections across versions seems very difficult with relatively little payoff. Thoughts?

@zm711
Copy link
Collaborator

zm711 commented Apr 16, 2024

@JoeZiminski I think that is fair. I know I don't have the bandwidth to pick this up. So far you and rat-h have been the only people to request this so it is desired but not super desired. Not sure if @alejoe91 wants to keep this around though.

@alejoe91
Copy link
Member

Maybe what we could do is to show that the SI preprocessing and the KS one (the temp_wh.dat) are the same?

@JoeZiminski
Copy link
Collaborator Author

Maybe what we could do is to show that the SI preprocessing and the KS one (the temp_wh.dat) are the same?

This is a nice idea! I think with the existing settings it is possible to run KS through SI with / without preprocessing (filter, CMR, whiten) and drift correction independently? (e.g. so it would be possible to test for example drift correction with / without whitening).

I wonder if it is also worth having some kind of backwards-compatibility tests that test running KS through SI vs. KS directly (e.g. the KS GUI). In theory these should be identical but I saw some issues in which there was slower performance through SI than KS directly (though maybe this is related only to KS4). I guess this would also protect against any future breaking changes. Do you think this additional layer of tests is worthwhile?

@alejoe91
Copy link
Member

yes, we have the do_correction=True/False to enable/disable drift and skip_kilosort_preprocessing parameters

@alejoe91
Copy link
Member

I wonder if it is also worth having some kind of backwards-compatibility tests that test running KS through SI vs. KS directly (e.g. the KS GUI). In theory these should be identical but I saw some issues in which there was slower performance through SI than KS directly (though maybe this is related only to KS4). I guess this would also protect against any future breaking changes. Do you think this additional layer of tests is worthwhile?

I think this would be tricky and a bit too kilosort-centric

@JoeZiminski
Copy link
Collaborator Author

Thanks! I agree the implementation may be too tricky to make feasiable. Theoretically though I would argue it should be kilosort-centric as it is essentially a wrapper around kilosort and so you would expect the output be 1:1 the same as running in kilosort directly.

For example a test such as:

  1. load spikeglx file in SI, don't do any preprocessing on it in SI.
  2. pass it directly to kilosort in SI for sorting
  3. compare output with a previously saved output from running in KS directly.

I am in general tentative about such tests to a fixed standard dataset (which is what the saved output from KS would be here) because intended changes elsewhere in the codebase can inadvertantly affect the thing you are testing and it is difficult to debug. In this case though I don't think it is an issue as there is basically no processing on the SI side unrelated to running KS, except for file loading, which should never change.

@alejoe91
Copy link
Member

I think these are good examples to run, but not as CI tests. Maybe creating some GISTS/Blogs?

@JoeZiminski
Copy link
Collaborator Author

but what if a future change to the KS-running machinery in SI (either the master scripts or something in the images) leads to a bug that results in KS in SI diverging from native KS? My understanding is at present this would not be caught in the CI?

@alejoe91
Copy link
Member

but what if a future change to the KS-running machinery in SI (either the master scripts or something in the images) leads to a bug that results in KS in SI diverging from native KS? My understanding is at present this would not be caught in the CI?

That is true. On the other hand, I think that having proper unit-tests for the different steps involved is the way to go.
I would vote not to be too Kilosort-dependent :)

@JoeZiminski
Copy link
Collaborator Author

Yes I agree it makes sense to make this as modular as possible. It's something that will come up soon as I look at bit more at the motion correction, will open an issue for discussion down the line. Cheers!

@Djoels
Copy link

Djoels commented Oct 30, 2024

I'm probably going to have to go through the effort of trying to reverse engineer KS2.5/4 preprocessing, in which case the effort that was done here could be helpful. (Luckily we are not performing drift correction as our in vitro setting doesn't require this).

Would there be any reason this can not be done if whitening_mat_inv.npy and whitening_mat.npy files are present?
The extractor could yield a FileNotFoundError if necessary files are not present?

@JoeZiminski
Copy link
Collaborator Author

JoeZiminski commented Nov 7, 2024

Hey @Djoels, cool let me know how you get on this this, it would be of interest to me also. I think if you are happy with the whitened version of the data, this is relatively straight forward. If I recall correctly, the data is written as binary and padded with zeros at the end, but would just require clipping these zeros off.

On the whitening side, I am not 100% sure what is going on across versions. In earlier KS versions (1-2), the inverse whitening matrix was saved, so unwhitening the preprocessed data should not be an issue. For 2.5-3 the inverse matrix is now identity (well, identity * (1/scaleproc) (default 200)). I am not sure why it was no longer saved. At some point, the templates were saved whitened and needed to be unwhited, e.g. here. I'm not sure what, but somewhere in the pipeline something has changed and so it seems no longer necessary to fully unwhiten the templates. And the inverse whitening matrix just used for scaling these templates. You can see this change was made here. However, this means it's not possible to unwhiten the raw data for these versions, as far as I know.

So, I think if you are happy with whitened data only, all versions are good. Otherwise, you will have to change this line in the kilosort source code to make it save the full inverse whitening matrix (I would save this to a file with a different name, so it doesn't mess up other workflows e.g. phy). I'm not sure what the situation is for kilosort 4 but would be interested in finding out.

I think this should be possible to hack for all versions, but gave up doing in a general way here as there are too many cases to reliably handle, and it seems the scaling is not fully resolved. Another option is to go into the rez.mat, there are more data accessible here, but I never delved in.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
extractors Related to extractors module
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants