Repository: SpikeInterface/spikeinterface
Pull request #1548 (2023-04-18 08:36:48)

Allow option to compute noise with std
In some cases, it can be useful to compute noise with std instead of mad.


Patch:
diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py
index 10ff38b55..7b9042b3d 100644
--- a/src/spikeinterface/core/recording_tools.py
+++ b/src/spikeinterface/core/recording_tools.py
@@ -1,3 +1,4 @@
+from typing import Literal
 import numpy as np
 
 
@@ -118,22 +119,27 @@ def get_closest_channels(recording, channel_ids=None, num_channels=None):
     return np.array(closest_channels_inds), np.array(dists)
 
 
-def get_noise_levels(recording, return_scaled=True, **random_chunk_kwargs):
+def get_noise_levels(recording: 'BaseRecording', return_scaled: bool = True, method: Literal['mad', 'std'] = "mad", **random_chunk_kwargs):
     """
     Estimate noise for each channel using MAD methods.
+    You can use standard deviation with `method='std'`
 
-    Internally it sample some chunk across segment.
+    Internally it samples some chunk across segment.
     And then, it use MAD estimator (more robust than STD)
 
     """
     random_chunks = get_random_data_chunks(
         recording, return_scaled=return_scaled, **random_chunk_kwargs
     )
-    med = np.median(random_chunks, axis=0, keepdims=True)
-    # hard-coded so that core doesn't depend on scipy
-    noise_levels = (
-        np.median(np.abs(random_chunks - med), axis=0) / 0.6744897501960817
-    )
+    
+    if method == "mad":
+        med = np.median(random_chunks, axis=0, keepdims=True)
+        # hard-coded so that core doesn't depend on scipy
+        noise_levels = (
+            np.median(np.abs(random_chunks - med), axis=0) / 0.6744897501960817
+        )
+    elif method == "std":
+        noise_levels = np.std(random_chunks, axis=0)
     return noise_levels
 
 


Test patch:
diff --git a/src/spikeinterface/core/tests/test_recording_tools.py b/src/spikeinterface/core/tests/test_recording_tools.py
index 709ab05bd..fcda4b2c6 100644
--- a/src/spikeinterface/core/tests/test_recording_tools.py
+++ b/src/spikeinterface/core/tests/test_recording_tools.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from spikeinterface.core import generate_recording
+from spikeinterface.core import NumpyRecording, generate_recording
 
 from spikeinterface.core.recording_tools import (get_random_data_chunks, get_chunk_with_margin,
                                                  get_closest_channels, get_channel_distances, get_noise_levels,
@@ -25,12 +25,22 @@ def test_get_noise_levels():
     rec = generate_recording(num_channels=2, sampling_frequency=1000., durations=[60.])
 
     noise_levels = get_noise_levels(rec, return_scaled=False)
-    print(noise_levels)
 
     rec.set_channel_gains(0.1)
     rec.set_channel_offsets(0)
     noise_levels = get_noise_levels(rec, return_scaled=True)
-    print(noise_levels)
+
+    noise_levels = get_noise_levels(rec, return_scaled=True, method="std")
+
+    # Generate a recording following a gaussian distribution to check the result of get_noise.
+    std = 6.0
+    traces = np.random.normal(loc=10.0, scale=std, size=(1_000_000, 2))
+    recording = NumpyRecording(traces, 30000)
+
+    assert np.allclose(get_noise_levels(recording, return_scaled=False), [std, std], rtol=1e-2, atol=1e-3)
+    assert np.allclose(get_noise_levels(recording, method="std", return_scaled=False), [std, std], rtol=1e-2, atol=1e-3)
+
+
 
 def test_get_chunk_with_margin():
     rec = generate_recording(num_channels=1, sampling_frequency=1000., durations=[10.])