Fix slicing channels of remove artifact median/avg modes
@@ -218,6 +218,9 @@ class RemoveArtifactsRecordingSegment(BasePreprocessorSegment):
def get_traces(self, start_frame, end_frame, channel_indices):
+ if self.mode in ['average', 'median']:
+ traces = self.parent_recording_segment.get_traces(start_frame, end_frame)
+ else:
traces = self.parent_recording_segment.get_traces(start_frame, end_frame, channel_indices)
traces = traces.copy()
@@ -400,6 +403,7 @@ class RemoveArtifactsRecordingSegment(BasePreprocessorSegment):
else:
traces[trace_slice] -= \
(best_amp * self.artifacts[label][artifact_slice]).astype(traces.dtype)
+ traces = traces[:, channel_indices]
return traces
