Unify params and extend tests

Author: alejoe91

spikeinterface/preprocessing/silence_periods.py

@@ -29,8 +29,7 @@ class SilencedPeriodsRecording(BasePreprocessor):
         - 'noise': The periods are filled with a gaussion noise that has the
                    same variance that the one in the recordings, on a per channel
                    basis
-    **random_chunk_kwargs
-        Random seed for random chunk, by default None
+    **random_chunk_kwargs: Keyword arguments for `spikeinterface.core.get_random_data_chunk()` function
 
     Returns
     -------
@@ -52,6 +51,8 @@ def __init__(self, recording, list_periods, mode='zeros',
                 list_periods = [list_periods]
 
         # some checks
+        assert mode in available_modes, f"mode {mode} is not an available mode: {available_modes}"
+
         assert isinstance(list_periods, list), "'list_periods' must be a list (one per segment)"
         assert len(list_periods) == num_seg, "'list_periods' must have the same length as the number of segments"
         assert all(isinstance(list_periods[i], (list, np.ndarray)) for i in range(num_seg)), \
@@ -60,9 +61,8 @@ def __init__(self, recording, list_periods, mode='zeros',
         for periods in list_periods:
             if len(periods) > 0:
                 assert np.all(np.diff(np.array(periods), axis=1) > 0), "t_stops should be larger than t_starts"
-                assert all(e < s for (_, e), (s, _) in zip(periods, periods[1:])), "Intervals should not overlap"
-
-        assert mode in available_modes, f"mode {mode} is not an available mode: {available_modes}"
+                assert np.all(periods[i][1] < periods[i + 1][0] for i in np.arange(len(periods) - 1)), \
+                    "Intervals should not overlap"
 
         if mode in ['noise']:
             noise_levels = get_noise_levels(recording, return_scaled=False, concatenated=True, **random_chunk_kwargs)
@@ -73,7 +73,7 @@ def __init__(self, recording, list_periods, mode='zeros',
         for seg_index, parent_segment in enumerate(recording._recording_segments):
             periods = list_periods[seg_index]
             periods = np.asarray(periods, dtype='int64')
-            periods = np.sort(self.periods, axis=0)
+            periods = np.sort(periods, axis=0)
             rec_segment = SilencedPeriodsRecordingSegment(parent_segment, periods, mode, noise_levels)
             self.add_recording_segment(rec_segment)
 
@@ -84,7 +84,6 @@ def __init__(self, recording, list_periods, mode='zeros',
 class SilencedPeriodsRecordingSegment(BasePreprocessorSegment):
 
     def __init__(self, parent_recording_segment, periods, mode, noise_levels):
-
         BasePreprocessorSegment.__init__(self, parent_recording_segment)
         self.periods = periods
         self.mode = mode
@@ -103,22 +102,23 @@ def get_traces(self, start_frame, end_frame, channel_indices):
 
         if len(self.periods) > 0:
             new_interval = np.array([start_frame, end_frame])
-            lower_index = np.searchsorted(self.periods[:,1], new_interval[0])
-            upper_index = np.searchsorted(self.periods[:,0], new_interval[1])
+            lower_index = np.searchsorted(self.periods[:, 1], new_interval[0])
+            upper_index = np.searchsorted(self.periods[:, 0], new_interval[1])
 
             if upper_index > lower_index:
 
-                intersection = self.periods[lower_index:upper_index]
+                periods_in_interval = self.periods[lower_index:upper_index]
 
-                for i in intersection:
+                for period in periods_in_interval:
 
-                    onset = max(0, i[0] - start_frame)
-                    offset = min(i[1] - start_frame, end_frame)
+                    onset = max(0, period[0] - start_frame)
+                    offset = min(period[1] - start_frame, end_frame)
 
                     if self.mode == 'zeros':
                         traces[onset:offset, :] = 0
                     elif self.mode == 'noise':
-                        traces[onset:offset, :] = self.noise_levels[channel_indices] * np.random.randn(offset-onset, num_channels)
+                        traces[onset:offset, :] = self.noise_levels[channel_indices] * \
+                            np.random.randn(offset - onset, num_channels)
 
         return traces
 

spikeinterface/preprocessing/tests/test_silence.py

@@ -25,21 +25,28 @@ def test_silence():
 
     rec0 = silence_periods(rec, list_periods=[[[0, 1000], [5000, 6000]], []], mode="zeros")
     rec0.save(verbose=False)
-    traces0 = rec0.get_traces(segment_index=0, start_frame=0, end_frame=1000)
-    traces1 = rec0.get_traces(segment_index=0, start_frame=5000, end_frame=6000)
-    assert np.all(traces0 == 0) and np.all(traces1 == 0)
+    traces_in0 = rec0.get_traces(segment_index=0, start_frame=0, end_frame=1000)
+    traces_in1 = rec0.get_traces(segment_index=0, start_frame=5000, end_frame=6000)
+    traces_out0 = rec0.get_traces(segment_index=0, start_frame=2000, end_frame=3000)
+    assert np.all(traces_in0 == 0)
+    assert np.all(traces_in1 == 0)
+    assert not np.all(traces_out0 == 0)
 
     rec1 = silence_periods(rec, list_periods=[[[0, 1000], [5000, 6000]], []], mode="noise")
-    rec1.save(verbose=False)
+    rec1 = rec1.save(folder=cache_folder / "rec_w_noise", verbose=False)
     noise_levels = get_noise_levels(rec, return_scaled=False)
-    traces0 = rec1.get_traces(segment_index=0, start_frame=0, end_frame=1000)
-    traces1 = rec1.get_traces(segment_index=0, start_frame=5000, end_frame=6000)
-    assert np.abs((np.std(traces0, axis=0) - noise_levels) < 0.1).sum() and np.abs((np.std(traces1, axis=0) - noise_levels)).sum() < 0.1
-
-    traces0 = rec0.get_traces(segment_index=0, start_frame=900, end_frame=5100)
-    traces = rec.get_traces(segment_index=0, start_frame=900, end_frame=5100)
-    assert np.all(traces[100:-100] == traces0[100:-100]) and np.all(traces0[:100] == 0) and np.all(traces0[-100:] == 0)
-
+    traces_in0 = rec1.get_traces(segment_index=0, start_frame=0, end_frame=1000)
+    traces_in1 = rec1.get_traces(segment_index=0, start_frame=5000, end_frame=6000)
+    assert np.abs((np.std(traces_in0, axis=0) - noise_levels) < 0.1).sum()
+    assert np.abs((np.std(traces_in1, axis=0) - noise_levels)).sum() < 0.1
+
+    traces_mix = rec0.get_traces(segment_index=0, start_frame=900, end_frame=5100)
+    traces_original = rec.get_traces(segment_index=0, start_frame=900, end_frame=5100)
+    assert np.all(traces_original[100:-100] == traces_mix[100:-100])
+    assert np.all(traces_mix[:100] == 0)
+    assert np.all(traces_mix[-100:] == 0)
+    assert not np.all(traces_mix[:200] == 0)
+    assert not np.all(traces_mix[:-200] == 0)
 
 
 if __name__ == '__main__':