Implementing sparsity for ComputeTemplates (#4212)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Garcia Samuel <[email protected]>
Co-authored-by: Alessio Buccino <[email protected]>

Author: 4 people

src/spikeinterface/core/analyzer_extension_core.py

@@ -437,6 +437,7 @@ def _run(self, verbose=False, **job_kwargs):
             return_in_uV = self.sorting_analyzer.return_in_uV
 
             return_std = "std" in self.params["operators"]
+            sparsity_mask = None if self.sparsity is None else self.sparsity.mask
             output = estimate_templates_with_accumulator(
                 recording,
                 some_spikes,
@@ -445,17 +446,24 @@ def _run(self, verbose=False, **job_kwargs):
                 self.nafter,
                 return_in_uV=return_in_uV,
                 return_std=return_std,
+                sparsity_mask=sparsity_mask,
                 verbose=verbose,
                 **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
+                data = dict(average=templates, std=stds)
             else:
-                self.data["average"] = output
+                templates = output
+                data = dict(average=templates)
+
+            if self.sparsity is not None:
+                # make average and std dense again
+                for k, arr in data.items():
+                    dense_arr = self.sparsity.densify_templates(arr)
+                    data[k] = dense_arr
+            self.data.update(data)
 
     def _compute_and_append_from_waveforms(self, operators):
         if not self.sorting_analyzer.has_extension("waveforms"):

src/spikeinterface/core/sparsity.py

@@ -247,6 +247,18 @@ def sparsify_templates(self, templates_array: np.ndarray) -> np.ndarray:
 
         return sparse_templates
 
+    def densify_templates(self, templates_array: np.ndarray) -> np.ndarray:
+        assert templates_array.shape[0] == self.num_units
+
+        densified_shape = (self.num_units, templates_array.shape[1], self.num_channels)
+        dense_templates = np.zeros(shape=densified_shape, dtype=templates_array.dtype)
+        for unit_index, unit_id in enumerate(self.unit_ids):
+            sparse_template = templates_array[unit_index, ...]
+            dense_template = self.densify_waveforms(waveforms=sparse_template[np.newaxis, :, :], unit_id=unit_id)
+            dense_templates[unit_index, :, :] = dense_template
+
+        return dense_templates
+
     @classmethod
     def from_unit_id_to_channel_ids(cls, unit_id_to_channel_ids, unit_ids, channel_ids):
         """