Merge pull request #2868 from h-mayorquin/fix_compute_templates_with_…

…acumulator_sa Avoid returning std when asking for only mean/average with templates extension
SpikeInterface · May 21, 2024 · 2e7c79c · 2e7c79c
2 parents 606a1cf + bc0cbd9
commit 2e7c79c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py
@@ -330,17 +330,26 @@ def _run(self, **job_kwargs):
 
             return_scaled = self.sorting_analyzer.return_scaled
 
-            self.data["average"], self.data["std"] = estimate_templates_with_accumulator(
+            return_std = "std" in self.params["operators"]
+            output = estimate_templates_with_accumulator(
                 recording,
                 some_spikes,
                 unit_ids,
                 self.nbefore,
                 self.nafter,
                 return_scaled=return_scaled,
-                return_std=True,
+                return_std=return_std,
                 **job_kwargs,
             )
 
+            # Output of estimate_templates_with_accumulator is either (templates,) or (templates, stds)
+            if return_std:
+                templates, stds = output
+                self.data["average"] = templates
+                self.data["std"] = stds
+            else:
+                self.data["average"] = output
+
     def _compute_and_append_from_waveforms(self, operators):
         if not self.sorting_analyzer.has_extension("waveforms"):
             raise ValueError(f"Computing templates with operators {operators} needs the 'waveforms' extension")

diff --git a/src/spikeinterface/sortingcomponents/clustering/tdc.py b/src/spikeinterface/sortingcomponents/clustering/tdc.py
@@ -11,7 +11,6 @@
     get_noise_levels,
     NumpySorting,
     get_channel_distances,
-    estimate_templates_with_accumulator,
     Templates,
     compute_sparsity,
     get_global_tmp_folder,