Merge pull request #3880 from catalystneuro/white_matter

White matter Recording Extractor Follow Up

Author: pauladkisson

src/spikeinterface/extractors/tests/test_whitematterrecordingextractor.py

@@ -1,9 +1,11 @@
 import pytest
 import numpy as np
 from pathlib import Path
+import pickle
 
 from spikeinterface.extractors import WhiteMatterRecordingExtractor, BinaryRecordingExtractor
 from spikeinterface.core.numpyextractors import NumpyRecording
+from spikeinterface.core.testing import check_recordings_equal
 from spikeinterface import get_global_dataset_folder, download_dataset
 
 
@@ -57,3 +59,22 @@ def test_on_data():
     assert recording.get_sampling_frequency() == sampling_frequency
     assert recording.get_num_channels() == num_channels
     assert recording.get_duration() == 1.0
+
+
+def test_pickling():
+    file_path = download_dataset(
+        repo=gin_repo, remote_path=remote_path, local_folder=local_folder, update_if_exists=True
+    )
+
+    sampling_frequency = 25_000.0
+    num_channels = 64
+    recording = WhiteMatterRecordingExtractor(
+        file_path=file_path,
+        sampling_frequency=sampling_frequency,
+        num_channels=num_channels,
+        is_filtered=True,
+    )
+    pickled_recording = pickle.dumps(recording)
+    unpickled_recording = pickle.loads(pickled_recording)
+
+    check_recordings_equal(recording, unpickled_recording)

src/spikeinterface/extractors/whitematterrecordingextractor.py

@@ -59,6 +59,14 @@ def __init__(
             channel_ids=channel_ids,
         )
 
+        self._kwargs = {
+            "file_path": file_path,
+            "sampling_frequency": sampling_frequency,
+            "num_channels": num_channels,
+            "channel_ids": channel_ids,
+            "is_filtered": is_filtered,
+        }
+
 
 # Define function equivalent for convenience
 read_whitematter = define_function_from_class(source_class=WhiteMatterRecordingExtractor, name="read_whitematter")