Repository: spikeinterface/spikeinterface
Issue: error during whitenning

Hello, 
I was trying to add whitenning in my preprocess chain, but I got the following error:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "C:\\Users\\juventin\\AppData\\Local\\Programs\\Python\\Python310\\lib\\multiprocessing\\spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "C:\\Users\\juventin\\AppData\\Local\\Programs\\Python\\Python310\\lib\\multiprocessing\\spawn.py", line 126, in _main
    self = reduction.pickle.load(from_parent)
  File "D:\\local_code_library\\spike_sorters\\spikeinterface\\spikeinterface\\core\\base.py", line 370, in from_dict
    extractor = _load_extractor_from_dict(d)
  File "D:\\local_code_library\\spike_sorters\\spikeinterface\\spikeinterface\\core\\base.py", line 844, in _load_extractor_from_dict
    kwargs[k] = _load_extractor_from_dict(v)
  File "D:\\local_code_library\\spike_sorters\\spikeinterface\\spikeinterface\\core\\base.py", line 861, in _load_extractor_from_dict
    extractor = cls(**kwargs)
  File "D:\\local_code_library\\spike_sorters\\spikeinterface\\spikeinterface\\preprocessing\\whiten.py", line 58, in __init__
    rec_segment = WhitenRecordingSegment(parent_segment, W, M, dtype_)
UnboundLocalError: local variable 'M' referenced before assignment

The issue is associated to the following code:
the recording is a SpikeGLXRecordingExtractor: 384 channels - 1 segments - 30.0kHz - 608.451s

`

    recording = spikeinterface.preprocessing.bandpass_filter(recording)
    recording = spikeinterface.preprocessing.phase_shift(recording)
    recording = spikeinterface.preprocessing.whiten(recording)
    recording = spikeinterface.preprocessing.common_reference(recording, reference='global', operator='median')

    peaks = detect_peaks(recording, method='locally_exclusive', **job_kwargs)

`
Surprisingly, if the replace the detect_peaks by a recording.get_traces(), I don't get the error.


Changed files:
spikeinterface/preprocessing/whiten.py, spikeinterface/preprocessing/tests/test_whiten.py

diff --git a/spikeinterface/preprocessing/tests/test_whiten.py b/spikeinterface/preprocessing/tests/test_whiten.py
index 5fbe34fa..6a5d5494 100644
--- a/spikeinterface/preprocessing/tests/test_whiten.py
+++ b/spikeinterface/preprocessing/tests/test_whiten.py
@@ -1,4 +1,5 @@
 import pytest
+import numpy as np
 from pathlib import Path
 
 from spikeinterface import set_global_tmp_folder
@@ -23,14 +24,13 @@ def test_whiten():
     # test dtype
     rec_int = scale(rec2, dtype="int16")
     rec3 = whiten(rec_int, dtype="float16")
+    rec3 = rec3.save(folder=cache_folder / "rec1")
     assert rec3.get_dtype() == "float16"
 
-    # ~ import matplotlib.pyplot as plt
-    # ~ from spikeinterface.widgets import plot_timeseries
-    # ~ fig, ax = plt.subplots()
-    # ~ ax.plot(rec.get_traces(segment_index=0)[:, 0], color='g')
-    # ~ ax.plot(rec2.get_traces(segment_index=0)[:, 0], color='r')
-    # ~ plt.show()
+    # test parallel
+    rec_par = rec3.save(folder=cache_folder / "rec_par", n_jobs=2)
+    np.testing.assert_array_equal(rec3.get_traces(segment_index=0),
+                                  rec_par.get_traces(segment_index=0))
 
 
 if __name__ == '__main__':
diff --git a/spikeinterface/preprocessing/whiten.py b/spikeinterface/preprocessing/whiten.py
index ff01a6ad..f2d01dbf 100644
--- a/spikeinterface/preprocessing/whiten.py
+++ b/spikeinterface/preprocessing/whiten.py
@@ -23,6 +23,8 @@ class WhitenRecording(BasePreprocessor):
         Random seed for random chunk, by default None
     W : 2d np.array
         Pre-computed whitening matrix, by default None
+    M : 1d np.array
+        Pre-computed means
 
     Returns
     -------
@@ -32,12 +34,14 @@ class WhitenRecording(BasePreprocessor):
     name = 'whiten'
 
     def __init__(self, recording, dtype="float32", num_chunks_per_segment=20,
-                 chunk_size=10000, seed=None, W=None):
+                 chunk_size=10000, seed=None, W=None, M=None):
         # fix dtype
         dtype_ = fix_dtype(recording, dtype)
 
         if W is not None:
+            assert M is not None, "W and M must be not None"
             W = np.array(W)
+            M = np.array(M)
         else:
             random_data = get_random_data_chunks(recording, num_chunks_per_segment=num_chunks_per_segment,
                                                 chunk_size=chunk_size, concatenated=True, seed=seed,
@@ -60,7 +64,8 @@ class WhitenRecording(BasePreprocessor):
 
         self._kwargs = dict(recording=recording.to_dict(), dtype=dtype,
                             num_chunks_per_segment=num_chunks_per_segment,
-                            chunk_size=chunk_size, seed=seed, W=W.tolist())
+                            chunk_size=chunk_size, seed=seed, 
+                            W=W.tolist(), M=M.tolist())
 
 
 class WhitenRecordingSegment(BasePreprocessorSegment):