Repository: SpikeInterface/spikeinterface
Pull request #2870 (2024-05-17 12:32:58)

Components of SC2
Modifications for components and SC2


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index d502d8302..89ea05e5b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -159,8 +159,8 @@ test = [
 ]
 
 docs = [
-    "Sphinx==5.1.1",
-    "sphinx_rtd_theme==1.0.0",
+    "Sphinx",
+    "sphinx_rtd_theme",
     "sphinx-gallery",
     "sphinx-design",
     "numpydoc",
@@ -173,6 +173,7 @@ docs = [
     "hdbscan>=0.8.33",   # For sorters spykingcircus2 + tridesclous
     "numba", # For many postprocessing functions
     "xarray", # For use of SortingAnalyzer zarr format
+    "networkx",
     # for release we need pypi, so this needs to be commented
     "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",  # We always build from the latest version
     "neo @ git+https://github.com/NeuralEnsemble/python-neo.git",  # We always build from the latest version
diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py
index e5691603a..83e035e84 100644
--- a/src/spikeinterface/core/analyzer_extension_core.py
+++ b/src/spikeinterface/core/analyzer_extension_core.py
@@ -330,17 +330,26 @@ class ComputeTemplates(AnalyzerExtension):
 
             return_scaled = self.sorting_analyzer.return_scaled
 
-            self.data["average"], self.data["std"] = estimate_templates_with_accumulator(
+            return_std = "std" in self.params["operators"]
+            output = estimate_templates_with_accumulator(
                 recording,
                 some_spikes,
                 unit_ids,
                 self.nbefore,
                 self.nafter,
                 return_scaled=return_scaled,
-                return_std=True,
+                return_std=return_std,
                 **job_kwargs,
             )
 
+            # Output of estimate_templates_with_accumulator is either (templates,) or (templates, stds)
+            if return_std:
+                templates, stds = output
+                self.data["average"] = templates
+                self.data["std"] = stds
+            else:
+                self.data["average"] = output
+
     def _compute_and_append_from_waveforms(self, operators):
         if not self.sorting_analyzer.has_extension("waveforms"):
             raise ValueError(f"Computing templates with operators {operators} needs the 'waveforms' extension")
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 85f83b081..7a77dc28c 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -5,6 +5,7 @@ from pathlib import Path
 from itertools import chain
 import os
 import json
+import math
 import pickle
 import weakref
 import shutil
@@ -237,7 +238,21 @@ class SortingAnalyzer:
         return_scaled=True,
     ):
         # some checks
-        assert sorting.sampling_frequency == recording.sampling_frequency
+        if sorting.sampling_frequency != recording.sampling_frequency:
+            if math.isclose(sorting.sampling_frequency, recording.sampling_frequency, abs_tol=1e-2, rel_tol=1e-5):
+                warnings.warn(
+                    "Sorting and Recording have a small difference in sampling frequency. "
+                    "This could be due to rounding of floats. Using the sampling frequency from the Recording."
+                )
+                # we make a copy here to change the smapling frequency
+                sorting = NumpySorting.from_sorting(sorting, with_metadata=True, copy_spike_vector=True)
+                sorting._sampling_frequency = recording.sampling_frequency
+            else:
+                raise ValueError(
+                    f"Sorting and Recording sampling frequencies are too different: "
+                    f"recording: {recording.sampling_frequency} - sorting: {sorting.sampling_frequency}. "
+                    "Ensure that you are associating the correct Recording and Sorting when creating a SortingAnalyzer."
+                )
         # check that multiple probes are non-overlapping
         all_probes = recording.get_probegroup().probes
         check_probe_do_not_overlap(all_probes)
diff --git a/src/spikeinterface/core/template_tools.py b/src/spikeinterface/core/template_tools.py
index c5fc55887..69706b3bd 100644
--- a/src/spikeinterface/core/template_tools.py
+++ b/src/spikeinterface/core/template_tools.py
@@ -56,7 +56,7 @@ def _get_nbefore(one_object):
 def get_template_amplitudes(
     templates_or_sorting_analyzer,
     peak_sign: "neg" | "pos" | "both" = "neg",
-    mode: "extremum" | "at_index" = "extremum",
+    mode: "extremum" | "at_index" | "peak_to_peak" = "extremum",
     return_scaled: bool = True,
     abs_value: bool = True,
 ):
@@ -67,11 +67,13 @@ def get_template_amplitudes(
     ----------
     templates_or_sorting_analyzer: Templates | SortingAnalyzer
         A Templates or a SortingAnalyzer object
-    peak_sign: "neg" | "pos" | "both", default: "neg"
-        Sign of the template to compute best channels
-    mode: "extremum" | "at_index", default: "extremum"
-        "extremum":  max or min
-        "at_index": take value at spike index
+    peak_sign:  "neg" | "pos" | "both"
+        Sign of the template to find extremum channels
+    mode: "extremum" | "at_index" | "peak_to_peak", default: "at_index"
+        Where the amplitude is computed
+        * "extremum": take the peak value (max or min depending on `peak_sign`)
+        * "at_index": take value at `nbefore` index
+        * "peak_to_peak": take the peak-to-peak amplitude
     return_scaled: bool, default True
         The amplitude is scaled or not.
     abs_value: bool = True
@@ -83,7 +85,7 @@ def get_template_amplitudes(
         Dictionary with unit ids as keys and template amplitudes as values
     """
     assert peak_sign in ("both", "neg", "pos"), "'peak_sign' must be 'both', 'neg', or 'pos'"
-    assert mode in ("extremum", "at_index"), "'mode' must be 'extremum' or 'at_index'"
+    assert mode in ("extremum", "at_index", "peak_to_peak"), "'mode' must be 'extremum', 'at_index', or 'peak_to_peak'"
 
     unit_ids = templates_or_sorting_analyzer.unit_ids
     before = _get_nbefore(templates_or_sorting_analyzer)
@@ -107,6 +109,8 @@ def get_template_amplitudes(
                 values = np.abs(template[before, :])
             elif peak_sign in ["neg", "pos"]:
                 values = template[before, :]
+        elif mode == "peak_to_peak":
+            values = np.ptp(template, axis=0)
 
         if abs_value:
             values = np.abs(values)
@@ -119,7 +123,7 @@ def get_template_amplitudes(
 def get_template_extremum_channel(
     templates_or_sorting_analyzer,
     peak_sign: "neg" | "pos" | "both" = "neg",
-    mode: "extremum" | "at_index" = "extremum",
+    mode: "extremum" | "at_index" | "peak_to_peak" = "extremum",
     outputs: "id" | "index" = "id",
 ):
     """
@@ -129,11 +133,13 @@ def get_template_extremum_channel(
     ----------
     templates_or_sorting_analyzer: Templates | SortingAnalyzer
         A Templates or a SortingAnalyzer object
-    peak_sign: "neg" | "pos" | "both", default: "neg"
-        Sign of the template to compute best channels
-    mode: "extremum" | "at_index", default: "extremum"
-        "extremum":  max or min
-        "at_index": take value at spike index
+    peak_sign:  "neg" | "pos" | "both"
+        Sign of the template to find extremum channels
+    mode: "extremum" | "at_index" | "peak_to_peak", default: "at_index"
+        Where the amplitude is computed
+        * "extremum": take the peak value (max or min depending on `peak_sign`)
+        * "at_index": take value at `nbefore` index
+        * "peak_to_peak": take the peak-to-peak amplitude
     outputs: "id" | "index", default: "id"
         * "id": channel id
         * "index": channel index
@@ -145,7 +151,7 @@ def get_template_extremum_channel(
         as values
     """
     assert peak_sign in ("both", "neg", "pos"), "`peak_sign` must be one of `both`, `neg`, or `pos`"
-    assert mode in ("extremum", "at_index"), "`mode` must be either `extremum` or `at_index`"
+    assert mode in ("extremum", "at_index", "peak_to_peak"), "'mode' must be 'extremum', 'at_index', or 'peak_to_peak'"
     assert outputs in ("id", "index"), "`outputs` must be either `id` or `index`"
 
     unit_ids = templates_or_sorting_analyzer.unit_ids
@@ -184,8 +190,8 @@ def get_template_extremum_channel_peak_shift(templates_or_sorting_analyzer, peak
     ----------
     templates_or_sorting_analyzer: Templates | SortingAnalyzer
         A Templates or a SortingAnalyzer object
-    peak_sign: "neg" | "pos" | "both", default: "neg"
-        Sign of the template to compute best channels
+    peak_sign:  "neg" | "pos" | "both"
+        Sign of the template to find extremum channels
 
     Returns
     -------
@@ -230,7 +236,7 @@ def get_template_extremum_channel_peak_shift(templates_or_sorting_analyzer, peak
 def get_template_extremum_amplitude(
     templates_or_sorting_analyzer,
     peak_sign: "neg" | "pos" | "both" = "neg",
-    mode: "extremum" | "at_index" = "at_index",
+    mode: "extremum" | "at_index" | "peak_to_peak" = "at_index",
     abs_value: bool = True,
 ):
     """
@@ -241,11 +247,12 @@ def get_template_extremum_amplitude(
     templates_or_sorting_analyzer: Templates | SortingAnalyzer
         A Templates or a SortingAnalyzer object
     peak_sign:  "neg" | "pos" | "both"
-        Sign of the template to compute best channels
-    mode: "extremum" | "at_index", default: "at_index"
+        Sign of the template to find extremum channels
+    mode: "extremum" | "at_index" | "peak_to_peak", default: "at_index"
         Where the amplitude is computed
-        "extremum":  max or min
-        "at_index": take value at spike index
+        * "extremum": take the peak value (max or min depending on `peak_sign`)
+        * "at_index": take value at `nbefore` index
+        * "peak_to_peak": take the peak-to-peak amplitude
     abs_value: bool = True
         Whether the extremum amplitude should be returned as an absolute value or not
 
@@ -256,7 +263,7 @@ def get_template_extremum_amplitude(
         Dictionary with unit ids as keys and amplitudes as values
     """
     assert peak_sign in ("both", "neg", "pos"), "'peak_sign' must be  'neg' or 'pos' or 'both'"
-    assert mode in ("extremum", "at_index"), "'mode' must be 'extremum' or 'at_index'"
+    assert mode in ("extremum", "at_index", "peak_to_peak"), "'mode' must be 'extremum', 'at_index', or 'peak_to_peak'"
     unit_ids = templates_or_sorting_analyzer.unit_ids
     channel_ids = templates_or_sorting_analyzer.channel_ids
 
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
index 1ba96defa..eed0a581b 100644
--- a/src/spikeinterface/core/zarrextractors.py
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from pathlib import Path
 import numpy as np
 import zarr
@@ -308,13 +309,14 @@ def get_default_zarr_compressor(clevel: int = 5):
     return Blosc(cname="zstd", clevel=clevel, shuffle=Blosc.BITSHUFFLE)
 
 
-def add_properties_and_annotations(
-    zarr_group: zarr.hierarchy.Group, recording_or_sorting: Union[BaseRecording, BaseSorting]
-):
+def add_properties_and_annotations(zarr_group: zarr.hierarchy.Group, recording_or_sorting: BaseRecording | BaseSorting):
     # save properties
     prop_group = zarr_group.create_group("properties")
     for key in recording_or_sorting.get_property_keys():
         values = recording_or_sorting.get_property(key)
+        if values.dtype.kind == "O":
+            warnings.warn(f"Property {key} not saved because it is a python Object type")
+            continue
         prop_group.create_dataset(name=key, data=values, compressor=None)
 
     # save annotations
diff --git a/src/spikeinterface/extractors/phykilosortextractors.py b/src/spikeinterface/extractors/phykilosortextractors.py
index 05aee160f..e26c2eb02 100644
--- a/src/spikeinterface/extractors/phykilosortextractors.py
+++ b/src/spikeinterface/extractors/phykilosortextractors.py
@@ -165,14 +165,18 @@ class BasePhyKilosortSortingExtractor(BaseSorting):
                 self.set_property(key="group", values=cluster_info[prop_name])
             elif prop_name == "cluster_id":
                 self.set_property(key="original_cluster_id", values=cluster_info[prop_name])
-            elif prop_name != "group":
-                self.set_property(key=prop_name, values=cluster_info[prop_name])
             elif prop_name == "group":
                 # rename group property to 'quality'
                 self.set_property(key="quality", values=cluster_info[prop_name])
             else:
                 if load_all_cluster_properties:
-                    self.set_property(key=prop_name, values=cluster_info[prop_name])
+                    # pandas loads strings as objects
+                    if cluster_info[prop_name].values.dtype.kind == "O":
+                        prop_dtype = type(cluster_info[prop_name].values[0])
+                        values_ = cluster_info[prop_name].values.astype(prop_dtype)
+                    else:
+                        values_ = cluster_info[prop_name].values
+                    self.set_property(key=prop_name, values=values_)
 
         self.annotate(phy_folder=str(phy_folder.resolve()))
 
diff --git a/src/spikeinterface/preprocessing/highpass_spatial_filter.py b/src/spikeinterface/preprocessing/highpass_spatial_filter.py
index f51beca73..c0bf86931 100644
--- a/src/spikeinterface/preprocessing/highpass_spatial_filter.py
+++ b/src/spikeinterface/preprocessing/highpass_spatial_filter.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import numpy as np
 
 from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
+from .filter import fix_dtype
 from ..core import order_channels_by_depth, get_chunk_with_margin
 from ..core.core_tools import define_function_from_class
 
@@ -47,6 +48,8 @@ class HighpassSpatialFilterRecording(BasePreprocessor):
         Order of spatial butterworth filter
     highpass_butter_wn : float, default: 0.01
         Critical frequency (with respect to Nyquist) of spatial butterworth filter
+    dtype : dtype, default: None
+        The dtype of the output traces. If None, the dtype is the same as the input traces
 
     Returns
     -------
@@ -73,6 +76,7 @@ class HighpassSpatialFilterRecording(BasePreprocessor):
         agc_window_length_s=0.1,
         highpass_butter_order=3,
         highpass_butter_wn=0.01,
+        dtype=None,
     ):
         BasePreprocessor.__init__(self, recording)
 
@@ -117,6 +121,8 @@ class HighpassSpatialFilterRecording(BasePreprocessor):
         butter_kwargs = dict(btype="highpass", N=highpass_butter_order, Wn=highpass_butter_wn)
         sos_filter = scipy.signal.butter(**butter_kwargs, output="sos")
 
+        dtype = fix_dtype(recording, dtype)
+
         for parent_segment in recording._recording_segments:
             rec_segment = HighPassSpatialFilterSegment(
                 parent_segment,
@@ -128,6 +134,7 @@ class HighpassSpatialFilterRecording(BasePreprocessor):
                 sos_filter,
                 order_f,
                 order_r,
+                dtype=dtype,
             )
             self.add_recording_segment(rec_segment)
 
@@ -155,6 +162,7 @@ class HighPassSpatialFilterSegment(BasePreprocessorSegment):
         sos_filter,
         order_f,
         order_r,
+        dtype,
     ):
         BasePreprocessorSegment.__init__(self, parent_recording_segment)
         self.parent_recording_segment = parent_recording_segment
@@ -178,6 +186,7 @@ class HighPassSpatialFilterSegment(BasePreprocessorSegment):
         self.order_r = order_r
         # get filter params
         self.sos_filter = sos_filter
+        self.dtype = dtype
 
     def get_traces(self, start_frame, end_frame, channel_indices):
         if channel_indices is None:
@@ -234,7 +243,7 @@ class HighPassSpatialFilterSegment(BasePreprocessorSegment):
             traces = traces[left_margin:-right_margin, channel_indices]
         else:
             traces = traces[left_margin:, channel_indices]
-        return traces
+        return traces.astype(self.dtype, copy=False)
 
 
 # function for API
diff --git a/src/spikeinterface/sorters/internal/spyking_circus2.py b/src/spikeinterface/sorters/internal/spyking_circus2.py
index 6575aba15..2af28fb17 100644
--- a/src/spikeinterface/sorters/internal/spyking_circus2.py
+++ b/src/spikeinterface/sorters/internal/spyking_circus2.py
@@ -45,18 +45,19 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
             "select_per_channel": False,
             "seed": 42,
         },
-        "drift_correction": {"preset": "nonrigid_fast_and_accurate"},
+        "apply_motion_correction": True,
+        "motion_correction": {"preset": "nonrigid_fast_and_accurate"},
         "merging": {
             "minimum_spikes": 10,
             "corr_diff_thresh": 0.5,
             "template_metric": "cosine",
             "censor_correlograms_ms": 0.4,
-            "num_channels": 5,
+            "num_channels": None,
         },
         "clustering": {"legacy": True},
-        "matching": {"method": "circus-omp-svd"},
+        "matching": {"method": "wobble"},
         "apply_preprocessing": True,
-        "matched_filtering": False,
+        "matched_filtering": True,
         "cache_preprocessing": {"mode": "memory", "memory_limit": 0.5, "delete_cache": True},
         "multi_units_only": False,
         "job_kwargs": {"n_jobs": 0.8},
@@ -121,6 +122,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         ms_before = params["general"].get("ms_before", 2)
         ms_after = params["general"].get("ms_after", 2)
         radius_um = params["general"].get("radius_um", 100)
+        exclude_sweep_ms = params["detection"].get("exclude_sweep_ms", max(ms_before, ms_after) / 2)
 
         ## First, we are filtering the data
         filtering_params = params["filtering"].copy()
@@ -133,15 +135,17 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
             recording_f.annotate(is_filtered=True)
 
         valid_geometry = check_probe_for_drift_correction(recording_f)
-        if params["drift_correction"] is not None:
+        if params["apply_motion_correction"]:
             if not valid_geometry:
-                print("Geometry of the probe does not allow 1D drift correction")
+                if verbose:
+                    print("Geometry of the probe does not allow 1D drift correction")
                 motion_folder = None
             else:
-                print("Motion correction activated (probe geometry compatible)")
+                if verbose:
+                    print("Motion correction activated (probe geometry compatible)")
                 motion_folder = sorter_output_folder / "motion"
-                params["drift_correction"].update({"folder": motion_folder})
-                recording_f = correct_motion(recording_f, **params["drift_correction"])
+                params["motion_correction"].update({"folder": motion_folder})
+                recording_f = correct_motion(recording_f, **params["motion_correction"])
         else:
             motion_folder = None
 
@@ -163,7 +167,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         detection_params.update(job_kwargs)
 
         detection_params["radius_um"] = detection_params.get("radius_um", 50)
-        detection_params["exclude_sweep_ms"] = detection_params.get("exclude_sweep_ms", 0.5)
+        detection_params["exclude_sweep_ms"] = exclude_sweep_ms
         detection_params["noise_levels"] = noise_levels
 
         fs = recording_w.get_sampling_frequency()
@@ -212,6 +216,8 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
             clustering_params["waveforms"]["ms_after"] = ms_after
             clustering_params["job_kwargs"] = job_kwargs
             clustering_params["noise_levels"] = noise_levels
+            clustering_params["ms_before"] = exclude_sweep_ms
+            clustering_params["ms_after"] = exclude_sweep_ms
             clustering_params["tmp_folder"] = sorter_output_folder / "clustering"
 
             legacy = clustering_params.get("legacy", True)
@@ -275,12 +281,10 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
             matching_params["templates"] = templates
             matching_job_params = job_kwargs.copy()
 
-            if matching_method == "circus-omp-svd":
-
-                for value in ["chunk_size", "chunk_memory", "total_memory", "chunk_duration"]:
-                    if value in matching_job_params:
-                        matching_job_params[value] = None
-                matching_job_params["chunk_duration"] = "100ms"
+            for value in ["chunk_size", "chunk_memory", "total_memory", "chunk_duration"]:
+                if value in matching_job_params:
+                    matching_job_params[value] = None
+            matching_job_params["chunk_duration"] = "100ms"
 
             spikes = find_spikes_from_templates(
                 recording_w, matching_method, method_kwargs=matching_params, **matching_job_params
@@ -308,7 +312,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         merging_params = params["merging"].copy()
 
         if len(merging_params) > 0:
-            if params["drift_correction"] and motion_folder is not None:
+            if params["motion_correction"] and motion_folder is not None:
                 from spikeinterface.preprocessing.motion import load_motion_info
 
                 motion_info = load_motion_info(motion_folder)
diff --git a/src/spikeinterface/sortingcomponents/clustering/circus.py b/src/spikeinterface/sortingcomponents/clustering/circus.py
index ce7a78e4c..06d2f1f6d 100644
--- a/src/spikeinterface/sortingcomponents/clustering/circus.py
+++ b/src/spikeinterface/sortingcomponents/clustering/circus.py
@@ -58,8 +58,8 @@ class CircusClustering:
         },
         "radius_um": 100,
         "n_svd": [5, 2],
-        "ms_before": 2,
-        "ms_after": 2,
+        "ms_before": 0.5,
+        "ms_after": 0.5,
         "noise_levels": None,
         "tmp_folder": None,
         "job_kwargs": {},
diff --git a/src/spikeinterface/sortingcomponents/clustering/tdc.py b/src/spikeinterface/sortingcomponents/clustering/tdc.py
index c9a7da453..90938f208 100644
--- a/src/spikeinterface/sortingcomponents/clustering/tdc.py
+++ b/src/spikeinterface/sortingcomponents/clustering/tdc.py
@@ -11,7 +11,6 @@ from spikeinterface.core import (
     get_noise_levels,
     NumpySorting,
     get_channel_distances,
-    estimate_templates_with_accumulator,
     Templates,
     compute_sparsity,
     get_global_tmp_folder,
@@ -37,8 +36,6 @@ from spikeinterface.sortingcomponents.clustering.tools import compute_template_f
 
 from sklearn.decomposition import TruncatedSVD
 
-import hdbscan
-
 
 class TdcClustering:
     """
@@ -63,6 +60,8 @@ class TdcClustering:
 
     @classmethod
     def main_function(cls, recording, peaks, params):
+        import hdbscan
+
         job_kwargs = params["job_kwargs"]
 
         if params["folder"] is None:
diff --git a/src/spikeinterface/widgets/spikes_on_traces.py b/src/spikeinterface/widgets/spikes_on_traces.py
index 9f1a3faeb..8b69c5b7b 100644
--- a/src/spikeinterface/widgets/spikes_on_traces.py
+++ b/src/spikeinterface/widgets/spikes_on_traces.py
@@ -92,7 +92,6 @@ class SpikesOnTracesWidget(BaseWidget):
         self.check_extensions(sorting_analyzer, "unit_locations")
 
         sorting: BaseSorting = sorting_analyzer.sorting
-        recording: BaseRecording = sorting_analyzer.recording
 
         if unit_ids is None:
             unit_ids = sorting.get_unit_ids()
@@ -219,9 +218,9 @@ class SpikesOnTracesWidget(BaseWidget):
 
                     nbefore = nafter = int(dp.spike_width_ms / 2 * sorting_analyzer.sampling_frequency / 1000)
                     waveform_idxs = spike_frames_to_plot[:, None] + np.arange(-nbefore, nafter) - frame_range[0]
-                    waveform_idxs = np.clip(waveform_idxs, 0, len(traces_widget.data_plot["times"]) - 1)
+                    waveform_idxs = np.clip(waveform_idxs, 0, len(traces_widget.data_plot["times_in_range"]) - 1)
 
-                    times = traces_widget.data_plot["times"][waveform_idxs]
+                    times = traces_widget.data_plot["times_in_range"][waveform_idxs]
 
                     # discontinuity
                     times[:, -1] = np.nan
diff --git a/src/spikeinterface/widgets/traces.py b/src/spikeinterface/widgets/traces.py
index f5e808e00..0282e5fe2 100644
--- a/src/spikeinterface/widgets/traces.py
+++ b/src/spikeinterface/widgets/traces.py
@@ -4,7 +4,7 @@ import warnings
 
 import numpy as np
 
-from ..core import BaseRecording, order_channels_by_depth
+from ..core import BaseRecording
 from .base import BaseWidget, to_attr
 from .utils import get_some_colors, array_to_image
 
@@ -33,6 +33,12 @@ class TracesWidget(BaseWidget):
         * "auto": auto switch depending on the channel count ("line" if less than 64 channels, "map" otherwise)
     return_scaled: bool, default: False
         If True and the recording has scaled traces, it plots the scaled traces
+    events: np.array | list[np.narray] or None, default: None
+        Events to display as vertical lines.
+        The numpy arrays cen either be of dtype float, with event times in seconds,
+        or a structured array with the "time" field,
+        and optional "duration" and "label" fields.
+        For multi-segment recordings, provide a list of numpy array events, one for each segment.
     cmap: matplotlib colormap, default: "RdBu_r"
         matplotlib colormap used in mode "map"
     show_channel_ids: bool, default: False
@@ -67,6 +73,9 @@ class TracesWidget(BaseWidget):
         return_scaled=False,
         cmap="RdBu_r",
         show_channel_ids=False,
+        events=None,
+        events_color="gray",
+        events_alpha=0.5,
         color_groups=False,
         color=None,
         clim=None,
@@ -114,6 +123,15 @@ class TracesWidget(BaseWidget):
         else:
             channel_locations = None
 
+        if not rec0.has_time_vector(segment_index=segment_index):
+            times = None
+            t_start = 0
+            t_end = rec0.get_duration(segment_index=segment_index)
+        else:
+            times = rec0.get_times(segment_index=segment_index)
+            t_start = times[0]
+            t_end = times[-1]
+
         layer_keys = list(recordings.keys())
 
         if segment_index is None:
@@ -123,14 +141,13 @@ class TracesWidget(BaseWidget):
 
         fs = rec0.get_sampling_frequency()
         if time_range is None:
-            time_range = (0, 1.0)
+            time_range = (t_start, t_start + 1.0)
         time_range = np.array(time_range)
-        if time_range[1] > rec0.get_duration(segment_index=segment_index):
+        if time_range[1] > t_end:
             warnings.warn(
-                "You have selected a time after the end of the segment. The range will be clipped to "
-                f"{rec0.get_duration(segment_index=segment_index)}"
+                "You have selected a time after the end of the segment. The range will be clipped to " f"{t_end}"
             )
-            time_range[1] = rec0.get_duration(segment_index=segment_index)
+            time_range[1] = t_end
 
         assert mode in ("auto", "line", "map"), 'Mode must be one of "auto","line", "map"'
         if mode == "auto":
@@ -141,8 +158,8 @@ class TracesWidget(BaseWidget):
         mode = mode
         cmap = cmap
 
-        times, list_traces, frame_range, channel_ids = _get_trace_list(
-            recordings, channel_ids, time_range, segment_index, return_scaled=return_scaled
+        times_in_range, list_traces, frame_range, channel_ids = _get_trace_list(
+            recordings, channel_ids, time_range, segment_index, return_scaled=return_scaled, times=times
         )
 
         list_traces = [traces * scale for traces in list_traces]
@@ -203,6 +220,30 @@ class TracesWidget(BaseWidget):
             else:
                 raise TypeError(f"'clim' can be None, tuple, or dict! Unsupported type {type(clim)}")
 
+        # check events
+        if events is not None:
+            num_segments = rec0.get_num_segments()
+            if num_segments > 1:
+                assert (
+                    len(events) == num_segments
+                ), f"events must be a list with the events for each of the {num_segments} segments"
+                evensts_w_dtype = events
+            else:
+                if isinstance(events, np.ndarray):
+                    events = [events]
+                evensts_w_dtype = []
+                for evt in events:
+                    if evt.dtype.names is None:
+                        evt_ = np.array(np.array(evt).astype("float64"), dtype=[("time", "float64")])
+                        evensts_w_dtype.append(evt_)
+                    else:
+                        evensts_w_dtype.append(evt)
+        else:
+            evensts_w_dtype = None
+
+        # keep aglobal ref of colorbar
+        self.cb = None
+
         plot_data = dict(
             recordings=recordings,
             segment_index=segment_index,
@@ -210,9 +251,12 @@ class TracesWidget(BaseWidget):
             channel_locations=channel_locations,
             time_range=time_range,
             frame_range=frame_range,
-            times=times,
+            times_in_range=times_in_range,
             layer_keys=layer_keys,
             list_traces=list_traces,
+            events=evensts_w_dtype,
+            events_color=events_color,
+            events_alpha=events_alpha,
             mode=mode,
             cmap=cmap,
             clims=clims,
@@ -240,6 +284,12 @@ class TracesWidget(BaseWidget):
 
         self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs)
 
+        if self.cb is not None:
+            try:
+                self.cb.remove()
+            except:
+                pass
+
         ax = self.ax
         n = len(dp.channel_ids)
         rec0 = dp.recordings[list(dp.recordings.keys())[0]]
@@ -259,7 +309,7 @@ class TracesWidget(BaseWidget):
                 for i, chan_id in enumerate(dp.channel_ids):
                     offset = dp.vspacing * i
                     color = dp.colors[layer_key][chan_id]
-                    ax.plot(dp.times, offset + traces[:, i], color=color)
+                    ax.plot(dp.times_in_range, offset + traces[:, i], color=color)
                 ax.get_lines()[-1].set_label(layer_key)
 
             if dp.show_channel_ids:
@@ -288,7 +338,7 @@ class TracesWidget(BaseWidget):
             im.set_clim(*clim)
 
             if dp.with_colorbar:
-                self.figure.colorbar(im, ax=ax)
+                self.cb = self.figure.colorbar(im, ax=ax)
 
             if dp.show_channel_ids:
                 ax.set_yticks(np.linspace(min_y, max_y, n) + (max_y - min_y) / n * 0.5)
@@ -298,29 +348,33 @@ class TracesWidget(BaseWidget):
                 ax.set_yticks([min_y, max_y])
                 ax.set_yticklabels([min_y, max_y])
 
+        if dp.events is not None:
+            events_segment = dp.events[dp.segment_index]
+            # find events in the time range
+            evt_mask = np.logical_and(
+                events_segment["time"] >= dp.time_range[0], events_segment["time"] < dp.time_range[1]
+            )
+            events_in_range = events_segment[evt_mask]
+            t0 = t1 = 0
+            for evt in events_in_range:
+                if "duration" in events_in_range.dtype.names and evt["duration"] is not None:
+                    t0 = evt["time"]
+                    t1 = evt["time"] + evt["duration"]
+                    ax.axvspan(t0, t1, alpha=dp.events_alpha, color=dp.events_color)
+                else:
+                    ax.axvline(evt["time"], alpha=dp.events_alpha, color=dp.events_color)
+
     def plot_ipywidgets(self, data_plot, **backend_kwargs):
         import matplotlib.pyplot as plt
         import ipywidgets.widgets as widgets
         from IPython.display import display
         import ipywidgets.widgets as W
-        from .utils_ipywidgets import (
-            check_ipywidget_backend,
-            # make_timeseries_controller,
-            # make_channel_controller,
-            # make_scale_controller,
-            TimeSlider,
-            ChannelSelector,
-            ScaleWidget,
-        )
+        from .utils_ipywidgets import check_ipywidget_backend, TimeSlider, ChannelSelector, ScaleWidget, EventSelector
 
         check_ipywidget_backend()
 
         self.next_data_plot = data_plot.copy()
-
         self.recordings = data_plot["recordings"]
-
-        # first layer
-        # rec0 = recordings[data_plot["layer_keys"][0]]
         rec0 = self.rec0 = self.recordings[self.data_plot["layer_keys"][0]]
 
         cm = 1 / 2.54
@@ -332,20 +386,40 @@ class TracesWidget(BaseWidget):
         with plt.ioff():
             output = widgets.Output()
             with output:
-                self.figure, self.ax = plt.subplots(figsize=(0.9 * ratios[1] * width_cm * cm, height_cm * cm))
+                self.figure, self.ax = plt.subplots(
+                    figsize=(0.9 * ratios[1] * width_cm * cm, height_cm * cm), layout="constrained"
+                )
                 plt.show()
 
+        if not self.rec0.has_time_vector(segment_index=data_plot["segment_index"]):
+            times = None
+        else:
+            times = [
+                np.array(self.rec0.get_times(segment_index=segment_index))
+                for segment_index in range(self.rec0.get_num_segments())
+            ]
+
         # some widgets
         self.time_slider = TimeSlider(
             durations=[rec0.get_duration(s) for s in range(rec0.get_num_segments())],
             sampling_frequency=rec0.sampling_frequency,
-            # layout=W.Layout(height="2cm"),
+            time_range=data_plot["time_range"],
+            times=times,
         )
+        # handle times
+        if data_plot["events"] is not None:
+            events_segment = data_plot["events"][data_plot["segment_index"]]
+            self.event_selector = EventSelector(events=events_segment, time_slider=self.time_slider)
+        else:
+            self.event_selector = None
 
-        start_frame = int(data_plot["time_range"][0] * rec0.sampling_frequency)
-        end_frame = int(data_plot["time_range"][1] * rec0.sampling_frequency)
-
-        self.time_slider.value = start_frame, end_frame, data_plot["segment_index"]
+        self.colorbar = W.Checkbox(
+            value=data_plot["with_colorbar"],
+            description="Colorbar",
+            indent=False,
+            layout=W.Layout(width="90%"),
+            align_items="center",
+        )
 
         _layer_keys = data_plot["layer_keys"]
         if len(_layer_keys) > 1:
@@ -364,18 +438,25 @@ class TracesWidget(BaseWidget):
         self.channel_selector = ChannelSelector(self.rec0.channel_ids)
         self.channel_selector.value = list(data_plot["channel_ids"])
 
+        left_sidebar_elements = [
+            W.Label(value="layer"),
+            self.layer_selector,
+            W.Label(value="mode"),
+            self.mode_selector,
+            self.scaler,
+            self.colorbar,
+        ]
         left_sidebar = W.VBox(
-            children=[
-                W.Label(value="layer"),
-                self.layer_selector,
-                W.Label(value="mode"),
-                self.mode_selector,
-                self.scaler,
-                # self.channel_selector,
-            ],
-            layout=W.Layout(width="3.5cm"),
+            children=left_sidebar_elements,
             align_items="center",
         )
+        if self.event_selector is None:
+            right_sidebar = self.channel_selector
+        else:
+            right_sidebar = W.VBox(
+                children=[self.channel_selector, self.event_selector],
+                align_items="center",
+            )
 
         self.return_scaled = data_plot["return_scaled"]
 
@@ -383,7 +464,7 @@ class TracesWidget(BaseWidget):
             center=self.figure.canvas,
             footer=self.time_slider,
             left_sidebar=left_sidebar,
-            right_sidebar=self.channel_selector,
+            right_sidebar=right_sidebar,
             pane_heights=[0, 6, 1],
             pane_widths=ratios,
         )
@@ -399,8 +480,13 @@ class TracesWidget(BaseWidget):
         self.channel_selector.observe(self._retrieve_traces, names="value", type="change")
         # other widgets only refresh
         self.scaler.observe(self._update_plot, names="value", type="change")
+        self.colorbar.observe(self._update_plot, names="value", type="change")
         # map is a special case because needs to check layer also
         self.mode_selector.observe(self._mode_changed, names="value", type="change")
+        # events
+        if self.event_selector is not None:
+            self.event_selector.observe(self._event_changed, names="value", type="change")
+        self.time_slider.segment_selector.observe(self._segment_changed, names="value", type="change")
 
         if backend_kwargs["display"]:
             # self.check_backend()
@@ -422,25 +508,46 @@ class TracesWidget(BaseWidget):
         else:
             self._update_plot()
 
+    def _event_changed(self, change=None):
+        window_size = self._time_range[1] - self._time_range[0]
+        events = self.event_selector.events
+        event_time = events["time"][self.event_selector.value]
+        t_start = event_time - window_size / 2
+        self.time_slider.time_label.value = f"{t_start}"
+        self._retrieve_traces()
+
+    def _segment_changed(self, change=None):
+        # when switching segments, update events if present
+        segment_index = self.time_slider.segment_selector.value
+        if self.event_selector is not None:
+            event_segment = self.data_plot["events"][segment_index]
+            self.event_selector.set_events(event_segment)
+
     def _retrieve_traces(self, change=None):
         channel_ids = np.array(self.channel_selector.value)
 
-        # if self.data_plot["order_channel_by_depth"]:
-        #     order, _ = order_channels_by_depth(self.rec0, channel_ids)
-        # else:
-        #     order = None
-
         start_frame, end_frame, segment_index = self.time_slider.value
-        time_range = np.array([start_frame, end_frame]) / self.rec0.sampling_frequency
+
+        if not self.rec0.has_time_vector(segment_index=segment_index):
+            times = None
+            time_range = np.array([start_frame, end_frame]) / self.rec0.sampling_frequency
+        else:
+            times = self.rec0.get_times(segment_index=segment_index)
+            time_range = np.array([times[start_frame], times[end_frame]])
 
         self._selected_recordings = {k: self.recordings[k] for k in self._get_layers()}
-        times, list_traces, frame_range, channel_ids = _get_trace_list(
-            self._selected_recordings, channel_ids, time_range, segment_index, return_scaled=self.return_scaled
+        times_in_range, list_traces, frame_range, channel_ids = _get_trace_list(
+            self._selected_recordings,
+            channel_ids,
+            time_range,
+            segment_index,
+            return_scaled=self.return_scaled,
+            times=times,
         )
 
         self._channel_ids = channel_ids
         self._list_traces = list_traces
-        self._times = times
+        self._times_in_range = times_in_range
         self._time_range = time_range
         self._frame_range = (start_frame, end_frame)
         self._segment_index = segment_index
@@ -457,7 +564,10 @@ class TracesWidget(BaseWidget):
         data_plot["mode"] = mode
         data_plot["frame_range"] = self._frame_range
         data_plot["time_range"] = self._time_range
-        data_plot["with_colorbar"] = False
+        if self.colorbar.value:
+            data_plot["with_colorbar"] = True
+        else:
+            data_plot["with_colorbar"] = False
         data_plot["recordings"] = self._selected_recordings
         data_plot["add_legend"] = False
 
@@ -474,14 +584,15 @@ class TracesWidget(BaseWidget):
 
         list_traces = [traces * self.scaler.value for traces in self._list_traces]
         data_plot["list_traces"] = list_traces
-        data_plot["times"] = self._times
+        data_plot["times_in_range"] = self._times_in_range
+        data_plot["segment_index"] = self._segment_index
 
         backend_kwargs = {}
         backend_kwargs["ax"] = self.ax
 
         self.ax.clear()
         self.plot_matplotlib(data_plot, **backend_kwargs)
-        self.ax.set_title("")
+        self.ax.set_title(layer_keys[0] if len(layer_keys) == 1 else "ALL")
 
         fig = self.ax.figure
         fig.canvas.draw()
@@ -551,7 +662,7 @@ class TracesWidget(BaseWidget):
         app.exec()
 
 
-def _get_trace_list(recordings, channel_ids, time_range, segment_index, return_scaled=False):
+def _get_trace_list(recordings, channel_ids, time_range, segment_index, return_scaled=False, times=None):
     # function also used in ipywidgets plotter
     k0 = list(recordings.keys())[0]
     rec0 = recordings[k0]
@@ -562,11 +673,15 @@ def _get_trace_list(recordings, channel_ids, time_range, segment_index, return_s
         assert all(
             rec.has_scaled() for rec in recordings.values()
         ), "Some recording layers do not have scaled traces. Use `return_scaled=False`"
-    frame_range = (time_range * fs).astype("int64", copy=False)
-    a_max = rec0.get_num_frames(segment_index=segment_index)
-    frame_range = np.clip(frame_range, 0, a_max)
-    time_range = frame_range / fs
-    times = np.arange(frame_range[0], frame_range[1]) / fs
+    if times is not None:
+        frame_range = np.searchsorted(times, time_range)
+        times = times[frame_range[0] : frame_range[1]]
+    else:
+        frame_range = (time_range * fs).astype("int64", copy=False)
+        a_max = rec0.get_num_frames(segment_index=segment_index)
+        frame_range = np.clip(frame_range, 0, a_max)
+        time_range = frame_range / fs
+        times = np.arange(frame_range[0], frame_range[1]) / fs
 
     list_traces = []
     for rec_name, rec in recordings.items():
diff --git a/src/spikeinterface/widgets/utils_ipywidgets.py b/src/spikeinterface/widgets/utils_ipywidgets.py
index 6d9114050..1032081c7 100644
--- a/src/spikeinterface/widgets/utils_ipywidgets.py
+++ b/src/spikeinterface/widgets/utils_ipywidgets.py
@@ -16,17 +16,25 @@ def check_ipywidget_backend():
 class TimeSlider(W.HBox):
     value = traitlets.Tuple(traitlets.Int(), traitlets.Int(), traitlets.Int())
 
-    def __init__(self, durations, sampling_frequency, time_range=(0, 1.0), **kwargs):
+    def __init__(self, durations, sampling_frequency, time_range, times=None, **kwargs):
         self.num_segments = len(durations)
         self.frame_limits = [int(sampling_frequency * d) for d in durations]
         self.sampling_frequency = sampling_frequency
-        start_frame = int(time_range[0] * sampling_frequency)
-        end_frame = int(time_range[1] * sampling_frequency)
+        self.segment_index = 0
+
+        if times is not None:
+            assert len(times) == len(durations)
+            times_segment = times[self.segment_index]
+            start_frame, end_frame = np.searchsorted(times_segment, time_range)
+            self.times = times
+        else:
+            start_frame = int(time_range[0] * sampling_frequency)
+            end_frame = int(time_range[1] * sampling_frequency)
+            self.times = None
 
         self.frame_range = (start_frame, end_frame)
 
-        self.segment_index = 0
-        self.value = (start_frame, end_frame, self.segment_index)
+        self.value = (int(start_frame), int(end_frame), self.segment_index)
 
         layout = W.Layout(align_items="center", width="2.5cm", height="1.cm")
         but_left = W.Button(description="", disabled=False, button_style="", icon="arrow-left", layout=layout)
@@ -126,7 +134,10 @@ class TimeSlider(W.HBox):
         if new_frame is None and new_time is None:
             start_frame = self.slider.value
         elif new_frame is None:
-            start_frame = int(new_time * self.sampling_frequency)
+            if self.times is not None:
+                start_frame = int(np.searchsorted(self.times[self.segment_index], [new_time])[0])
+            else:
+                start_frame = int(new_time * self.sampling_frequency)
         else:
             start_frame = new_frame
         delta_s = self.window_sizer.value
@@ -139,7 +150,10 @@ class TimeSlider(W.HBox):
 
         end_frame = min(self.frame_limits[self.segment_index], end_frame)
 
-        start_time = start_frame / self.sampling_frequency
+        if self.times is not None:
+            start_time = self.times[self.segment_index][start_frame]
+        else:
+            start_time = start_frame / self.sampling_frequency
 
         if update_label:
             self.time_label.unobserve(self.time_label_changed, names="value", type="change")
@@ -359,3 +373,76 @@ class UnitSelector(W.VBox):
         self.selector.unobserve(self.on_selector_changed, names=["value"], type="change")
         self.selector.value = change["new"]
         self.selector.observe(self.on_selector_changed, names=["value"], type="change")
+
+
+class EventSelector(W.VBox):
+    value = traitlets.Int()
+
+    def __init__(self, events, **kwargs):
+        layout = W.Layout(align_items="center", width="2.5cm", height="1.cm")
+        self.previous_evt = W.Button(description="", disabled=False, button_style="", icon="arrow-left", layout=layout)
+        self.next_evt = W.Button(description="", disabled=False, button_style="", icon="arrow-right", layout=layout)
+
+        self.previous_evt.on_click(self.move_left)
+        self.next_evt.on_click(self.move_right)
+
+        self.events_title = W.Label("Events", layout=W.Layout(justify_content="center"))
+        self.event_options = [f"{i}" for i in range(len(events))]
+        self.events_list = W.Dropdown(
+            options=self.event_options, value=self.event_options[0], disable=False, layout=W.Layout(width="90%")
+        )
+        self.events_list.observe(self.on_selector_changed, names=["value"], type="change")
+        self.set_events(events)
+
+        if "label" in events.dtype.names:
+            self.events_label = W.Label("Evt. index (label)", layout=W.Layout(width="95%"))
+        else:
+            self.events_label = W.Label("Evt. index", layout=W.Layout(width="95%"))
+
+        self.skip_events = W.IntText(
+            value=1,
+            description="",
+            layout=W.Layout(width="1.5cm"),
+        )
+
+        super(W.VBox, self).__init__(
+            children=[
+                self.events_title,
+                self.events_label,
+                self.events_list,
+                W.HBox([self.previous_evt, self.skip_events, self.next_evt]),
+            ],
+            **kwargs,
+        )
+
+        self.observe(self.value_changed, names=["value"], type="change")
+
+    def set_events(self, events):
+        self.events_list.unobserve(self.on_selector_changed, names=["value"], type="change")
+        self.events = events
+        if "label" in self.events.dtype.names:
+            self.event_options = [f"{i} ({self.events['label'][i]})" for i in range(len(self.events))]
+        else:
+            self.event_options = [f"{i}" for i in range(len(self.events))]
+        self.events_list.options = self.event_options
+        self.value = 0
+        self.events_list.observe(self.on_selector_changed, names=["value"], type="change")
+
+    def value_changed(self, change=None):
+        pass
+
+    def on_selector_changed(self, change=None):
+        events_index = self.event_options.index(self.events_list.value)
+        self.value = events_index
+
+    def move_left(self, change=None):
+        events_index = self.value - self.skip_events.value
+        events_index = events_index if events_index >= 0 else 0
+        self.value = events_index
+        self.events_list.value = self.event_options[events_index]
+
+    def move_right(self, change=None):
+        events_index = self.value + self.skip_events.value
+        events_index = events_index if events_index < len(self.events) else len(self.events["time"]) - 1
+        self.value = events_index
+        self.events_list.value = self.event_options[events_index]
diff --git a/src/spikeinterface/widgets/utils_matplotlib.py b/src/spikeinterface/widgets/utils_matplotlib.py
index c370039c9..9eb5f275b 100644
--- a/src/spikeinterface/widgets/utils_matplotlib.py
+++ b/src/spikeinterface/widgets/utils_matplotlib.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -11,7 +12,10 @@ def make_mpl_figure(figure=None, ax=None, axes=None, ncols=None, num_axes=None,
     if figure is not None:
         assert ax is None and axes is None, "figure/ax/axes : only one of then can be not None"
         if num_axes is None:
-            ax = figure.add_subplot(111)
+            if "ipympl" not in matplotlib.get_backend():
+                ax = figure.add_subplot(111)
+            else:
+                ax = figure.add_subplot(111, layout="constrained")
             axes = np.array([[ax]])
         else:
             assert ncols is not None


Test patch:
diff --git a/src/spikeinterface/core/tests/test_template_tools.py b/src/spikeinterface/core/tests/test_template_tools.py
index 6ef826774..970cc3138 100644
--- a/src/spikeinterface/core/tests/test_template_tools.py
+++ b/src/spikeinterface/core/tests/test_template_tools.py
@@ -1,4 +1,5 @@
 import pytest
+import numpy as np
 
 from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
 
@@ -54,10 +55,10 @@ def _get_templates_object_from_sorting_analyzer(sorting_analyzer):
 
 def test_get_template_amplitudes(sorting_analyzer):
     peak_values = get_template_amplitudes(sorting_analyzer)
-    print(peak_values)
     templates = _get_templates_object_from_sorting_analyzer(sorting_analyzer)
-    peak_values = get_template_amplitudes(templates)
-    print(peak_values)
+    peak_values = get_template_amplitudes(templates, abs_value=True)
+    peak_to_peak_values = get_template_amplitudes(templates, mode="peak_to_peak")
+    assert np.all(ptp > p for ptp, p in zip(peak_to_peak_values.values(), peak_values.values()))
 
 
 def test_get_template_extremum_channel(sorting_analyzer):
diff --git a/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py b/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py
index 5b2bec2ed..5c843e7c0 100644
--- a/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py
+++ b/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py
@@ -109,6 +109,34 @@ def test_highpass_spatial_filter_synthetic_data(num_channels, ntr_pad, ntr_tap,
             assert raw_traces.shape == si_filtered.shape
 
 
+@pytest.mark.parametrize("dtype", [np.int16, np.float32, np.float64])
+def test_dtype_stability(dtype):
+    """
+    Check that the dtype of the recording and
+    output data is as expected, as data is cast to float32
+    during filtering.
+    """
+    num_chan = 32
+    si_recording = generate_recording(num_channels=num_chan, durations=[2])
+    si_recording.set_property("gain_to_uV", np.ones(num_chan))
+    si_recording.set_property("offset_to_uV", np.ones(num_chan))
+    si_recording = spre.astype(si_recording, dtype)
+
+    assert si_recording.dtype == dtype
+
+    highpass_spatial_filter = spre.highpass_spatial_filter(si_recording, n_channel_pad=2)
+
+    assert highpass_spatial_filter.dtype == dtype
+
+    filtered_data_unscaled = highpass_spatial_filter.get_traces(return_scaled=False)
+
+    assert filtered_data_unscaled.dtype == dtype
+
+    filtered_data_scaled = highpass_spatial_filter.get_traces(return_scaled=True)
+
+    assert filtered_data_scaled.dtype == np.float32
+
+
 # ----------------------------------------------------------------------------------------------------------------------
 # Test Utils
 # ----------------------------------------------------------------------------------------------------------------------
@@ -125,7 +153,7 @@ def get_ibl_si_data():
     ibl_data = ibl_recording.read(slice(None), slice(None), sync=False)[:, :-1].T  # cut sync channel
 
     si_recording = se.read_spikeglx(local_path, stream_id="imec0.ap")
-    si_recording = spre.scale(si_recording, dtype="float32")
+    si_recording = spre.astype(si_recording, dtype="float32")
 
     return ibl_data, si_recording