Repository: SpikeInterface/spikeinterface
Pull request #249 (2021-09-01 14:41:19)

Waveform extractor std
@alejoe91 @samuelgarcia 
creating method to extract waveform standard deviation for every time point. 
Just like templates which are either a mean or median across different spikes, this method finds the standard deviation across the spikes for every time point. 
I think the get_templates method will work for mean.


Patch:
diff --git a/spikeinterface/core/waveform_extractor.py b/spikeinterface/core/waveform_extractor.py
index 14417172f..4b75a05f3 100644
--- a/spikeinterface/core/waveform_extractor.py
+++ b/spikeinterface/core/waveform_extractor.py
@@ -64,8 +64,10 @@ class WaveformExtractor:
 
         # cache in memory
         self._waveforms = {}
+        self._template_std = {}
         self._template_average = {}
         self._template_median = {}
+        self._template_quantile = {}
         self._params = {}
 
         if (self.folder / 'params.json').is_file():
@@ -110,8 +112,10 @@ class WaveformExtractor:
 
     def _reset(self):
         self._waveforms = {}
+        self._template_std = {}
         self._template_average = {}
         self._template_median = {}
+        self._template_quantile = {}
         self._params = {}
 
         waveform_folder = self.folder / 'waveforms'
@@ -197,7 +201,7 @@ class WaveformExtractor:
         Returns
         -------
         wfs: np.array
-            The returned waveform (num_spikes, num_channels, num_samples)
+            The returned waveform (num_spikes, num_samples, num_channels)
         indices: np.array
             If 'with_index' is True, the spike indices corresponding to the waveforms extracted
         """
@@ -219,7 +223,7 @@ class WaveformExtractor:
         else:
             return wfs
 
-    def get_template(self, unit_id, mode='median'):
+    def get_template(self, unit_id, mode='median', quantile_value=0.5):
         """
         Return template (average waveform)
 
@@ -228,14 +232,16 @@ class WaveformExtractor:
         unit_id: int
             Unit id to retrieve waveforms for
         mode: str
-            'mean' or 'median' (default)
+            'mean', 'median' (default), 'std'(standard deviation), 'quantile'
+        quantile_value: float
+            quantile value for argument to np.quantile
 
         Returns
         -------
         template: np.array
             The returned template (num_samples, num_channels)
         """
-        assert mode in ('median', 'average')
+        assert mode in ('median', 'average', 'std', 'quantile')
         assert unit_id in self.sorting.unit_ids
 
         if mode == 'median':
@@ -254,8 +260,27 @@ class WaveformExtractor:
                 template = np.average(wfs, axis=0)
                 self._template_average[unit_id] = template
                 return template
+        elif mode == 'std':
+            if unit_id in self._template_std:
+                return self._template_std[unit_id]
+            else:
+                wfs = self.get_waveforms(unit_id)
+                template = np.std(wfs, axis=0)
+                self._template_std[unit_id] = template
+                return template
+        elif mode == 'quantile':
+            if quantile_value in self._template_quantile:
+                if unit_id in self._template_quantile[quantile_value]:
+                    return self._template_quantile[quantile_value][unit_id]
+            else:
+                wfs = self.get_waveforms(unit_id)
+                template = np.quantile(wfs, quantile_value, axis=0)
+                if quantile_value not in self._template_quantile:
+                    self._template_quantile[quantile_value] = dict()
+                self._template_quantile[quantile_value][unit_id] = template
+                return template
 
-    def get_all_templates(self, unit_ids=None, mode='median'):
+    def get_all_templates(self, unit_ids=None, mode='median', quantile_value=0.5):
         """
         Return several templates (average waveform)
 
@@ -264,8 +289,9 @@ class WaveformExtractor:
         unit_ids: list or None
             Unit ids to retrieve waveforms for
         mode: str
-            'mean' or 'median' (default)
-
+            'mean' or 'median' (default), 'std', 'quantile'
+        quantile_value: float
+            quantile value as argument to np.quantile
         Returns
         -------
         templates: np.array
@@ -278,7 +304,7 @@ class WaveformExtractor:
         dtype = self._params['dtype']
         templates = np.zeros((len(unit_ids), self.nsamples, num_chans), dtype=dtype)
         for i, unit_id in enumerate(unit_ids):
-            templates[i, :, :] = self.get_template(unit_id, mode=mode)
+            templates[i, :, :] = self.get_template(unit_id, mode=mode, quantile_value=quantile_value)
         return templates
 
     def sample_spikes(self):


Test patch:
diff --git a/spikeinterface/core/tests/test_waveform_extractor.py b/spikeinterface/core/tests/test_waveform_extractor.py
index 04b4cf673..bd9b8b05f 100644
--- a/spikeinterface/core/tests/test_waveform_extractor.py
+++ b/spikeinterface/core/tests/test_waveform_extractor.py
@@ -64,6 +64,15 @@ def test_WaveformExtractor():
     templates = we.get_all_templates()
     assert templates.shape == (5, 210, 2)
 
+    wf_std = we.get_template(0, mode='std')
+    assert wf_std.shape == (210, 2)
+    wfs_std = we.get_all_templates(mode='std')
+    assert wfs_std.shape == (5, 210, 2)
+
+    wf_qnt = we.get_template(0, mode='quantile', quantile_value = 0.9)
+    assert wf_qnt.shape == (210, 2)
+    wf_qnt = we.get_all_templates(mode='quantile')
+    assert wf_qnt.shape == (5, 210, 2)
 
 def test_extract_waveforms():
     # 2 segments