put back extra_outputs in find_spikes_from_templates()

Author: samuelgarcia

src/spikeinterface/sortingcomponents/clustering/clustering_tools.py

@@ -602,21 +602,11 @@ def detect_mixtures(templates, method_kwargs={}, job_kwargs={}, tmp_folder=None,
 
             sub_recording = recording.frame_slice(t_start, t_stop)
             local_params.update({"ignore_inds": ignore_inds + [i]})
-            spikes, computed = find_spikes_from_templates(
+            
+            spikes, more_outputs = find_spikes_from_templates(
                 sub_recording, method="circus-omp-svd", method_kwargs=local_params, extra_outputs=True, **job_kwargs
             )
-            local_params.update(
-                {
-                    "overlaps": computed["overlaps"],
-                    "normed_templates": computed["normed_templates"],
-                    "norms": computed["norms"],
-                    "temporal": computed["temporal"],
-                    "spatial": computed["spatial"],
-                    "singular": computed["singular"],
-                    "units_overlaps": computed["units_overlaps"],
-                    "unit_overlaps_indices": computed["unit_overlaps_indices"],
-                }
-            )
+            local_params["precomputed"] = more_outputs
             valid = (spikes["sample_index"] >= 0) * (spikes["sample_index"] < duration + 2 * margin)
 
             if np.sum(valid) > 0:

src/spikeinterface/sortingcomponents/matching/base.py

@@ -28,4 +28,8 @@ def get_trace_margin(self):
         raise NotImplementedError  
 
     def compute(self, traces, start_frame, end_frame, segment_index, max_margin, *args):
-        raise NotImplementedError    
+        raise NotImplementedError
+    
+    def get_extra_outputs(self):
+        # can be overwritten if need to ouput some variables with a dict
+        return None

src/spikeinterface/sortingcomponents/matching/circus.py

@@ -120,6 +120,17 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
         of template temporal width)
     -----
     """
+
+    _more_output_keys = [
+                "norms",
+                "temporal",
+                "spatial",
+                "singular",
+                "units_overlaps",
+                "unit_overlaps_indices",
+                "normed_templates",
+            ]
+
     def __init__(self, recording, return_output=True, parents=None,
         templates=None,
         amplitudes=[0.6, np.inf],
@@ -130,6 +141,7 @@ def __init__(self, recording, return_output=True, parents=None,
         rank=5,
         ignore_inds=[],
         vicinity=3,
+        precomputed=None,
         ):
 
         BaseTemplateMatching.__init__(self, recording, templates, return_output=True, parents=None)
@@ -162,7 +174,12 @@ def __init__(self, recording, return_output=True, parents=None,
         #         "unit_overlaps_indices",
         #     ]:
         #         assert d[key] is not None, "If templates are provided, %d should also be there" % key
-        self._prepare_templates()
+        if precomputed is None:
+            self._prepare_templates()
+        else:
+            for key in self._more_output_keys:
+                assert precomputed[key] is not None, "If templates are provided, %d should also be there" % key
+                setattr(self, key, precomputed[key])
 
 
         self.ignore_inds = np.array(ignore_inds)
@@ -245,6 +262,14 @@ def _prepare_templates(self):
         self.temporal = np.moveaxis(self.temporal, [0, 1, 2], [1, 2, 0])
         self.singular = self.singular.T[:, :, np.newaxis]
 
+    def get_extra_outputs(self):
+        output = {}
+        for key in self._more_output_keys:
+            output[key] = getattr(self, key)
+        return output
+
+
+
 
     def get_trace_margin(self):
         return self.margin    

src/spikeinterface/sortingcomponents/matching/main.py

@@ -25,7 +25,7 @@ def find_spikes_from_templates(
     method_kwargs : dict, optional
         Keyword arguments for the chosen method
     extra_outputs : bool
-        If True then method_kwargs is also returned
+        If True then a dict is also returned is also returned
     **job_kwargs : dict
         Parameters for ChunkRecordingExecutor
     verbose : Bool, default: False
@@ -35,9 +35,8 @@ def find_spikes_from_templates(
     -------
     spikes : ndarray
         Spikes found from templates.
-    method_kwargs:
+    outputs:
         Optionaly returns for debug purpose.
-
     """
     from .method_list import matching_methods
 
@@ -58,9 +57,8 @@ def find_spikes_from_templates(
         squeeze_output=True,
     )
     if extra_outputs:
-        # TODO deprecated extra_outputs
-        method_kwargs = {}
-        return spikes, method_kwargs
+        outputs = node0.get_extra_outputs()
+        return spikes, outputs
     else:
         return spikes
 

src/spikeinterface/sortingcomponents/tests/test_template_matching.py

@@ -53,7 +53,10 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
     # }
 
     method_kwargs.update(method_kwargs_all)
-    spikes = find_spikes_from_templates(recording, method=method, method_kwargs=method_kwargs, **job_kwargs)
+    spikes, info = find_spikes_from_templates(recording, method=method, 
+                                              method_kwargs=method_kwargs, extra_outputs=True, **job_kwargs)
+
+    # print(info)
 
     # DEBUG = True
 
@@ -66,7 +69,7 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
 
     #     gt_sorting = sorting_analyzer.sorting
 
-    #     sorting = NumpySorting.from_times_labels(spikes["sample_index"], spikes["cluster_index"], sampling_frequency)
+    #     sorting = NumpySorting.from_times_labels(spikes["sample_index"], spikes["cluster_index"], recording.sampling_frequency)
 
     #     ##metrics = si.compute_quality_metrics(sorting_analyzer, metric_names=["snr"])
 
@@ -81,7 +84,7 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
     sorting_analyzer = get_sorting_analyzer()
     # method = "naive"
     # method = "tdc-peeler"
-    method =  "circus"
-    # method = "circus-omp-svd"
+    # method =  "circus"
+    method = "circus-omp-svd"
     # method = "wobble"
     test_find_spikes_from_templates(method, sorting_analyzer)