Repository: SpikeInterface/spikeinterface
Pull request #511 (2022-02-14 16:25:50)

Start refactor waveformextractor with waveform_tools.

TODO:
* [x] refactoring with a submodule waveform_tools.py
* [x] use spikes vector
* [x] extract to memory
* [x] add sparsity



Patch:
diff --git a/spikeinterface/core/basesorting.py b/spikeinterface/core/basesorting.py
index eb061a808..0963ae0e4 100644
--- a/spikeinterface/core/basesorting.py
+++ b/spikeinterface/core/basesorting.py
@@ -115,7 +115,8 @@ class BaseSorting(BaseExtractor):
             NpzSortingExtractor.write_sorting(self, save_path)
             cached = NpzSortingExtractor(save_path)
             if self.has_recording():
-                warnings.warn("The registered recording will not be persistent on disk, but only available in memory")
+                warnings.warn(
+                    "The registered recording will not be persistent on disk, but only available in memory")
                 cached.register_recording(self._recording)
         else:
             raise ValueError(f'format {format} not supported')
@@ -129,13 +130,15 @@ class BaseSorting(BaseExtractor):
 
     def select_units(self, unit_ids, renamed_unit_ids=None):
         from spikeinterface import UnitsSelectionSorting
-        sub_sorting = UnitsSelectionSorting(self, unit_ids, renamed_unit_ids=renamed_unit_ids)
+        sub_sorting = UnitsSelectionSorting(
+            self, unit_ids, renamed_unit_ids=renamed_unit_ids)
         return sub_sorting
 
     # add tests
     def frame_slice(self, start_frame, end_frame):
         from spikeinterface import FrameSliceSorting
-        sub_sorting = FrameSliceSorting(self, start_frame=start_frame, end_frame=end_frame)
+        sub_sorting = FrameSliceSorting(
+            self, start_frame=start_frame, end_frame=end_frame)
         return sub_sorting
 
     def get_all_spike_trains(self, outputs='unit_id'):
@@ -148,7 +151,8 @@ class BaseSorting(BaseExtractor):
             spike_times = []
             spike_labels = []
             for i, unit_id in enumerate(self.unit_ids):
-                st = self.get_unit_spike_train(unit_id=unit_id, segment_index=segment_index)
+                st = self.get_unit_spike_train(
+                    unit_id=unit_id, segment_index=segment_index)
                 spike_times.append(st)
                 if outputs == 'unit_id':
                     spike_labels.append(np.array([unit_id] * st.size))
@@ -162,6 +166,35 @@ class BaseSorting(BaseExtractor):
             spikes.append((spike_times, spike_labels))
         return spikes
 
+    def to_spike_vector(self):
+        """
+        Construct a unique structured numpy vector concatenating all spikes 
+        with several fields: sample_ind, unit_index, segment_index.
+
+        See also `get_all_spike_trains()`
+        
+        Returns
+        -------
+        spikes: np.array
+            Structured numpy array ('sample_ind', 'unit_index', 'segment_index') with all spikes
+        """
+        spikes_ = self.get_all_spike_trains(outputs='unit_index')
+
+        n = np.sum([e[0].size for e in spikes_])
+        spike_dtype = [('sample_ind', 'int64'), ('unit_ind',
+                                                 'int64'), ('segment_ind', 'int64')]
+        spikes = np.zeros(n, dtype=spike_dtype)
+
+        pos = 0
+        for segment_index, (spike_times, spike_labels) in enumerate(spikes_):
+            n = spike_times.size
+            spikes[pos:pos+n]['sample_ind'] = spike_times
+            spikes[pos:pos+n]['unit_ind'] = spike_labels
+            spikes[pos:pos+n]['segment_ind'] = segment_index
+            pos += n
+
+        return spikes
+
 
 class BaseSortingSegment(BaseSegment):
     """
diff --git a/spikeinterface/core/core_tools.py b/spikeinterface/core/core_tools.py
index 024c92d3f..9c590ad32 100644
--- a/spikeinterface/core/core_tools.py
+++ b/spikeinterface/core/core_tools.py
@@ -337,7 +337,8 @@ def make_shared_array(shape, dtype):
         raise Exception('SharedMemory is available only for python>=3.8')
 
     dtype = np.dtype(dtype)
-    nbytes = shape[0] * shape[1] * dtype.itemsize
+    # nbytes = shape[0] * shape[1] * dtype.itemsize
+    nbytes = np.prod(shape) * dtype.itemsize
     shm = SharedMemory(name=None, create=True, size=nbytes)
     arr = np.ndarray(shape=shape, dtype=dtype, buffer=shm.buf)
     arr[:] = 0
diff --git a/spikeinterface/core/waveform_extractor.py b/spikeinterface/core/waveform_extractor.py
index 3ceb480f3..cdceecb82 100644
--- a/spikeinterface/core/waveform_extractor.py
+++ b/spikeinterface/core/waveform_extractor.py
@@ -7,7 +7,8 @@ import numpy as np
 from .base import load_extractor
 
 from .core_tools import check_json
-from .job_tools import ChunkRecordingExecutor, ensure_n_jobs, _shared_job_kwargs_doc
+from .job_tools import _shared_job_kwargs_doc
+from spikeinterface.core.waveform_tools import allocate_waveforms, distribute_waveforms_to_buffers
 
 _possible_template_modes = ('average', 'std', 'median')
 
@@ -646,12 +647,10 @@ class WaveformExtractor:
         nbefore = self.nbefore
         nafter = self.nafter
         return_scaled = self.return_scaled
-
-        n_jobs = ensure_n_jobs(self.recording, job_kwargs.get('n_jobs', None))
+        unit_ids = self.sorting.unit_ids
 
         selected_spikes = self.sample_spikes()
 
-        # get spike times
         selected_spike_times = {}
         for unit_id in self.sorting.unit_ids:
             selected_spike_times[unit_id] = []
@@ -659,29 +658,30 @@ class WaveformExtractor:
                 spike_times = self.sorting.get_unit_spike_train(unit_id=unit_id, segment_index=segment_index)
                 sel = selected_spikes[unit_id][segment_index]
                 selected_spike_times[unit_id].append(spike_times[sel])
+        
+        spikes = []
+        for segment_index in range(self.sorting.get_num_segments()):
+            num_in_seg = np.sum([selected_spikes[unit_id][segment_index].size for unit_id in unit_ids])
+            spike_dtype = [('sample_ind', 'int64'), ('unit_ind', 'int64'), ('segment_ind', 'int64')]
+            spikes_ = np.zeros(num_in_seg,  dtype=spike_dtype)
+            pos = 0
+            for unit_ind, unit_id in enumerate(unit_ids):
+                spike_times = self.sorting.get_unit_spike_train(unit_id=unit_id, segment_index=segment_index)
+                sel = selected_spikes[unit_id][segment_index]
+                n = sel.size
+                spikes_[pos:pos+n]['sample_ind'] = spike_times[sel]
+                spikes_[pos:pos+n]['unit_ind'] = unit_ind
+                spikes_[pos:pos+n]['segment_ind'] = segment_index
+                pos += n
+            order = np.argsort(spikes_)
+            spikes_ = spikes_[order]
+            spikes.append(spikes_)
+        spikes = np.concatenate(spikes)
+
+        wf_folder = self.folder / 'waveforms'
+        wfs_arrays, wfs_arrays_info = allocate_waveforms(self.recording, spikes, unit_ids, nbefore, nafter, mode='memmap', folder=wf_folder, dtype=p['dtype'])
+        distribute_waveforms_to_buffers(self.recording, spikes, unit_ids, wfs_arrays_info, nbefore, nafter, return_scaled, **job_kwargs)
 
-        # prepare memmap
-        wfs_memmap = {}
-        for unit_id in self.sorting.unit_ids:
-            file_path = self.folder / 'waveforms' / f'waveforms_{unit_id}.npy'
-            n_spikes = np.sum([e.size for e in selected_spike_times[unit_id]])
-            shape = (n_spikes, self.nsamples, num_chans)
-            wfs = np.zeros(shape, dtype=p['dtype'])
-            np.save(file_path, wfs)
-            # wfs = np.load(file_path, mmap_mode='r+')
-            wfs_memmap[unit_id] = file_path
-
-        # and run
-        func = _waveform_extractor_chunk
-        init_func = _init_worker_waveform_extractor
-        if n_jobs == 1:
-            init_args = (self.recording, self.sorting,)
-        else:
-            init_args = (self.recording.to_dict(), self.sorting.to_dict(),)
-        init_args = init_args + (wfs_memmap, selected_spikes, selected_spike_times, nbefore, nafter, return_scaled)
-        processor = ChunkRecordingExecutor(self.recording, func, init_func, init_args, job_name='extract waveforms',
-                                           **job_kwargs)
-        processor.run()
 
 
 def select_random_spikes_uniformly(recording, sorting, max_spikes_per_unit, nbefore=None, nafter=None):
@@ -726,89 +726,6 @@ def select_random_spikes_uniformly(recording, sorting, max_spikes_per_unit, nbef
     return selected_spikes
 
 
-# used by WaveformExtractor + ChunkRecordingExecutor
-def _init_worker_waveform_extractor(recording, sorting, wfs_memmap,
-                                    selected_spikes, selected_spike_times, nbefore, nafter, return_scaled):
-    # create a local dict per worker
-    worker_ctx = {}
-    if isinstance(recording, dict):
-        from spikeinterface.core import load_extractor
-        recording = load_extractor(recording)
-    worker_ctx['recording'] = recording
-
-    if isinstance(sorting, dict):
-        from spikeinterface.core import load_extractor
-        sorting = load_extractor(sorting)
-    worker_ctx['sorting'] = sorting
-
-    worker_ctx['wfs_memmap_files'] = wfs_memmap
-    worker_ctx['selected_spikes'] = selected_spikes
-    worker_ctx['selected_spike_times'] = selected_spike_times
-    worker_ctx['nbefore'] = nbefore
-    worker_ctx['nafter'] = nafter
-    worker_ctx['return_scaled'] = return_scaled
-
-    num_seg = sorting.get_num_segments()
-    unit_cum_sum = {}
-    for unit_id in sorting.unit_ids:
-        # spike per segment
-        n_per_segment = [selected_spikes[unit_id][i].size for i in range(num_seg)]
-        cum_sum = [0] + np.cumsum(n_per_segment).tolist()
-        unit_cum_sum[unit_id] = cum_sum
-    worker_ctx['unit_cum_sum'] = unit_cum_sum
-
-    return worker_ctx
-
-
-# used by WaveformExtractor + ChunkRecordingExecutor
-def _waveform_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):
-    # recover variables of the worker
-    recording = worker_ctx['recording']
-    sorting = worker_ctx['sorting']
-    wfs_memmap_files = worker_ctx['wfs_memmap_files']
-    selected_spikes = worker_ctx['selected_spikes']
-    selected_spike_times = worker_ctx['selected_spike_times']
-    nbefore = worker_ctx['nbefore']
-    nafter = worker_ctx['nafter']
-    return_scaled = worker_ctx['return_scaled']
-    unit_cum_sum = worker_ctx['unit_cum_sum']
-
-    seg_size = recording.get_num_samples(segment_index=segment_index)
-
-    to_extract = {}
-    for unit_id in sorting.unit_ids:
-        spike_times = selected_spike_times[unit_id][segment_index]
-        i0 = np.searchsorted(spike_times, start_frame)
-        i1 = np.searchsorted(spike_times, end_frame)
-        if i0 != i1:
-            # protect from spikes on border :  spike_time<0 or spike_time>seg_size
-            # useful only when max_spikes_per_unit is not None
-            # waveform will not be extracted and a zeros will be left in the memmap file
-            while (spike_times[i0] - nbefore) < 0 and (i0!=i1):
-                i0 = i0 + 1
-            while (spike_times[i1-1] + nafter) > seg_size and (i0!=i1):
-                i1 = i1 - 1
-
-        if i0 != i1:
-            to_extract[unit_id] = i0, i1, spike_times[i0:i1]
-
-    if len(to_extract) > 0:
-        start = min(st[0] for _, _, st in to_extract.values()) - nbefore
-        end = max(st[-1] for _, _, st in to_extract.values()) + nafter
-        start = int(start)
-        end = int(end)
-
-        # load trace in memory
-        traces = recording.get_traces(start_frame=start, end_frame=end, segment_index=segment_index,
-                                      return_scaled=return_scaled)
-
-        for unit_id, (i0, i1, local_spike_times) in to_extract.items():
-            wfs = np.load(wfs_memmap_files[unit_id], mmap_mode="r+")
-            for i in range(local_spike_times.size):
-                st = local_spike_times[i]
-                st = int(st)
-                pos = unit_cum_sum[unit_id][segment_index] + i0 + i
-                wfs[pos, :, :] = traces[st - start - nbefore:st - start + nafter, :]
 
 
 def extract_waveforms(recording, sorting, folder,
diff --git a/spikeinterface/core/waveform_tools.py b/spikeinterface/core/waveform_tools.py
new file mode 100644
index 000000000..07adc74f8
--- /dev/null
+++ b/spikeinterface/core/waveform_tools.py
@@ -0,0 +1,254 @@
+"""
+This module contains low-level functions to extract snippets of traces (aka "spike waveforms").
+
+This is internally used by WaveformExtractor, but can also be used as a sorting component.
+
+It is a 2-step approach:
+  1. allocate buffers (shared file or memory)
+  2. extract and distribute snippets into buffers (optionally in parallel)
+
+"""
+from pathlib import Path
+
+import numpy as np
+
+from .job_tools import ChunkRecordingExecutor, ensure_n_jobs, _shared_job_kwargs_doc
+from .core_tools import make_shared_array
+
+
+def allocate_waveforms(recording, spikes, unit_ids, nbefore, nafter, mode='memmap', folder=None, dtype=None,
+                       sparsity_mask=None):
+    """
+    Allocate memmap or shered memory buffers before snippet extraction.
+
+    Important note: for the shared memory mode wfs_arrays_info contains reference to
+    the shared memmory buffer, this variable must be reference as long as arrays as used.
+
+    Parameters
+    ----------
+    recording: recording
+        The recording object
+    spikes: 1d numpy array with several fields
+        Spikes handled as a unique vector.
+        This vector can be obtained with: `spikes = Sorting.to_spike_vector()`
+    unit_ids: list ot numpy
+        List of unit_ids
+    nbefore: int
+        N samples before spike
+    nafter: int
+        N samples after spike
+    mode: str
+        Mode to use ('memmap' | 'shared_memory')
+    folder: str or path
+        In case of memmap mode, folder to save npy files
+    dtype: numpy.dtype
+        dtype for waveforms buffer
+    sparsity_mask: None or array of bool
+        If not None shape must be must be (len(unit_ids), len(channel_ids)
+
+    Returns
+    -------
+    wfs_arrays: dict of arrays
+        Arrays for all units (memmap or shared_memmep
+    wfs_arrays_info: dict of info
+        Dictionary to "construct" array in workers process (memmap file or sharemem)
+    """
+
+    nsamples = nbefore + nafter
+
+    dtype = np.dtype(dtype)
+    if mode == 'shared_memory':
+        assert folder is None
+    else:
+        folder = Path(folder)
+
+    # prepare buffers
+    wfs_arrays = {}
+    wfs_arrays_info = {}
+    for unit_ind, unit_id in enumerate(unit_ids):
+        n_spikes = np.sum(spikes['unit_ind'] == unit_ind)
+        if sparsity_mask is None:
+            num_chans = recording.get_num_channels()
+        else:
+            num_chans = np.sum(sparsity_mask[unit_ind, :])
+        shape = (n_spikes, nsamples, num_chans)
+
+        if mode == 'memmap':
+            filename = str(folder / f'waveforms_{unit_id}.npy')
+            arr = np.lib.format.open_memmap(
+                filename, mode='w+', dtype=dtype, shape=shape)
+            wfs_arrays[unit_id] = arr
+            wfs_arrays_info[unit_id] = filename
+        elif mode == 'shared_memory':
+            arr, shm = make_shared_array(shape, dtype)
+            wfs_arrays[unit_id] = arr
+            wfs_arrays_info[unit_id] = (shm, shm.name, dtype.str, shape)
+        else:
+            raise ValueError('allocate_waveforms bad mode')
+
+    return wfs_arrays, wfs_arrays_info
+
+
+def distribute_waveforms_to_buffers(recording, spikes, unit_ids, wfs_arrays_info, nbefore, nafter, return_scaled,
+                                    mode='memmap', sparsity_mask=None, **job_kwargs):
+    """
+    Distribute snippets of traces into corresponding buffers.
+
+    Buffers must be pre-allocated with the `allocate_waveforms()` function.
+
+    Important note, for "shared_memory" mode wfs_arrays_info contain reference to
+    the shared memmory buffer, this variable must be reference as long as arrays as used.
+
+    Parameters
+    ----------
+    recording: recording
+        The recording object
+    spikes: 1d numpy array with several field
+        Spikes handled as a unique vector.
+        This vector can be spikes = Sorting.to_spike_vector()
+    unit_ids: list ot numpy
+        List of unit_ids
+    wfs_arrays_info: dict
+        Dictionary to "construct" array in workers process (memmap file or sharemem)
+    nbefore: int
+        N samples before spike
+    nafter: int
+        N samples after spike
+    return_scaled: bool
+        Scale traces before exporting to buffer or not.
+    mode: str
+        Mode to use ('memmap' | 'shared_memory')
+    sparsity_mask: None or array of bool
+        If not None shape must be must be (len(unit_ids), len(channel_ids)
+
+    {}
+
+    """
+
+    n_jobs = ensure_n_jobs(recording, job_kwargs.get('n_jobs', None))
+
+    inds_by_unit = {}
+    for unit_ind, unit_id in enumerate(unit_ids):
+        inds,  = np.nonzero(spikes['unit_ind'] == unit_ind)
+        inds_by_unit[unit_id] = inds
+
+    # and run
+    func = _waveform_extractor_chunk
+    init_func = _init_worker_waveform_extractor
+    if n_jobs == 1:
+        init_args = (recording, )
+    else:
+        init_args = (recording.to_dict(), )
+    init_args = init_args + (unit_ids, spikes, wfs_arrays_info, nbefore,
+                             nafter, return_scaled, inds_by_unit, mode, sparsity_mask)
+    processor = ChunkRecordingExecutor(recording, func, init_func, init_args, job_name=f'extract waveforms {mode}',
+                                       **job_kwargs)
+    processor.run()
+
+
+distribute_waveforms_to_buffers.__doc__ = distribute_waveforms_to_buffers.__doc__.format(
+    _shared_job_kwargs_doc)
+
+# used by ChunkRecordingExecutor
+
+
+def _init_worker_waveform_extractor(recording, unit_ids, spikes, wfs_arrays_info, nbefore, nafter, return_scaled,
+                                    inds_by_unit, mode, sparsity_mask):
+    # create a local dict per worker
+    worker_ctx = {}
+    if isinstance(recording, dict):
+        from spikeinterface.core import load_extractor
+        recording = load_extractor(recording)
+    worker_ctx['recording'] = recording
+
+    if mode == 'memmap':
+        # ~ if not platform.system().lower().startswith('linux'):
+        # For OSX and windows : need to re open all npy files in r+ mode for each worker
+        wfs_arrays = {}
+        for unit_id, filename in wfs_arrays_info.items():
+            wfs_arrays[unit_id] = np.load(filename, mmap_mode='r+')
+    elif mode == 'shared_memory':
+
+        from multiprocessing.shared_memory import SharedMemory
+        wfs_arrays = {}
+        shms = {}
+        for unit_id, (sm, shm_name, dtype, shape) in wfs_arrays_info.items():
+            shm = SharedMemory(shm_name)
+            arr = np.ndarray(shape=shape, dtype=dtype, buffer=shm.buf)
+            wfs_arrays[unit_id] = arr
+            # we need a reference to all sham otherwise we get segment fault!!!
+            shms[unit_id] = shm
+        worker_ctx['shms'] = shms
+
+    worker_ctx['unit_ids'] = unit_ids
+    worker_ctx['spikes'] = spikes
+    worker_ctx['wfs_arrays'] = wfs_arrays
+    worker_ctx['nbefore'] = nbefore
+    worker_ctx['nafter'] = nafter
+    worker_ctx['return_scaled'] = return_scaled
+    worker_ctx['inds_by_unit'] = inds_by_unit
+    worker_ctx['sparsity_mask'] = sparsity_mask
+
+    return worker_ctx
+
+
+# used by ChunkRecordingExecutor
+def _waveform_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):
+    # recover variables of the worker
+    recording = worker_ctx['recording']
+    unit_ids = worker_ctx['unit_ids']
+    spikes = worker_ctx['spikes']
+    wfs_arrays = worker_ctx['wfs_arrays']
+    nbefore = worker_ctx['nbefore']
+    nafter = worker_ctx['nafter']
+    return_scaled = worker_ctx['return_scaled']
+    inds_by_unit = worker_ctx['inds_by_unit']
+    sparsity_mask = worker_ctx['sparsity_mask']
+
+    seg_size = recording.get_num_samples(segment_index=segment_index)
+
+    # take only spikes with the correct segment_ind
+    # this is a slice so no copy!!
+    s0 = np.searchsorted(spikes['segment_ind'], segment_index)
+    s1 = np.searchsorted(spikes['segment_ind'], segment_index + 1)
+    in_seg_spikes = spikes[s0:s1]
+
+    # take only spikes in range [start_frame, end_frame]
+    # this is a slice so no copy!!
+    i0 = np.searchsorted(in_seg_spikes['sample_ind'], start_frame)
+    i1 = np.searchsorted(in_seg_spikes['sample_ind'], end_frame)
+    if i0 != i1:
+        # protect from spikes on border :  spike_time<0 or spike_time>seg_size
+        # useful only when max_spikes_per_unit is not None
+        # waveform will not be extracted and a zeros will be left in the memmap file
+        while (in_seg_spikes[i0]['sample_ind'] - nbefore) < 0 and (i0 != i1):
+            i0 = i0 + 1
+        while (in_seg_spikes[i1-1]['sample_ind'] + nafter) > seg_size and (i0 != i1):
+            i1 = i1 - 1
+
+    # slice in absolut in spikes vector
+    l0 = i0 + s0
+    l1 = i1 + s0
+
+    if l1 > l0:
+        start = spikes[l0]['sample_ind'] - nbefore
+        end = spikes[l1-1]['sample_ind'] + nafter
+
+        # load trace in memory
+        traces = recording.get_traces(start_frame=start, end_frame=end, segment_index=segment_index,
+                                      return_scaled=return_scaled)
+
+        for unit_ind, unit_id in enumerate(unit_ids):
+            wfs = wfs_arrays[unit_id]
+            # find pos
+            inds = inds_by_unit[unit_id]
+            in_chunk_pos,  = np.nonzero((inds >= l0) & (inds < l1))
+            for pos in in_chunk_pos:
+                sample_ind = spikes[inds[pos]]['sample_ind']
+                wf = traces[sample_ind - start -
+                            nbefore:sample_ind - start + nafter, :]
+
+                if sparsity_mask is None:
+                    wfs[pos, :, :] = wf
+                else:
+                    wfs[pos, :, :] = wf[:, sparsity_mask[unit_ind]]


Test patch:
diff --git a/spikeinterface/core/tests/test_basesorting.py b/spikeinterface/core/tests/test_basesorting.py
index e1dd83107..142ff5ed2 100644
--- a/spikeinterface/core/tests/test_basesorting.py
+++ b/spikeinterface/core/tests/test_basesorting.py
@@ -70,7 +70,10 @@ def test_BaseSorting():
     sorting3 = BaseExtractor.load(folder)
 
     spikes = sorting.get_all_spike_trains()
-    print(spikes)
+    # print(spikes)
+    
+    spikes = sorting.to_spike_vector()
+    # print(spikes)
 
 
 if __name__ == '__main__':
diff --git a/spikeinterface/core/tests/test_waveform_tools.py b/spikeinterface/core/tests/test_waveform_tools.py
new file mode 100644
index 000000000..0706dbf23
--- /dev/null
+++ b/spikeinterface/core/tests/test_waveform_tools.py
@@ -0,0 +1,124 @@
+from pathlib import Path
+import shutil
+import platform
+
+import numpy as np
+
+from spikeinterface.core.testing_tools import generate_recording, generate_sorting
+#~ from spikeinterface import WaveformExtractor, extract_waveforms
+from spikeinterface.core.waveform_tools import allocate_waveforms, distribute_waveforms_to_buffers
+
+
+def _clean_all():
+    folders = ["wf_rec1", "test_waveform_tools", "test_waveform_tools_sparse",]
+    for folder in folders:
+        if Path(folder).exists():
+            shutil.rmtree(folder)
+
+
+def setup_module():
+    _clean_all()
+
+
+def teardown_module():
+    _clean_all()
+
+
+def _check_all_wf_equal(list_wfs_arrays):
+    wfs_arrays0 = list_wfs_arrays[0]
+    for i, wfs_arrays in enumerate(list_wfs_arrays):
+        for unit_id in wfs_arrays.keys():
+            assert np.array_equal(wfs_arrays[unit_id], wfs_arrays0[unit_id])
+
+
+def test_waveform_tools():
+
+
+    durations = [30, 40]
+    sampling_frequency = 30000.
+
+    # 2 segments
+    num_channels = 2
+    recording = generate_recording(num_channels=num_channels, durations=durations, 
+                                   sampling_frequency=sampling_frequency)
+    recording.annotate(is_filtered=True)
+    folder_rec = "wf_rec1"
+    #~ recording = recording.save(folder=folder_rec)
+    num_units = 15
+    sorting = generate_sorting(num_units=num_units, sampling_frequency=sampling_frequency, durations=durations)
+
+    # test with dump !!!!
+    recording = recording.save()
+    sorting = sorting.save()
+
+    
+
+    #~ we = WaveformExtractor.create(recording, sorting, folder)
+    
+    nbefore = int(3. * sampling_frequency / 1000.)
+    nafter = int(4. * sampling_frequency / 1000.)
+    
+    dtype = recording.get_dtype()
+    return_scaled = False
+    
+    spikes = sorting.to_spike_vector()
+    
+    unit_ids = sorting.unit_ids
+
+    some_job_kwargs = [
+        {},
+        {'n_jobs': 1, 'chunk_size': 3000, 'progress_bar':True},
+        {'n_jobs': 2, 'chunk_size': 3000, 'progress_bar':True},
+    ]
+    
+    # memmap mode 
+    list_wfs = []
+    for j, job_kwargs in enumerate(some_job_kwargs):
+        wf_folder = Path(f'test_waveform_tools_{j}')
+        if wf_folder.is_dir():
+            shutil.rmtree(wf_folder)
+        wf_folder.mkdir()
+        wfs_arrays, wfs_arrays_info = allocate_waveforms(recording, spikes, unit_ids, nbefore, nafter, mode='memmap', folder=wf_folder, dtype=dtype)
+        distribute_waveforms_to_buffers(recording, spikes, unit_ids, wfs_arrays_info, nbefore, nafter, return_scaled, **job_kwargs)
+        for unit_ind, unit_id in enumerate(unit_ids):
+            wf = wfs_arrays[unit_id]
+            assert wf.shape[0] == np.sum(spikes['unit_ind'] == unit_ind)
+        list_wfs.append({unit_id: wfs_arrays[unit_id].copy() for unit_id in unit_ids})
+    _check_all_wf_equal(list_wfs)
+    
+
+    # memory
+    if platform.system() != 'Windows':
+        # shared memory on windows is buggy...
+        list_wfs = []
+        for job_kwargs in some_job_kwargs:
+            wfs_arrays, wfs_arrays_info = allocate_waveforms(recording, spikes, unit_ids, nbefore, nafter, mode='shared_memory', folder=None, dtype=dtype)
+            distribute_waveforms_to_buffers(recording, spikes, unit_ids, wfs_arrays_info, nbefore, nafter, return_scaled, mode='shared_memory', **job_kwargs)
+            for unit_ind, unit_id in enumerate(unit_ids):
+                wf = wfs_arrays[unit_id]
+                assert wf.shape[0] == np.sum(spikes['unit_ind'] == unit_ind)
+            list_wfs.append({unit_id: wfs_arrays[unit_id].copy() for unit_id in unit_ids})
+            # to avoid warning we need to first destroy arrays then sharedmemm object
+            del wfs_arrays
+            del wfs_arrays_info
+        _check_all_wf_equal(list_wfs)
+
+    
+    # with sparsity
+    wf_folder = Path('test_waveform_tools_sparse')
+    if wf_folder.is_dir():
+        shutil.rmtree(wf_folder)
+    wf_folder.mkdir()
+    
+    
+    sparsity_mask = np.random.randint(0, 2, size=(unit_ids.size, recording.channel_ids.size), dtype='bool')
+    
+    wfs_arrays, wfs_arrays_info = allocate_waveforms(recording, spikes, unit_ids, nbefore, nafter, mode='memmap', folder=wf_folder, dtype=dtype, sparsity_mask=sparsity_mask)
+    job_kwargs = {'n_jobs': 1, 'chunk_size': 3000, 'progress_bar':True}
+    distribute_waveforms_to_buffers(recording, spikes, unit_ids, wfs_arrays_info, nbefore, nafter, return_scaled, sparsity_mask=sparsity_mask, **job_kwargs)
+    
+
+if __name__ == '__main__':
+    setup_module()
+    test_waveform_tools()
+    
\ No newline at end of file