ibl-neuropixel 1.8.1__py3-none-any.whl → 1.9.1__py3-none-any.whl

{ibl_neuropixel-1.8.1.dist-info → ibl_neuropixel-1.9.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ibl-neuropixel
-Version: 1.8.1
+Version: 1.9.1
 Summary: Collection of tools for Neuropixel 1.0 and 2.0 probes data
 Home-page: https://github.com/int-brain-lab/ibl-neuropixel
 Author: The International Brain Laboratory
@@ -41,6 +41,53 @@ Minimum Python version supported is 3.10
 
 ## Destriping
 ### Getting started
+
+#### Compress a binary file losslessly using `mtscomp`
+
+The mtscomp util implements fast chunked compression for neurophysiology data in a single shard.
+Package repository is [here](https://github.com/int-brain-lab/mtscomp).
+
+
+```python
+from pathlib import Path
+import spikeglx
+file_spikeglx = Path('/datadisk/neuropixel/file.imec0.ap.bin')
+sr = spikeglx.Reader(file_spikeglx)
+sr.compress_file()
+# note: you can use sr.compress_file(keep_original=False) to also remove the orginal bin file
+```
+
+#### Reading raw spikeglx file and manipulating arrays
+
+The mtscomp util implements fast chunked compression for neurophysiology data in a single shard.
+Package repository is [here](https://github.com/int-brain-lab/mtscomp).
+
+```python
+from pathlib import Path
+import spikeglx
+
+import ibldsp.voltage
+
+file_spikeglx = Path('/datadisk/Data/neuropixel/human/Pt01.imec0.ap.bin')
+sr = spikeglx.Reader(file_spikeglx)
+
+# reads in 300ms of data
+raw = sr[10_300_000:10_310_000, :sr.nc - sr.nsync].T
+destripe = ibldsp.voltage.destripe(raw, fs=sr.fs, neuropixel_version=1)
+
+# display with matplotlib backend
+import ibldsp.plots
+ibldsp.plots.voltageshow(raw, fs=sr.fs, title='raw')
+ibldsp.plots.voltageshow(destripe, fs=sr.fs, title='destripe')
+
+# display with QT backend
+from viewephys.gui import viewephys
+eqc = {}
+eqc['raw'] = viewephys(raw, fs=sr.fs, title='raw')
+eqc['destripe'] = viewephys(destripe, fs=sr.fs, title='destripe')
+```
+
+#### Destripe a binary file
 This relies on a fast fourier transform external library: `pip install pyfftw`.
 
 Minimal working example to destripe a neuropixel binary file. 
@@ -71,22 +118,4 @@ The following describes the methods implemented in this repository.
 https://doi.org/10.6084/m9.figshare.19705522
 
 ## Contribution
-Contribution checklist:
-- run tests
-- ruff format
-- PR to main
-
-
-Pypi Release checklist:
-- Edit the version number in `setup.py`
-- add release notes in `release_notes.md`
-
-
-```shell
-ruff format
-tag=X.Y.Z
-git tag -a $tag 
-git push origin $tag
-```
-
-Create new release with tag X.Y.Z (will automatically publish to PyPI)
+Please see our [contribution guidelines](CONTRIBUTING.md) for details on how to contribute to this project.

{ibl_neuropixel-1.8.1.dist-info → ibl_neuropixel-1.9.1.dist-info}/RECORD

@@ -1,6 +1,6 @@
 neuropixel.py,sha256=P7sIBAtGIqKReK7OqMBqdwPaTeHjhHMyfyBRL_AvuQY,37987
-spikeglx.py,sha256=OPvkZdnMguBAuajA8XjJko9N6-UDo4GdKeYp3DhIClU,40865
-ibl_neuropixel-1.8.1.dist-info/licenses/LICENSE,sha256=JJCjBeS78UPiX7TZpE-FnMjNNpCyrFb4s8VDGG2wD10,1087
+spikeglx.py,sha256=4TPXnFGhJahClxr4fA9HwTeiiHBQS9ZEfkWl6t20q2s,41068
+ibl_neuropixel-1.9.1.dist-info/licenses/LICENSE,sha256=JJCjBeS78UPiX7TZpE-FnMjNNpCyrFb4s8VDGG2wD10,1087
 ibldsp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ibldsp/cadzow.py,sha256=pAtxDxBwoNhoxFNc2R5WLwUrmKsq4rQuaglRNgW2Lj8,7251
 ibldsp/cuda_tools.py,sha256=6LpVhYOCuOXEEg8kJ3aOCE4hzA1Yq1dojsbbBQmQCF4,2387
@@ -8,26 +8,28 @@ ibldsp/destripe_gpu.py,sha256=I5jzFocpsYw36kMMd533YThbrQaZix5e1sHqsUjHvO4,2824
 ibldsp/filter_gpu.py,sha256=DPrPBLRXeCh_6BcJWJnPFaxS9Q6kX4nPENZg-c2q5rc,5789
 ibldsp/fourier.py,sha256=RI58nhs4ZZXx1M6EtuhA0vbtkNaBRS2QNv7tPkVomao,10608
 ibldsp/icsd.py,sha256=y9NWOXBB4Nfb5A1fQMKlOu0PdVDVOZ39v2pwk2zzB84,44923
-ibldsp/plots.py,sha256=lgSqnGXMKnJ7fAa3ru30oeIZIAAd-Fz9Cgzx-p2w04k,2064
+ibldsp/plots.py,sha256=XmYC4yca_seZYNEmC5hE5wBiJAl_fi_KU00DbNcM6jI,4577
 ibldsp/raw_metrics.py,sha256=Ie4b7unuFc-XiFc9-tpTsUkph29G-20NvM7iJ25jAPI,5198
 ibldsp/smooth.py,sha256=m_mByXHG_JyFErnYsZ27gXjcqpfwCEuWa6eOb9eFuyg,8033
 ibldsp/spiketrains.py,sha256=lYP1PD4l6T-4KhFu8ZXlbnUUnEQLOriGxN1szacolPY,6878
-ibldsp/utils.py,sha256=p3yvxXdfW36PNmN8qZQ237ZlkPvNvrJ3qCWkDFuy5Q8,13398
-ibldsp/voltage.py,sha256=ID1FDpA9s4qhr6GBy-1SqQlsuU7YjvX5-T25bUaGpDI,39815
-ibldsp/waveform_extraction.py,sha256=IZWMDmsDnC7FdwvhYKzCwriq0rbWw4esGZYKPMQWkhY,26544
+ibldsp/utils.py,sha256=uvEPw1adkppiGXuYBkM_fuuX5owq7LRmA6vm438rrYc,17959
+ibldsp/voltage.py,sha256=Iias93xAvxfRDrzgZT-aw-w4xfWtykx2zWLhI2CxzVI,45408
+ibldsp/waveform_extraction.py,sha256=yKrldgHqpwQ_Dq6xdoSCceKkfrL9FUXnpwKJUM3R41M,26570
 ibldsp/waveforms.py,sha256=5OBLYuM902WS_9WGDDmiTh4BpYWGe7-bQYTMxc2mYII,35166
 neurowaveforms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 neurowaveforms/model.py,sha256=YOPWMMNNS_Op5TyK4Br1i9_Ni41jLSqHie5r1vb5VjY,6729
 tests/integration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/integration/csd_experiments.py,sha256=bddMl2SCzeEM_QnBrZGypUYMKxFVDc6qderyUyX-iew,3158
-tests/integration/test_destripe.py,sha256=6OwqWWz3hJSPGAeEGDcJJkG4bZMnNeaU80AlH7vyrno,6170
+tests/integration/test_destripe.py,sha256=ZV7gasuFib2AbVb63WczgQvc13PbIX9t4pQgamBMgRY,6161
 tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/test_ephys_np2.py,sha256=1wsgS_C5W8tUO_qDyORBRUKGsrB0Gq3wMLAjJcjrNZ4,15599
-tests/unit/test_ibldsp.py,sha256=bekaSW02sZxdn4xML_7WRssUOdVpWPTXmj3IPxyLlWA,24872
 tests/unit/test_neuropixel.py,sha256=ZFKrvTYaYgK5WgOfoHa5x9BNUpRomACPiIm6Kr-A3gw,2511
-tests/unit/test_spikeglx.py,sha256=p0cATRg7xK4WsS_PP_fng_qMATSx6FbRz9kBe8bgIUk,33130
+tests/unit/test_plots.py,sha256=PhCxrEN1Zd1jTgmiwd16_dEghcI7kwmHT3AQmAPpzkA,850
+tests/unit/test_spikeglx.py,sha256=9PrSOPGrYAAQEeJPAOmqc3Rhgia6ftv-zihVWXglhqw,34388
+tests/unit/test_utils.py,sha256=37XQDUqcABYrrsdX17kX54H4e5jld7GOn1ISxtgoa5U,21859
+tests/unit/test_voltage.py,sha256=Nr6KqNGn2yOGPJYnvVzxdM5IiEHvK2FicDR_7fzvTHQ,6228
 tests/unit/test_waveforms.py,sha256=VnFvUi1pteROwwbC5Ebp2lqSxF3a8a7eXHpD8OUeuTg,16237
-ibl_neuropixel-1.8.1.dist-info/METADATA,sha256=nctIH01vFPOU42nUCLmbqWG6ByMZQoZGBqQgCgswQ0c,2505
-ibl_neuropixel-1.8.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ibl_neuropixel-1.8.1.dist-info/top_level.txt,sha256=WtVcEUptnwU6BT72cgGmrWYFGM9d9qCEqe3LwR9FIw4,48
-ibl_neuropixel-1.8.1.dist-info/RECORD,,
+ibl_neuropixel-1.9.1.dist-info/METADATA,sha256=RaS1xeg11qze-sAmPqKVqdKOgUcSk-5l01HNrkX9kIw,3746
+ibl_neuropixel-1.9.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ibl_neuropixel-1.9.1.dist-info/top_level.txt,sha256=WtVcEUptnwU6BT72cgGmrWYFGM9d9qCEqe3LwR9FIw4,48
+ibl_neuropixel-1.9.1.dist-info/RECORD,,

ibldsp/plots.py

@@ -1,9 +1,17 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
+AP_RANGE_UV = 75
+LF_RANGE_UV = 250
+
 
 def show_channels_labels(
-    raw, fs, channel_labels, xfeats, similarity_threshold, psd_hf_threshold=0.02
+    raw,
+    fs,
+    channel_labels,
+    xfeats,
+    similarity_threshold=(-0.5, 1),
+    psd_hf_threshold=0.02,
 ):
     """
     Shows the features side by side a snippet of raw data
@@ -13,14 +21,20 @@ def show_channels_labels(
     nc, ns = raw.shape
     raw = raw - np.mean(raw, axis=-1)[:, np.newaxis]  # removes DC offset
     ns_plot = np.minimum(ns, 3000)
-    vaxis_uv = 250 if fs < 2600 else 75
     fig, ax = plt.subplots(
         1, 5, figsize=(18, 6), gridspec_kw={"width_ratios": [1, 1, 1, 8, 0.2]}
     )
     ax[0].plot(xfeats["xcor_hf"], np.arange(nc))
-    ax[0].plot(
+    ax[0].plot(  # plot channel below the similarity threshold as dead in black
         xfeats["xcor_hf"][(iko := channel_labels == 1)], np.arange(nc)[iko], "k*"
     )
+    ax[0].plot(  # plot the values above the similarity threshold as noisy in red
+        xfeats["xcor_hf"][
+            (iko := np.where(xfeats["xcor_hf"] > similarity_threshold[1]))
+        ],
+        np.arange(nc)[iko],
+        "r*",
+    )
     ax[0].plot(similarity_threshold[0] * np.ones(2), [0, nc], "k--")
     ax[0].plot(similarity_threshold[1] * np.ones(2), [0, nc], "r--")
     ax[0].set(
@@ -30,7 +44,11 @@ def show_channels_labels(
         title="a) dead channel",
     )
     ax[1].plot(xfeats["psd_hf"], np.arange(nc))
-    ax[1].plot(xfeats["psd_hf"][(iko := channel_labels == 2)], np.arange(nc)[iko], "r*")
+    ax[1].plot(
+        xfeats["psd_hf"][(iko := xfeats["psd_hf"] > psd_hf_threshold)],
+        np.arange(nc)[iko],
+        "r*",
+    )
     ax[1].plot(psd_hf_threshold * np.array([1, 1]), [0, nc], "r--")
     ax[1].set(yticklabels=[], xlabel="PSD", ylim=[0, nc], title="b) noisy channel")
     ax[1].sharey(ax[0])
@@ -41,18 +59,77 @@ def show_channels_labels(
     ax[2].plot([-0.75, -0.75], [0, nc], "y--")
     ax[2].set(yticklabels=[], xlabel="LF coherence", ylim=[0, nc], title="c) outside")
     ax[2].sharey(ax[0])
-    im = ax[3].imshow(
-        raw[:, :ns_plot] * 1e6,
-        origin="lower",
-        cmap="PuOr",
-        aspect="auto",
-        vmin=-vaxis_uv,
-        vmax=vaxis_uv,
-        extent=[0, ns_plot / fs * 1e3, 0, nc],
-    )
-    ax[3].set(yticklabels=[], title="d) Raw data", xlabel="time (ms)", ylim=[0, nc])
-    ax[3].grid(False)
+    voltageshow(raw[:, :ns_plot], fs, ax=ax[3], cax=ax[4])
     ax[3].sharey(ax[0])
-    plt.colorbar(im, cax=ax[4], shrink=0.8).ax.set(ylabel="(uV)")
     fig.tight_layout()
     return fig, ax
+
+
+def voltageshow(
+    raw,
+    fs,
+    cmap="PuOr",
+    ax=None,
+    cax=None,
+    cbar_label="Voltage (uV)",
+    scaling=1e6,
+    vrange=None,
+    **axis_kwargs,
+):
+    """
+    Visualizes electrophysiological voltage data as a heatmap.
+
+    This function displays raw voltage data as a color-coded image with appropriate
+    scaling based on the sampling frequency. It automatically selects voltage range
+    based on whether the data is low-frequency (LF) or action potential (AP) data.
+
+    Parameters
+    ----------
+    raw : numpy.ndarray
+        Raw voltage data array with shape (channels, samples), in Volts
+    fs : float
+        Sampling frequency in Hz, used to determine time axis scaling and voltage range.
+    cmap : str, optional
+        Matplotlib colormap name for the heatmap. Default is 'PuOr'.
+    ax : matplotlib.axes.Axes, optional
+        Axes object to plot on. If None, a new figure and axes are created.
+    cax : matplotlib.axes.Axes, optional
+        Axes object for the colorbar. If None and ax is None, a new colorbar axes is created.
+    cbar_label : str, optional
+        Label for the colorbar. Default is 'Voltage (uV)'.
+    vrange: float, optional
+        Voltage range for the colorbar. Defaults to +/- 75 uV for AP and +/- 250 uV for LF.
+    scaling: float, optional
+        Unit transform: default is 1e6: we expect Volts but plot uV.
+    **axis_kwargs: optional
+        Additional keyword arguments for the axis properties, fed to the ax.set() method.
+    Returns
+    -------
+    matplotlib.image.AxesImage
+        The image object created by imshow, which can be used for further customization.
+    """
+    if ax is None:
+        fig, axs = plt.subplots(1, 2, gridspec_kw={"width_ratios": [1, 0.05]})
+        ax, cax = axs
+    nc, ns = raw.shape
+    default_vrange = LF_RANGE_UV if fs < 2600 else AP_RANGE_UV
+    vrange = vrange if vrange is not None else default_vrange
+    im = ax.imshow(
+        raw * scaling,
+        origin="lower",
+        cmap=cmap,
+        aspect="auto",
+        vmin=-vrange,
+        vmax=vrange,
+        extent=[0, ns / fs, 0, nc],
+    )
+    # set the axis properties: we use defaults values that can be overridden by user-provided ones
+    axis_kwargs = (
+        dict(ylim=[0, nc], xlabel="Time (s)", ylabel="Depth (μm)") | axis_kwargs
+    )
+    ax.set(**axis_kwargs)
+    ax.grid(False)
+    if cax is not None:
+        plt.colorbar(im, cax=cax, shrink=0.8).ax.set(ylabel=cbar_label)
+
+    return im

ibldsp/utils.py

@@ -89,7 +89,7 @@ def parabolic_max(x):
     # for 2D arrays, operate along the last dimension
     ns = x.shape[-1]
     axis = -1
-    imax = np.argmax(x, axis=axis)
+    imax = np.nanargmax(x, axis=axis)
 
     if x.ndim == 1:
         v010 = x[np.maximum(np.minimum(imax + np.array([-1, 0, 1]), ns - 1), 0)]
@@ -268,12 +268,64 @@ def make_channel_index(geom, radius=200.0, pad_val=None):
 
 class WindowGenerator(object):
     """
-    `wg = WindowGenerator(ns, nswin, overlap)`
+    A utility class for generating sliding windows for signal processing applications.
 
-    Provide sliding windows indices generator for signal processing applications.
-    For straightforward spectrogram / periodogram implementation, prefer scipy methods !
+    WindowGenerator provides various methods to iterate through windows of a signal
+    with configurable window size and overlap. It's particularly useful for operations
+    like spectrograms, filtering, or any processing that requires windowed analysis.
 
-    Example of implementations in test_dsp.py.
+    Parameters
+    ----------
+    ns : int
+        Total number of samples in the signal to be windowed.
+    nswin : int
+        Number of samples in each window.
+    overlap : int
+        Number of samples that overlap between consecutive windows.
+
+    Attributes
+    ----------
+    ns : int
+        Total number of samples in the signal.
+    nswin : int
+        Number of samples in each window.
+    overlap : int
+        Number of samples that overlap between consecutive windows.
+    nwin : int
+        Total number of windows.
+    iw : int or None
+        Current window index during iteration.
+
+    Notes
+    -----
+    For straightforward spectrogram or periodogram implementation,
+    scipy methods are recommended over this class.
+
+    Examples
+    --------
+    # straight windowing without overlap
+    >>> wg = WindowGenerator(ns=1000, nwin=111)
+    >>> signal = np.random.randn(1000)
+    >>> for window_slice in wg.slice:
+    ...     window_data = signal[window_slice]
+    ...     # Process window_data
+
+    # windowing with overlap (ie. buffers for apodization)
+    >>> for win_slice, valid_slice, win_valid_slice in wg.slices_valid:
+    ...     window = signal[win_slice]
+    ...     # Process window
+    ...     processed = some_function_with_edge_effect(window)
+    ...     # Only use the valid portion for reconstruction
+    ...     recons[valid_slice] = processed[win_valid_slice]
+
+    # splicing add a fade-in / fade-out in the overlap so that reconstruction has unit amplitude
+    >>> recons = np.zeros_like(signal)
+    >>> for win_slice, amplitude in wg.splice:
+    ...     window = signal[win_slice]
+    ...     # Process window
+    ...     processed = some_function(window)
+    ...     # The processed windows is weighted with the amplitude and added to the reconstructed signal
+    ...     recons[win_slice] = recons[win_slice] + processed * amplitude
     """
 
     def __init__(self, ns, nswin, overlap):
@@ -289,14 +341,35 @@ class WindowGenerator(object):
         self.iw = None
 
     @property
-    def firstlast_splicing(self):
+    def splice(self):
         """
-        Generator that yields the indices as well as an amplitude function that can be used
-        to splice the windows together.
-        In the overlap, the amplitude function gradually transitions the amplitude from one window
-        to the next. The amplitudes always sum to one (ie. windows are symmetrical)
+        Generator that yields slices and amplitude arrays for windowed signal processing with splicing.
+
+        This property provides a convenient way to iterate through all windows with their
+        corresponding amplitude arrays for proper signal reconstruction. The amplitude arrays
+        contain tapering values (from a Hann window) at the overlapping regions to ensure
+        unit amplitude of all samples of the original signal
+
+        Yields
+        ------
+        tuple
+            A tuple containing:
+            - slice: A Python slice object representing the current window
+            - amp: A numpy array containing amplitude values for proper splicing/tapering
+              at overlap regions
+
+        Notes
+        -----
+        This is particularly useful for overlap-add methods where windows need to be
+        properly weighted before being combined in the reconstruction process.
+        """
+        for first, last, amp in self.firstlast_splicing:
+            yield slice(first, last), amp
 
-        :return: tuple of (first_index, last_index, amplitude_vector]
+    @property
+    def firstlast_splicing(self):
+        """
+        cf. self.splice
         """
         w = scipy.signal.windows.hann((self.overlap + 1) * 2 + 1, sym=True)[
             1 : self.overlap + 1
@@ -310,7 +383,7 @@ class WindowGenerator(object):
             yield (first, last, amp)
 
     @property
-    def firstlast_valid(self):
+    def firstlast_valid(self, discard_edges=False):
         """
         Generator that yields a tuple of first, last, first_valid, last_valid index of windows
         The valid indices span up to half of the overlap
@@ -318,12 +391,18 @@ class WindowGenerator(object):
         """
         assert self.overlap % 2 == 0, "Overlap must be even"
         for first, last in self.firstlast:
-            first_valid = 0 if first == 0 else first + self.overlap // 2
-            last_valid = last if last == self.ns else last - self.overlap // 2
+            first_valid = (
+                0 if first == 0 and not discard_edges else first + self.overlap // 2
+            )
+            last_valid = (
+                last
+                if last == self.ns and not discard_edges
+                else last - self.overlap // 2
+            )
             yield (first, last, first_valid, last_valid)
 
     @property
-    def firstlast(self, return_valid=False):
+    def firstlast(self):
         """
         Generator that yields first and last index of windows
 
@@ -343,13 +422,51 @@ class WindowGenerator(object):
     @property
     def slice(self):
         """
-        Generator that yields slices of windows
-
-        :return: a slice of the window
+        Generator that yields slice objects for each window in the signal.
+
+        This property provides a convenient way to iterate through all windows
+        defined by the WindowGenerator parameters. Each yielded slice can be
+        used directly to index into the original signal array.
+
+        Yields
+        ------
+        slice
+            A Python slice object representing the current window, defined by
+            its first and last indices. The slice can be used to extract the
+            corresponding window from the original signal.
         """
         for first, last in self.firstlast:
             yield slice(first, last)
 
+    @property
+    def slices_valid(self):
+        """
+        Generator that yields slices for windowed signal processing with valid regions.
+
+        This method generates tuples of slice objects that can be used to extract windows
+        from a signal and identify the valid (non-overlapping) portions within each window.
+        It's particularly useful for reconstruction operations where overlapping regions
+        need special handling.
+
+        Yields
+        ------
+        tuple
+            A tuple containing three slice objects:
+            - slice(first, last): The full window slice
+            - slice(first_valid, last_valid): The valid portion of the signal in absolute indices
+            - slice_window_valid: The valid portion relative to the window (for use within the window)
+
+        Notes
+        -----
+        This generator relies on the firstlast_valid property which provides the
+        indices for both the full windows and their valid regions.
+        """
+        for first, last, first_valid, last_valid in self.firstlast_valid:
+            slice_window_valid = slice(
+                first_valid - first, None if (lv := -(last - last_valid)) == 0 else lv
+            )
+            yield slice(first, last), slice(first_valid, last_valid), slice_window_valid
+
     def slice_array(self, sig, axis=-1):
         """
         Provided an array or sliceable object, generator that yields

ibldsp/voltage.py

@@ -3,6 +3,8 @@ Module to work with raw voltage traces. Spike sorting pre-processing functions.
 """
 
 import inspect
+import joblib
+import tqdm
 from pathlib import Path
 
 import numpy as np
@@ -217,6 +219,7 @@ def kfilt(
         xf, gain = agc(x, wl=lagc, si=1.0, gpu=gpu)
     if ntr_pad > 0:
         # pad the array with a mirrored version of itself and apply a cosine taper
+        ntr_pad = np.min([ntr_pad, xf.shape[0]])
         xf = gp.r_[gp.flipud(xf[:ntr_pad]), xf, gp.flipud(xf[-ntr_pad:])]
     if ntr_tap > 0:
         taper = fourier.fcn_cosine([0, ntr_tap], gpu=gpu)(gp.arange(nxp))  # taper up
@@ -266,6 +269,120 @@ def saturation(
     return saturation, mute
 
 
+def saturation_samples_to_intervals(
+    _saturation: np.ndarray, output_file: Path = None
+) -> pd.DataFrame:
+    """
+    Convert a flat npy file to a table with saturation intervals.
+    :param _saturation: np.ndarray: Boolean array with saturation samples set as True
+    :return:
+    """
+    assert not _saturation[0]
+    ind, pol = ibldsp.utils.fronts(_saturation.astype(np.int8))
+    # if the last sample is positive, make sure the interval is closed by providing an even number of events
+    if len(pol) > 0 and pol[-1] == 1:
+        pol = np.r_[pol, -1]
+        ind = np.r_[ind, _saturation.shape[0] - 1]
+    df_saturation = pd.DataFrame(
+        np.c_[ind[::2], ind[1::2]], columns=["start_sample", "stop_sample"]
+    )
+    if output_file is not None:
+        df_saturation.to_parquet(output_file)
+    return df_saturation
+
+
+def saturation_cbin(
+    bin_file_path: Path,
+    file_saturation: Path = None,
+    max_voltage=None,
+    n_jobs=4,
+    v_per_sec=1e-8,
+    proportion=0.2,
+    mute_window_samples=7,
+) -> Path:
+    """
+    Detect saturation in a compressed binary (cbin) electrophysiology file and save the results.
+
+    This function processes a SpikeGLX binary file in chunks to identify saturated samples
+    and saves the results as a memory-mapped boolean array. Processing is parallelized
+    for improved performance.
+
+    Parameters
+    ----------
+    bin_file_path : Path | spikeglx.Reader
+        Path to the SpikeGLX binary file to be processed or spikeglx.Reader object
+    file_saturation : Path, optional
+        Path where the saturation data will be saved. If None, defaults to
+        "_iblqc_ephysSaturation.samples.npy" in the same directory as the input file
+    max_voltage : np.float, optional
+        one-sided maximum voltage range (V), if not provided will use the spikeglx metadata
+    n_jobs : int, optional
+        Number of parallel jobs to use for processing, defaults to 4
+    v_per_sec : float, optional
+        Maximum derivative of the voltage in V/s (or units/s), defaults to 1e-8
+    proportion : float, optional
+        Threshold proportion (0-1) of channels that must be above threshold to consider
+        a sample as saturated, defaults to 0.2
+    mute_window_samples : int, optional
+        Number of samples for the cosine taper applied to the saturation, defaults to 7
+
+    Returns
+    -------
+    Path
+        Path to the file where the saturation data was saved
+    """
+    if isinstance(bin_file_path, spikeglx.Reader):
+        sr = bin_file_path
+        bin_file_path = sr.file_bin
+    else:
+        sr = spikeglx.Reader(bin_file_path)
+    file_saturation = (
+        file_saturation
+        if file_saturation is not None
+        else bin_file_path.parent.joinpath("_iblqc_ephysSaturation.samples.npy")
+    )
+    max_voltage = max_voltage if max_voltage is not None else sr.range_volts[:-1]
+    # Create a memory-mapped array
+    _saturation = np.lib.format.open_memmap(
+        file_saturation, dtype=bool, mode="w+", shape=(sr.ns,)
+    )
+    _saturation[:] = False  # Initialize all values to False
+    _saturation.flush()  # Make sure to flush to disk
+
+    wg = ibldsp.utils.WindowGenerator(ns=sr.ns, nswin=2**17, overlap=16)
+
+    # we can parallelize this as there is no conflict on output
+    def _saturation_slice(slice_win, slice_valid, slice_relative_valid):
+        sr = spikeglx.Reader(bin_file_path)
+        data = sr[slice_win, : sr.nc - sr.nsync].T
+        satwin, _ = ibldsp.voltage.saturation(
+            data,
+            max_voltage=max_voltage,
+            fs=sr.fs,
+            v_per_sec=v_per_sec,
+            proportion=proportion,
+            mute_window_samples=mute_window_samples,
+        )
+        _saturation[slice_valid] = satwin[slice_relative_valid]
+        _saturation.flush()
+        # getting the list of jobs as a generator allows running tqdm to monitor progress
+
+    jobs = [
+        joblib.delayed(_saturation_slice)(slw, slv, slrv)
+        for (slw, slv, slrv) in wg.slices_valid
+    ]
+    list(
+        tqdm.tqdm(
+            joblib.Parallel(return_as="generator", n_jobs=n_jobs)(jobs), total=wg.nwin
+        )
+    )
+
+    _ = saturation_samples_to_intervals(
+        _saturation, output_file=file_saturation.with_suffix(".pqt")
+    )
+    return file_saturation.with_suffix(".pqt")
+
+
 def interpolate_bad_channels(
     data, channel_labels=None, x=None, y=None, p=1.3, kriging_distance_um=20, gpu=False
 ):
@@ -655,6 +772,9 @@ def decompress_destripe_cbin(
         saturation_data = np.load(file_saturation)
         assert rms_data.shape[0] == time_data.shape[0] * ncv
         rms_data = rms_data.reshape(time_data.shape[0], ncv)
+        # Save the rms data using the original channel index
+        unsort = np.argsort(sr.raw_channel_order)[: -sr.nsync]
+        rms_data = rms_data[:, unsort]
         output_qc_path = (
             output_file.parent if output_qc_path is None else output_qc_path
         )
@@ -662,10 +782,11 @@ def decompress_destripe_cbin(
         np.save(
             output_qc_path.joinpath("_iblqc_ephysTimeRmsAP.timestamps.npy"), time_data
         )
-        np.save(
-            output_qc_path.joinpath("_iblqc_ephysSaturation.samples.npy"),
+        saturation_samples_to_intervals(
             saturation_data,
+            output_file=output_qc_path.joinpath("_iblqc_ephysSaturation.samples.pqt"),
         )
+        file_saturation.unlink()
 
 
 def detect_bad_channels(
@@ -781,7 +902,7 @@ def detect_bad_channels(
     window_size = 25  # Choose based on desired smoothing (e.g., 25 samples)
     kernel = np.ones(window_size) / window_size
     # Apply convolution
-    signal_filtered = np.convolve(signal_noisy, kernel, mode='same')
+    signal_filtered = np.convolve(signal_noisy, kernel, mode="same")
 
     diff_x = np.diff(signal_filtered)
     indx = np.where(diff_x < -0.02)[0]  # hardcoded threshold
@@ -934,16 +1055,39 @@ def stack(data, word, fcn_agg=np.nanmean, header=None):
 
 def current_source_density(lfp, h, n=2, method="diff", sigma=1 / 3):
     """
-    Compute the current source density (CSD) of a given LFP signal recorded on neuropixel 1 or 2
-    :param data: LFP signal (n_channels, n_samples)
-    :param h: trace header dictionary
-    :param n: the n derivative
-    :param method: diff (straight double difference) or kernel CSD (needs the KCSD python package)
-    :param sigma: conductivity, defaults to 1/3 S.m-1
-    :return:
+    Compute the current source density (CSD) of a given LFP signal recorded on Neuropixel probes.
+
+    The CSD estimates the location of current sources and sinks in neural tissue based on
+    the spatial distribution of local field potentials (LFPs). This implementation supports
+    both the standard double-derivative method and kernel CSD method.
+
+    The CSD is computed for each column of the Neuropixel probe layout separately.
+
+    Parameters
+    ----------
+    lfp : numpy.ndarray
+        LFP signal array with shape (n_channels, n_samples)
+    h : dict
+        Trace header dictionary containing probe geometry information with keys:
+        'x', 'y' for electrode coordinates, 'col' for column indices, and 'row' for row indices
+    n : int, optional
+        Order of the derivative for the 'diff' method, defaults to 2
+    method : str, optional
+        Method to compute CSD:
+        - 'diff': standard finite difference method (default)
+        - 'kcsd': kernel CSD method (requires the KCSD Python package)
+    sigma : float, optional
+        Tissue conductivity in Siemens per meter, defaults to 1/3 S.m-1
+
+    Returns
+    -------
+    numpy.ndarray
+        Current source density with the same shape as the input LFP array.
+        Positive values indicate current sources, negative values indicate sinks.
+        Units are in A.m-3 (amperes per cubic meter).
     """
     csd = np.zeros(lfp.shape, dtype=np.float64) * np.nan
-    xy = h["x"] + 1j * h["y"]
+    xy = (h["x"] + 1j * h["y"]) / 1e6
     for col in np.unique(h["col"]):
         ind = np.where(h["col"] == col)[0]
         isort = np.argsort(h["row"][ind])
@@ -990,7 +1134,6 @@ def _svd_denoise(datr, rank):
 
 def svd_denoise_npx(datr, rank=None, collection=None):
     """
-
     :param datr: [nc, ns]
     :param rank:
     :param collection:

ibldsp/waveform_extraction.py

@@ -280,6 +280,7 @@ def extract_wfs_cbin(
     chunksize_samples=int(30_000),
     reader_kwargs=None,
     n_jobs=None,
+    wfs_dtype=np.float32,
     preprocess_steps=None,
     seed=None,
     scratch_dir=None,

spikeglx.py

@@ -144,8 +144,8 @@ class Reader:
         sglx_file = str(self.file_bin)
         if self.is_mtscomp:
             self._raw = mtscomp.Reader()
-            ch_file = self.ch_file or _get_companion_file(sglx_file, ".ch")
-            self._raw.open(self.file_bin, ch_file)
+            self.ch_file = self._parse_ch_file()
+            self._raw.open(self.file_bin, self.ch_file)
             if self._raw.shape != (self.ns, self.nc):
                 ftsec = self._raw.shape[0] / self.fs
                 if not self.ignore_warnings:  # avoid the checks for streaming data
@@ -392,10 +392,8 @@ class Reader:
         if "out" not in kwargs:
             kwargs["out"] = self.file_bin.with_suffix(".bin")
         assert self.is_mtscomp
-        if file_ch is None:
-            file_ch = self.file_bin.with_suffix(".ch")
-
-        r = mtscomp.decompress(self.file_bin, file_ch, **kwargs)
+        ch_file = self._parse_ch_file(file_ch)
+        r = mtscomp.decompress(self.file_bin, ch_file, **kwargs)
         r.close()
         if not keep_original:
             self.close()
@@ -411,14 +409,15 @@ class Reader:
         """
         if file_meta is None:
             file_meta = Path(self.file_bin).with_suffix(".meta")
-
+        file_ch = file_ch if file_ch is not None else self.ch_file
         if scratch_dir is None:
             bin_file = Path(self.file_bin).with_suffix(".bin")
         else:
             scratch_dir.mkdir(exist_ok=True, parents=True)
-            bin_file = scratch_dir / Path(self.file_bin).with_suffix(".bin").name
-            file_meta_scratch = scratch_dir / file_meta.name
-            shutil.copy(self.file_meta_data, file_meta_scratch)
+            bin_file = (
+                Path(scratch_dir).joinpath(self.file_bin.name).with_suffix(".bin")
+            )
+            shutil.copy(self.file_meta_data, bin_file.parent / self.file_meta_data.name)
         if not bin_file.exists():
             t0 = time.time()
             _logger.info("File is compressed, decompressing to a temporary file...")
@@ -460,6 +459,12 @@ class Reader:
         log_func(f"SHA1 computed: {sc}")
         return sm == sc
 
+    def _parse_ch_file(self, ch_file=None):
+        ch_file = (
+            _get_companion_file(self.file_bin, ".ch") if ch_file is None else ch_file
+        )
+        return ch_file
+
 
 class OnlineReader(Reader):
     @property
@@ -995,7 +1000,7 @@ def _mock_spikeglx_file(
     meta_file,
     ns,
     nc,
-    sync_depth,
+    sync_depth=16,
     random=False,
     int2volts=0.6 / 32768,
     corrupt=False,

tests/integration/test_destripe.py

@@ -4,6 +4,7 @@ import logging
 import shutil
 import unittest
 from pathlib import Path
+import pandas as pd
 
 import neuropixel
 import spikeglx
@@ -84,11 +85,11 @@ class TestEphysSpikeSortingMultiProcess(unittest.TestCase):
             shutil.rmtree(self.file_path.parent)
 
     def _assert_qc(self):
-        sr = spikeglx.Reader(self.file_path)
-        saturated = np.load(
-            self.file_path.parent.joinpath("_iblqc_ephysSaturation.samples.npy")
+        df_saturated = pd.read_parquet(
+            self.file_path.parent.joinpath("_iblqc_ephysSaturation.samples.pqt")
         )
-        self.assertEqual(sr.ns, saturated.size)
+        self.assertTrue(df_saturated.shape[1] == 2)
+
         self.assertTrue(
             self.file_path.parent.joinpath("_iblqc_ephysTimeRmsAP.rms.npy").exists()
         )

tests/unit/test_plots.py

@@ -0,0 +1,30 @@
+import unittest
+
+import numpy as np
+
+import ibldsp.plots
+import ibldsp.voltage
+
+
+class TestPlots(unittest.TestCase):
+    def test_voltage(self):
+        ibldsp.plots.voltageshow(
+            (np.random.rand(384, 2000) - 0.5) / 1e6 * 20, fs=30_000
+        )
+
+    def test_bad_channels(self):
+        np.random.seed(0)
+        raw = np.random.randn(384, 2000) / 1e6 * 15
+        raw += np.random.randn(1, 2000) / 1e6 * 2
+        raw[66] *= 2
+        raw[166] = 0
+        fs = 30_000
+        labels, features = ibldsp.voltage.detect_bad_channels(raw, fs)
+        ibldsp.plots.show_channels_labels(
+            raw=raw,
+            fs=30_000,
+            channel_labels=labels,
+            xfeats=features,
+        )
+        np.testing.assert_array_equal(np.argwhere(labels == 2), 66)
+        np.testing.assert_array_equal(np.argwhere(labels == 1), 166)

tests/unit/test_spikeglx.py

@@ -2,6 +2,7 @@ from pathlib import Path
 import shutil
 import tempfile
 import unittest
+import uuid
 
 import numpy as np
 from iblutil.io import hashfile
@@ -243,6 +244,8 @@ class TestsSpikeGLX_compress(unittest.TestCase):
             with spikeglx.Reader(self.file_cbin, open=False) as sc:
                 self.assertTrue(sc.is_mtscomp)
                 compare_data(sr_ref, sc)
+                # here we make sure the chunks file has been registered as a property
+                self.assertEqual(sc.ch_file, self.file_cbin.with_suffix(".ch"))
 
             # test decompression in-place
             sc.decompress_file(keep_original=False, overwrite=True)
@@ -669,9 +672,29 @@ class TestsBasicReader(unittest.TestCase):
     Tests the basic usage where there is a flat binary and no metadata associated
     """
 
-    def test_get_companion_file(self):
-        import uuid
+    def test_integration_companion_files_cbin(self):
+        with tempfile.TemporaryDirectory() as td:
+            bin_file_orig = Path(td) / "toto.ap.bin"
+            meta_file = Path(td) / f"toto.ap.{str(uuid.uuid4())}.meta"
+            ch_file = Path(td) / f"toto.ap.{str(uuid.uuid4())}.ch"
+            spikeglx._mock_spikeglx_file(
+                bin_file_orig,
+                meta_file=Path(TEST_PATH).joinpath("sample3B_g0_t0.imec1.ap.meta"),
+                ns=90_000,
+                nc=385,
+            )
+            sr = spikeglx.Reader(bin_file_orig)
+            sr.compress_file(keep_original=False)
+            cbin_file = Path(td) / f"toto.ap.{str(uuid.uuid4())}.cbin"
+            shutil.move(bin_file_orig.with_suffix(".cbin"), cbin_file)
+            shutil.move(bin_file_orig.with_suffix(".ch"), ch_file)
+            shutil.move(bin_file_orig.with_suffix(".meta"), meta_file)
+            sr = spikeglx.Reader(cbin_file)
+            self.assertEqual(sr.file_bin, cbin_file)
+            self.assertEqual(sr.file_meta_data, meta_file)
+            self.assertEqual(sr.ch_file, ch_file)
 
+    def test_get_companion_file(self):
         with tempfile.TemporaryDirectory() as td:
             sglx_file = Path(td) / f"sample3A_g0_t0.imec.ap.{str(uuid.uuid4())}.bin"
             meta_file = Path(td) / f"sample3A_g0_t0.imec.ap.{str(uuid.uuid4())}.meta"

tests/unit/{test_ibldsp.py → test_utils.py}

@@ -7,7 +7,6 @@ import spikeglx
 import ibldsp.fourier as fourier
 import ibldsp.utils as utils
 import ibldsp.voltage as voltage
-import ibldsp.cadzow as cadzow
 import ibldsp.smooth as smooth
 import ibldsp.spiketrains as spiketrains
 import ibldsp.raw_metrics as raw_metrics
@@ -70,8 +69,8 @@ class TestSyncTimestamps(unittest.TestCase):
 
 class TestParabolicMax(unittest.TestCase):
     # expected values
-    maxi = np.array([np.nan, 0, 3.04166667, 3.04166667, 5, 5])
-    ipeak = np.array([np.nan, 0, 5.166667, 2.166667, 0, 7])
+    maxi = np.array([0.0, 0.0, 3.04166667, 3.04166667, 5, 5])
+    ipeak = np.array([0.0, 0.0, 5.166667, 2.166667, 0, 7])
     # input
     x = np.array(
         [
@@ -368,6 +367,13 @@ class TestWindowGenerator(unittest.TestCase):
         for first, last, amp in wg.firstlast_splicing:
             sig_out[first:last] = sig_out[first:last] + amp * sig_in[first:last]
         np.testing.assert_allclose(sig_out, sig_in)
+        # now performs the same operation with the new interface
+        sig_in = np.random.randn(600)
+        sig_out = np.zeros_like(sig_in)
+        wg = utils.WindowGenerator(ns=600, nswin=100, overlap=20)
+        for slicewin, amp in wg.splice:
+            sig_out[slicewin] = sig_out[slicewin] + amp * sig_in[slicewin]
+        np.testing.assert_allclose(sig_out, sig_in)
 
     def test_firstlast_valid(self):
         sig_in = np.random.randn(600)
@@ -377,6 +383,15 @@ class TestWindowGenerator(unittest.TestCase):
             sig_out[first_valid:last_valid] = sig_in[first_valid:last_valid]
         np.testing.assert_array_equal(sig_out, sig_in)
 
+    def test_slices_valid(self):
+        sig_in = np.random.randn(600)
+        sig_out = np.zeros_like(sig_in)
+        wg = utils.WindowGenerator(ns=600, nswin=39, overlap=20)
+        for slice_win, slice_valid, slice_win_valid in wg.slices_valid:
+            win = sig_in[slice_win]
+            sig_out[slice_valid] = win[slice_win_valid]
+        np.testing.assert_array_equal(sig_out, sig_in)
+
     def test_tscale(self):
         wg = utils.WindowGenerator(ns=500, nswin=100, overlap=50)
         ts = wg.tscale(fs=1000)
@@ -430,103 +445,6 @@ class TestFrontDetection(unittest.TestCase):
         np.testing.assert_array_equal(utils.rises(a, step=3, analog=True), 283)
 
 
-class TestVoltage(unittest.TestCase):
-    def test_destripe_parameters(self):
-        import inspect
-
-        _, _, spatial_fcn = voltage._get_destripe_parameters(
-            30_000, None, None, k_filter=True
-        )
-        assert "kfilt" in inspect.getsource(spatial_fcn)
-        _, _, spatial_fcn = voltage._get_destripe_parameters(
-            2_500, None, None, k_filter=False
-        )
-        assert "car" in inspect.getsource(spatial_fcn)
-        _, _, spatial_fcn = voltage._get_destripe_parameters(
-            2_500, None, None, k_filter=None
-        )
-        assert "dat: dat" in inspect.getsource(spatial_fcn)
-        _, _, spatial_fcn = voltage._get_destripe_parameters(
-            2_500, None, None, k_filter=lambda dat: 3 * dat
-        )
-        assert "lambda dat: 3 * dat" in inspect.getsource(spatial_fcn)
-
-    def test_fk(self):
-        """
-        creates a couple of plane waves and separate them using the velocity HP filter
-        """
-        ntr, ns, sr, dx, v1, v2 = (500, 2000, 0.002, 5, 2000, 1000)
-        data = np.zeros((ntr, ns), np.float32)
-        data[:, :100] = utils.ricker(100, 4)
-        offset = np.arange(ntr) * dx
-        offset = np.abs(offset - np.mean(offset))
-        data_v1 = fourier.fshift(data, offset / v1 / sr)
-        data_v2 = fourier.fshift(data, offset / v2 / sr)
-
-        noise = np.random.randn(ntr, ns) / 60
-        fk = voltage.fk(
-            data_v1 + data_v2 + noise,
-            si=sr,
-            dx=dx,
-            vbounds=[1200, 1500],
-            ntr_pad=10,
-            ntr_tap=15,
-            lagc=0.25,
-        )
-        fknoise = voltage.fk(
-            noise, si=sr, dx=dx, vbounds=[1200, 1500], ntr_pad=10, ntr_tap=15, lagc=0.25
-        )
-        # at least 90% of the traces should be below 50dB and 98% below 40 dB
-        assert np.mean(20 * np.log10(utils.rms(fk - data_v1 - fknoise)) < -50) > 0.9
-        assert np.mean(20 * np.log10(utils.rms(fk - data_v1 - fknoise)) < -40) > 0.98
-        # test the K option
-        kbands = np.sin(np.arange(ns) / ns * 8 * np.pi) / 10
-        fkk = voltage.fk(
-            data_v1 + data_v2 + kbands,
-            si=sr,
-            dx=dx,
-            vbounds=[1200, 1500],
-            ntr_pad=40,
-            ntr_tap=15,
-            lagc=0.25,
-            kfilt={"bounds": [0, 0.01], "btype": "hp"},
-        )
-        assert np.mean(20 * np.log10(utils.rms(fkk - data_v1)) < -40) > 0.9
-        # from easyqc.gui import viewseis
-        # a = viewseis(data_v1 + data_v2 + kbands, .002, title='input')
-        # b = viewseis(fkk, .002, title='output')
-        # c = viewseis(data_v1 - fkk, .002, title='test')
-
-    def test_saturation(self):
-        np.random.seed(7654)
-        data = (np.random.randn(384, 30_000).astype(np.float32) + 20) * 1e-6
-        saturated, mute = voltage.saturation(data, max_voltage=1200)
-        np.testing.assert_array_equal(saturated, 0)
-        np.testing.assert_array_equal(mute, 1.0)
-        # now we stick a big waveform in the middle of the recorder and expect some saturation
-        w = utils.ricker(100, 4)
-        w = np.minimum(1200, w / w.max() * 1400)
-        data[:, 13_600:13700] = data[0, 13_600:13700] + w * 1e-6
-        saturated, mute = voltage.saturation(
-            data,
-            max_voltage=np.ones(
-                384,
-            )
-            * 1200
-            * 1e-6,
-        )
-        self.assertGreater(np.sum(saturated), 5)
-        self.assertGreater(np.sum(mute == 0), np.sum(saturated))
-
-
-class TestCadzow(unittest.TestCase):
-    def test_trajectory_matrixes(self):
-        assert np.all(
-            cadzow.traj_matrix_indices(4) == np.array([[1, 0], [2, 1], [3, 2]])
-        )
-        assert np.all(cadzow.traj_matrix_indices(3) == np.array([[1, 0], [2, 1]]))
-
-
 class TestStack(unittest.TestCase):
     def test_simple_stack(self):
         ntr, ns = (24, 400)

tests/unit/test_voltage.py

@@ -0,0 +1,160 @@
+import numpy as np
+import tempfile
+from pathlib import Path
+import unittest
+
+import pandas as pd
+
+import spikeglx
+import ibldsp.voltage
+import ibldsp.fourier
+import ibldsp.utils
+import ibldsp.cadzow
+
+
+class TestDestripe(unittest.TestCase):
+    def test_destripe_parameters(self):
+        import inspect
+
+        _, _, spatial_fcn = ibldsp.voltage._get_destripe_parameters(
+            30_000, None, None, k_filter=True
+        )
+        assert "kfilt" in inspect.getsource(spatial_fcn)
+        _, _, spatial_fcn = ibldsp.voltage._get_destripe_parameters(
+            2_500, None, None, k_filter=False
+        )
+        assert "car" in inspect.getsource(spatial_fcn)
+        _, _, spatial_fcn = ibldsp.voltage._get_destripe_parameters(
+            2_500, None, None, k_filter=None
+        )
+        assert "dat: dat" in inspect.getsource(spatial_fcn)
+        _, _, spatial_fcn = ibldsp.voltage._get_destripe_parameters(
+            2_500, None, None, k_filter=lambda dat: 3 * dat
+        )
+        assert "lambda dat: 3 * dat" in inspect.getsource(spatial_fcn)
+
+    def test_fk(self):
+        """
+        creates a couple of plane waves and separate them using the velocity HP filter
+        """
+        ntr, ns, sr, dx, v1, v2 = (500, 2000, 0.002, 5, 2000, 1000)
+        data = np.zeros((ntr, ns), np.float32)
+        data[:, :100] = ibldsp.utils.ricker(100, 4)
+        offset = np.arange(ntr) * dx
+        offset = np.abs(offset - np.mean(offset))
+        data_v1 = ibldsp.fourier.fshift(data, offset / v1 / sr)
+        data_v2 = ibldsp.fourier.fshift(data, offset / v2 / sr)
+
+        noise = np.random.randn(ntr, ns) / 60
+        fk = ibldsp.voltage.fk(
+            data_v1 + data_v2 + noise,
+            si=sr,
+            dx=dx,
+            vbounds=[1200, 1500],
+            ntr_pad=10,
+            ntr_tap=15,
+            lagc=0.25,
+        )
+        fknoise = ibldsp.voltage.fk(
+            noise, si=sr, dx=dx, vbounds=[1200, 1500], ntr_pad=10, ntr_tap=15, lagc=0.25
+        )
+        # at least 90% of the traces should be below 50dB and 98% below 40 dB
+        assert (
+            np.mean(20 * np.log10(ibldsp.utils.rms(fk - data_v1 - fknoise)) < -50) > 0.9
+        )
+        assert (
+            np.mean(20 * np.log10(ibldsp.utils.rms(fk - data_v1 - fknoise)) < -40)
+            > 0.98
+        )
+        # test the K option
+        kbands = np.sin(np.arange(ns) / ns * 8 * np.pi) / 10
+        fkk = ibldsp.voltage.fk(
+            data_v1 + data_v2 + kbands,
+            si=sr,
+            dx=dx,
+            vbounds=[1200, 1500],
+            ntr_pad=40,
+            ntr_tap=15,
+            lagc=0.25,
+            kfilt={"bounds": [0, 0.01], "btype": "hp"},
+        )
+        assert np.mean(20 * np.log10(ibldsp.utils.rms(fkk - data_v1)) < -40) > 0.9
+        # from easyqc.gui import viewseis
+        # a = viewseis(data_v1 + data_v2 + kbands, .002, title='input')
+        # b = viewseis(fkk, .002, title='output')
+        # c = viewseis(data_v1 - fkk, .002, title='test')
+
+
+class TestSaturation(unittest.TestCase):
+    def test_saturation_cbin(self):
+        nsat = 252
+        ns, nc = (350_072, 384)
+        s2v = np.float32(2.34375e-06)
+        sat = ibldsp.utils.fcn_cosine([0, 100])(
+            np.arange(nsat)
+        ) - ibldsp.utils.fcn_cosine([150, 250])(np.arange(nsat))
+        range_volt = 0.0012
+        sat = (sat / s2v * 0.0012).astype(np.int16)
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_bin = Path(temp_dir) / "binary.bin"
+            data = np.memmap(file_bin, dtype=np.int16, mode="w+", shape=(ns, nc))
+            data[50_000 : 50_000 + nsat, :] = sat[:, np.newaxis]
+
+            _sr = spikeglx.Reader(
+                file_bin, fs=30_000, dtype=np.int16, nc=nc, nsync=0, s2v=s2v
+            )
+            file_saturation = ibldsp.voltage.saturation_cbin(
+                _sr, max_voltage=range_volt, n_jobs=1
+            )
+            df_sat = pd.read_parquet(file_saturation)
+            assert np.sum(df_sat["stop_sample"] - df_sat["start_sample"]) == 67
+
+    def test_saturation(self):
+        np.random.seed(7654)
+        data = (np.random.randn(384, 30_000).astype(np.float32) + 20) * 1e-6
+        saturated, mute = ibldsp.voltage.saturation(data, max_voltage=1200)
+        np.testing.assert_array_equal(saturated, 0)
+        np.testing.assert_array_equal(mute, 1.0)
+        # now we stick a big waveform in the middle of the recorder and expect some saturation
+        w = ibldsp.utils.ricker(100, 4)
+        w = np.minimum(1200, w / w.max() * 1400)
+        data[:, 13_600:13700] = data[0, 13_600:13700] + w * 1e-6
+        saturated, mute = ibldsp.voltage.saturation(
+            data,
+            max_voltage=np.ones(
+                384,
+            )
+            * 1200
+            * 1e-6,
+        )
+        self.assertGreater(np.sum(saturated), 5)
+        self.assertGreater(np.sum(mute == 0), np.sum(saturated))
+
+    def test_saturation_intervals_output(self):
+        saturation = np.zeros(50_000, dtype=bool)
+        # we test empty files, make sure we can read/write from empty parquet
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Create a file path within the temporary directory
+            temp_file = Path(temp_dir).joinpath("saturation.pqt")
+            df_nothing = ibldsp.voltage.saturation_samples_to_intervals(
+                saturation, output_file=Path(temp_dir).joinpath("saturation.pqt")
+            )
+            df_nothing2 = pd.read_parquet(temp_file)
+        self.assertEqual(df_nothing.shape[0], 0)
+        self.assertEqual(df_nothing2.shape[0], 0)
+        # for the case with saturation intervals, we simply test the number of rows correspond to the events
+        saturation[3441:3509] = True
+        saturation[45852:45865] = True
+        df_sat = ibldsp.voltage.saturation_samples_to_intervals(saturation)
+        self.assertEqual(81, np.sum(df_sat["stop_sample"] - df_sat["start_sample"]))
+
+
+class TestCadzow(unittest.TestCase):
+    def test_trajectory_matrixes(self):
+        assert np.all(
+            ibldsp.cadzow.traj_matrix_indices(4) == np.array([[1, 0], [2, 1], [3, 2]])
+        )
+        assert np.all(
+            ibldsp.cadzow.traj_matrix_indices(3) == np.array([[1, 0], [2, 1]])
+        )