Simplify MockWaveformExtractor
@@ -3,7 +3,6 @@ import numpy as np
import scipy.signal
import scipy.spatial
-from ..core import WaveformExtractor
from ..core.template_tools import get_template_extremum_channel
from ..postprocessing import compute_correlograms
from ..qualitymetrics import compute_refrac_period_violations, compute_firing_rates
@@ -393,11 +392,21 @@ def compute_templates_diff(sorting, templates, num_channels=5, num_shift=5, pair
return templates_diff
-class MockWaveformExtractor(WaveformExtractor):
- def __init__(self, we, sorting):
- WaveformExtractor.__init__(self, recording=None, sorting=sorting, folder=None,
- rec_attributes=we._rec_attributes, allow_unfiltered=True,
- sparsity=None)
+class MockWaveformExtractor:
+ """
+ Mock WaveformExtractor to be able to run compute_refrac_period_violations()
+ needed for the auto_merge() function.
+ """
+ def __init__(self, waveform_extractor, sorting):
+ self.waveform_extractor = waveform_extractor
+ self._recording = None
+ self.sorting = sorting
+
+ def get_total_samples(self):
+ return self.waveform_extractor.get_total_samples()
+
+ def get_total_duration(self):
+ return self.waveform_extractor.get_total_duration()
def check_improve_contaminations_score(we, pair_mask, contaminations,
