Add kilosort drift maps (that work directly on kilosort output)

[JoeZiminski]

This PR adds kilosort-style drift maps calculated directly from the kilosort output, ported from Nick Steinmetz's spikes repository (with permission from Nick, and indicated in relevant docstrings).

The benefit of this functionality is that it works directly on the kilosort output and provides a nice overview of all detected spikes and drift. This is very useful for assessing the outputs of various preprocessing steps in SpikeInterface and kilosort. At present, if you want to see the effects of skipping a kilosort preprocessing step it is hard to assess the results (the kilosort preprocessed temp_wh.dat is a bit tricky to read consistency across KS versions) and there is no ground truth for real data so comparing the sorting outputs directly is not so helpful. The plot is also very pretty!

@alejoe91 @samuelgarcia this PR uses gridspec to define the size of the subplots (gs = fig.add_gridspec(1, 2, width_ratios=[1, 5])) and I couldn't find a way to do this with make_mpl_figure so I did not use it for this widget. I am sure not using make_mpl_figure is a bad idea for many architectural reasons I don't know about it, but on of off chance it's okay, I ask: a) may I not use make_mpl_figure b) if it is required, is it okay to add an option to define the relative width of subplots?

Example code to run:

import spikeinterface.full as si
import matplotlib.pyplot as plt

si.plot_kilosort_drift_map(
    "/Users/joeziminski/data/bombcelll/sorter_output",
    only_include_large_amplitude_spikes=False,
    decimate=None,
    add_histogram_plot=False,
)
plt.show()

Example plots

Original MATLAB version

Python version from this PR

TODO

• add tests

[alejoe91]

@JoeZiminski we have a very similar plotting functio: sw.plot_drift_raster_map. I think it would be better to extend that function and add the depth histograms?

Also can you remove the npy files from the PR?

[JoeZiminski]

Hi @alejoe91 thanks for this, I removed those test files!

That makes sense, I was in two minds about this as the input data formats and computations are quite closely linked to the kilosort drift map implementation and would bloat the module, and make the signature very long. I was thinking about discarding them but it would be nice to have these options so MATLAB users who are porting to SpikeInterface have the same functionality. However I agree there is loads of code duplication so its not ideal.

How about one of these two possibilities:

Have two functions DriftRasterMapWidget and KilosortSortDriftMap. The first plots raster map and optional left-hand histogram (without peaks, boundaries etc) and works with SpikeInterface-generated inputs. The latter loads directly kilosort output, coerces them to SpikeInterface peaks/peak locations, calls DriftRasterMapWidget then paints some additional features on the histogram and raster plot (peak colouring, boundary computation, drift events, that require the majority of the code and new arguments).

Alternatively, all features from the KilosortDriftMap could be ported to DriftRasterMapWidget, even if it results in quite a heavy module and longer signature. Then, KilosortDriftMap can be a very light wrapper to load data from kilosort files, coerce them to SpikeInterface format and call DriftRasterMapWidget. Would be happy with either approach.

[samuelgarcia]

I think I agree with Alessio. I would prefer to enhance the existing plotting function we already have.

Also about having plotting function that are sorters specific, this is not in main idea of spikeinterface which is globally sorter neutral. What we could have ks specific utils function that tranform the ks folder into spikeinterface objects (analyser, motion, ...) and then make plotting function on top on spikeinterface objects.

In short I would do:

ks_analyzer, motion = si.read_kilosort_folder('path/to ks_folder')
si.plot_drift_raster_map(ks_analyzer, motion)

with this in mind, the enhancement for plot_drift_raster_map would be also beneficial for all sorters.

[JoeZiminski]

Thanks @samuelgarcia, I agree this is a much nicer approach. Would you be happy for all features from the KS drift map (e.g. histogram, peak coloring, drift events) to be added to the plot_drift_raster_map?

It would be awesome to fully load the kilosort output directory to an analyzer object, however I think the work that would take is beyond what I could do at this time. Would you be happy for now with just loading the peaks? Even though it's not as nice overall, it would would be sufficient for the raster map e.g.:

peaks, peak_locations, sampling_frequency = si.load_peaks_from_kilosort_output(sorter_path)
si.plot_drift_raster_map(peaks, peak_locations, sampling_frequency=sampling_frequency, add_histogram_plot=True, ...)