kilosort-matching in si

Author: samuelgarcia

src/spikeinterface/sortingcomponents/matching/method_list.py

@@ -12,3 +12,11 @@
     "circus-omp-svd": CircusOMPSVDPeeler,
     "wobble": WobbleMatch,
 }
+
+try:
+    # kilosort licence (GPL 3) is forcing use to make an external package
+    from spikeinterface_kilosort_reshape import KiloSortMatching
+    matching_methods["kilosort-matching"] = KiloSortMatching
+except ImportError:
+    pass
+

src/spikeinterface/sortingcomponents/tests/test_template_matching.py

@@ -45,8 +45,35 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
         "templates": templates,
     }
     method_kwargs = {}
-    if method in ("naive", "tdc-peeler", "circus", "tdc-peeler2"):
+    if method in ("naive", "tdc-peeler", "circus",):
         method_kwargs["noise_levels"] = noise_levels
+    
+    if method == "kilosort-matching":
+        from spikeinterface.sortingcomponents.peak_selection import select_peaks
+        from spikeinterface.sortingcomponents.tools import extract_waveform_at_max_channel
+        from spikeinterface.sortingcomponents.peak_detection import detect_peaks
+
+        peaks = detect_peaks(sorting_analyzer.recording, method="locally_exclusive", skip_after_n_peaks=5000)
+        few_wfs = extract_waveform_at_max_channel(sorting_analyzer.recording, peaks, ms_before=1, ms_after=2)
+
+        wfs = few_wfs[:, :, 0]
+        import numpy as np
+
+        n_components = 5
+        from sklearn.cluster import KMeans
+
+        wfs /= np.linalg.norm(wfs, axis=1)[:, None]
+        model = KMeans(n_clusters=n_components, n_init=10).fit(wfs)
+        temporal_components = model.cluster_centers_
+        temporal_components = temporal_components / np.linalg.norm(temporal_components[:, None])
+        temporal_components = temporal_components.astype(np.float32)
+        from sklearn.decomposition import TruncatedSVD
+
+        model = TruncatedSVD(n_components=n_components).fit(wfs)
+        spatial_components = model.components_.astype(np.float32)
+        method_kwargs["spatial_components"] = spatial_components
+        method_kwargs["temporal_components"] = temporal_components
+
 
     # method_kwargs["wobble"] = {
     #     "templates": waveform_extractor.get_all_templates(),
@@ -91,5 +118,7 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
     # method = "tdc-peeler"
     # method =  "circus"
     # method = "circus-omp-svd"
-    method = "wobble"
+    # method = "wobble"
+    method = "kilosort-matching"
+
     test_find_spikes_from_templates(method, sorting_analyzer)