Add cache_to_mem option to WaveformExtractor
@@ -49,7 +49,7 @@ class WaveformExtractor:
"""
extensions = []
- def __init__(self, recording, sorting, folder):
+ def __init__(self, recording, sorting, folder, cache_to_mem=False):
assert recording.get_num_segments() == sorting.get_num_segments(), \
"The recording and sorting objects must have the same number of segments!"
np.testing.assert_almost_equal(recording.get_sampling_frequency(),
@@ -63,12 +63,11 @@ class WaveformExtractor:
self.sorting = sorting
self.folder = Path(folder)
-
-
- # cache in memory
+ # cache in memory/memmap
self._waveforms = {}
self._template_cache = {}
self._params = {}
+ self._cache_to_mem = cache_to_mem
if (self.folder / 'params.json').is_file():
with open(str(self.folder / 'params.json'), 'r') as f:
@@ -102,7 +101,7 @@ class WaveformExtractor:
return we
@classmethod
- def create(cls, recording, sorting, folder, remove_if_exists=False):
+ def create(cls, recording, sorting, folder, cache_to_mem=False, remove_if_exists=False):
folder = Path(folder)
if folder.is_dir():
if remove_if_exists:
@@ -116,7 +115,7 @@ class WaveformExtractor:
if sorting.is_dumpable:
sorting.dump(folder / 'sorting.json', relative_to=None)
- return cls(recording, sorting, folder)
+ return cls(recording, sorting, folder, cache_to_mem)
@classmethod
def register_extension(cls, extension_class):
@@ -399,7 +398,9 @@ class WaveformExtractor:
if not waveform_file.is_file():
raise Exception('Waveforms not extracted yet: '
'please do WaveformExtractor.run_extract_waveforms() first')
-
+ if not self._cache_to_mem:
+ wfs = np.load(waveform_file, mmap_mode="r")
+ else:
wfs = np.load(waveform_file)
self._waveforms[unit_id] = wfs
@@ -799,6 +800,7 @@ def _waveform_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx)
def extract_waveforms(recording, sorting, folder,
load_if_exists=False,
+ cache_to_mem=False,
precompute_template=('average', ),
ms_before=3., ms_after=4.,
max_spikes_per_unit=500,
@@ -821,6 +823,8 @@ def extract_waveforms(recording, sorting, folder,
load_if_exists: bool
If True and waveforms have already been extracted in the specified folder, they are loaded
and not recomputed.
+ cache_to_mem: bool
+ If True, waveforms are loaded and cached in memory (default False)
precompute_template: None or list
Precompute average/std/median for template. If None not precompute.
ms_before: float
@@ -854,7 +858,7 @@ def extract_waveforms(recording, sorting, folder,
if load_if_exists and folder.is_dir():
we = WaveformExtractor.load_from_folder(folder)
else:
- we = WaveformExtractor.create(recording, sorting, folder)
+ we = WaveformExtractor.create(recording, sorting, folder, cache_to_mem)
we.set_params(ms_before=ms_before, ms_after=ms_after, max_spikes_per_unit=max_spikes_per_unit, dtype=dtype,
return_scaled=return_scaled)
we.run_extract_waveforms(**job_kwargs)
