Repository: SpikeInterface/spikeinterface
Pull request #1583 (2023-04-26 13:47:42)

Caching noise levels
As we just discussed, here is a proposal to cache noise levels in the recording. I think this makes a lot of sense, because otherwise, we might face some variability in these levels is estimated several times during a pipeline. However, not familiar with the whole architecture of SI. If new recordings are made from a given recordings, I am assuming this property is not inherited. Is it correct?


Patch:
diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py
index bfeb60f4f..9b300e478 100644
--- a/src/spikeinterface/core/base.py
+++ b/src/spikeinterface/core/base.py
@@ -34,6 +34,9 @@ class BaseExtractor:
     _main_annotations = []
     _main_properties = []
 
+    # these properties are skipped by default in copy_metadata
+    _skip_properties = []
+
     installed = True
     installation_mesg = ""
     is_writable = False
@@ -246,12 +249,25 @@ class BaseExtractor:
             raise Exception(f"{key} is not a property key")
 
     def copy_metadata(
-        self, other: "BaseExtractor", only_main: bool = False, ids: Union[Iterable, slice, None] = None
+        self,
+        other: "BaseExtractor",
+        only_main: bool = False,
+        ids: Union[Iterable, slice, None] = None,
+        skip_properties: Optional[Iterable[str]] = None,
     ) -> None:
         """
-        Copy annotations/properties to another extractor.
+        Copy metadata (annotations/properties) to another extractor (`other`).
 
-        If 'only main' is True, then only "main" annotations/properties one are copied.
+        Parameters
+        ----------
+        other: BaseExtractor
+            The extractor to copy the metadata to.
+        only_main: bool
+            If True, only the main annotations/properties are copied.
+        ids: list
+            List of ids to copy the metadata to. If None, all ids are copied.
+        skip_properties: list
+            List of properties to skip. Default is None.
         """
 
         if ids is None:
@@ -269,7 +285,15 @@ class BaseExtractor:
             prop_keys = self._properties.keys()
 
         other._annotations = deepcopy({k: self._annotations[k] for k in ann_keys})
+
+        # skip properties based on target "other" extractor
+        skip_properties_all = other._skip_properties
+        if skip_properties is not None:
+            skip_properties_all = skip_properties_all + skip_properties
+
         for k in prop_keys:
+            if k in skip_properties_all:
+                continue
             values = self._properties[k]
             if values is not None:
                 other.set_property(k, values[inds])
diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py
index 1bfb2e576..d4d0beb5f 100644
--- a/src/spikeinterface/core/baserecording.py
+++ b/src/spikeinterface/core/baserecording.py
@@ -25,6 +25,8 @@ class BaseRecording(BaseRecordingSnippets):
     _main_properties = ["group", "location", "gain_to_uV", "offset_to_uV"]
     _main_features = []  # recording do not handle features
 
+    _skip_properties = ["noise_level_raw", "noise_level_scaled"]
+
     def __init__(self, sampling_frequency: float, channel_ids: List, dtype):
         BaseRecordingSnippets.__init__(
             self, channel_ids=channel_ids, sampling_frequency=sampling_frequency, dtype=dtype
diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py
index 1b6c76745..865e5cc28 100644
--- a/src/spikeinterface/core/recording_tools.py
+++ b/src/spikeinterface/core/recording_tools.py
@@ -79,27 +79,6 @@ def get_random_data_chunks(
         return chunk_list
 
 
-def get_noise_levels(
-    recording: "BaseRecording", return_scaled: bool = True, method: Literal["mad", "std"] = "mad", **random_chunk_kwargs
-):
-    """
-    Estimate signal std for each channel using median absolute deviation (MAD) and std.
-
-    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)
-
-    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
-
-
 def get_channel_distances(recording):
     """
     Distance between channel pairs
@@ -147,6 +126,61 @@ 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: "BaseRecording",
+    return_scaled: bool = True,
+    method: Literal["mad", "std"] = "mad",
+    force_recompute: bool = False,
+    **random_chunk_kwargs,
+):
+    """
+    Estimate noise for each channel using MAD methods.
+    You can use standard deviation with `method='std'`
+
+    Internally it samples some chunk across segment.
+    And then, it use MAD estimator (more robust than STD)
+
+    Parameters
+    ----------
+
+    recording: BaseRecording
+        The recording extractor to get noise levels
+    return_scaled: bool
+        If True, returned noise levels are scaled to uV
+    method: str
+        'mad' or 'std'
+    force_recompute: bool
+        If True, noise levels are recomputed even if they are already stored in the recording extractor
+    random_chunk_kwargs: dict
+        Kwargs for get_random_data_chunks
+
+    Returns
+    -------
+    noise_levels: array
+        Noise levels for each channel
+    """
+
+    if return_scaled:
+        key = "noise_level_scaled"
+    else:
+        key = "noise_level_raw"
+
+    if key in recording.get_property_keys() and not force_recompute:
+        noise_levels = recording.get_property(key=key)
+    else:
+        random_chunks = get_random_data_chunks(recording, return_scaled=return_scaled, **random_chunk_kwargs)
+
+        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)
+        recording.set_property(key, noise_levels)
+
+    return noise_levels
+
+
 def get_chunk_with_margin(
     rec_segment,
     start_frame,


Test patch:
diff --git a/src/spikeinterface/core/tests/test_baserecording.py b/src/spikeinterface/core/tests/test_baserecording.py
index 89693a550..230ff40b4 100644
--- a/src/spikeinterface/core/tests/test_baserecording.py
+++ b/src/spikeinterface/core/tests/test_baserecording.py
@@ -264,7 +264,7 @@ def test_BaseRecording():
     rec_2d = rec_3d.planarize(axes="zy")
     assert np.allclose(rec_2d.get_channel_locations(), locations_3d[:, [2, 1]])
 
-    # Test save to zarr
+    # test save to zarr
     compressor = get_default_zarr_compressor()
     rec_zarr = rec2.save(format="zarr", folder=cache_folder / "recording", compressor=compressor)
     rec_zarr_loaded = load_extractor(cache_folder / "recording.zarr")
diff --git a/src/spikeinterface/core/tests/test_noise_levels_propagation.py b/src/spikeinterface/core/tests/test_noise_levels_propagation.py
new file mode 100644
index 000000000..d6dbd08ab
--- /dev/null
+++ b/src/spikeinterface/core/tests/test_noise_levels_propagation.py
@@ -0,0 +1,47 @@
+import pytest
+from pathlib import Path
+import shutil
+
+from spikeinterface import set_global_tmp_folder, get_noise_levels
+from spikeinterface.core import generate_recording, concatenate_recordings, aggregate_channels
+
+import numpy as np
+
+if hasattr(pytest, "global_test_folder"):
+    cache_folder = pytest.global_test_folder / "preprocessing"
+else:
+    cache_folder = Path("cache_folder") / "preprocessing"
+
+set_global_tmp_folder(cache_folder)
+
+
+def test_skip_noise_levels_propagation():
+    rec = generate_recording(durations=[5.0], num_channels=4)
+    rec.set_property("test", ["1", "2", "3", "4"])
+    rec = rec.save()
+    noise_level_raw = get_noise_levels(rec, return_scaled=False)
+    assert "noise_level_raw" in rec.get_property_keys()
+
+    rec_frame_slice = rec.frame_slice(start_frame=0, end_frame=1000)
+    assert "noise_level_raw" not in rec_frame_slice.get_property_keys()
+    assert "test" in rec_frame_slice.get_property_keys()
+
+    # make scaled
+    rec.set_channel_gains([100] * 4)
+    rec.set_channel_offsets([0] * 4)
+    noise_level_scaled = get_noise_levels(rec, return_scaled=True)
+    assert "noise_level_scaled" in rec.get_property_keys()
+
+    rec_frame_slice = rec.frame_slice(start_frame=0, end_frame=1000)
+    rec_concat = concatenate_recordings([rec] * 5)
+
+    assert "noise_level_raw" not in rec_concat.get_property_keys()
+    assert "noise_level_scaled" not in rec_concat.get_property_keys()
+    assert "noise_level_raw" not in rec_frame_slice.get_property_keys()
+    assert "noise_level_scaled" not in rec_frame_slice.get_property_keys()
+    assert "test" in rec_frame_slice.get_property_keys()
+    assert "test" in rec_concat.get_property_keys()
+
+
+if __name__ == "__main__":
+    test_skip_noise_levels_propagation()
diff --git a/src/spikeinterface/core/tests/test_recording_tools.py b/src/spikeinterface/core/tests/test_recording_tools.py
index 828a25f95..6e92d155f 100644
--- a/src/spikeinterface/core/tests/test_recording_tools.py
+++ b/src/spikeinterface/core/tests/test_recording_tools.py
@@ -29,14 +29,26 @@ def test_get_closest_channels():
 def test_get_noise_levels():
     rec = generate_recording(num_channels=2, sampling_frequency=1000.0, durations=[60.0])
 
-    noise_levels = get_noise_levels(rec, return_scaled=False)
+    noise_levels_1 = get_noise_levels(rec, return_scaled=False)
+    noise_levels_2 = get_noise_levels(rec, return_scaled=False)
 
     rec.set_channel_gains(0.1)
     rec.set_channel_offsets(0)
-    noise_levels = get_noise_levels(rec, return_scaled=True)
+    noise_levels = get_noise_levels(rec, return_scaled=True, force_recompute=True)
 
     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
+    seed = 0
+    rng = np.random.default_rng(seed=seed)
+    traces = rng.normal(loc=10.0, scale=std, size=(1_000_000, 2))
+    recording = NumpyRecording(traces, 30000)
+
+    assert np.all(noise_levels_1 == noise_levels_2)
+    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_noise_levels_output():
     # Generate a recording following a gaussian distribution to check the result of get_noise.