Update DriftRasterMapWidget

Author: chrishalcrow

src/spikeinterface/widgets/amplitudes.py

@@ -29,6 +29,10 @@ class AmplitudesWidget(BaseRasterWidget):
         The segment index (or None if mono-segment)
     max_spikes_per_unit : int or None, default: None
         Number of max spikes per unit to display. Use None for all spikes
+    y_lim : tuple or None, default: None
+        The min and max depth to display, if None (min and max of the amplitudes).
+    scatter_decimate : int, default: 1
+        If > 1, the scatter points are decimated.
     hide_unit_selector : bool, default: False
         If True the unit selector is not displayed
         (sortingview backend)
@@ -49,6 +53,8 @@ def __init__(
         unit_colors=None,
         segment_index=None,
         max_spikes_per_unit=None,
+        y_lim=None,
+        scatter_decimate=1,
         hide_unit_selector=False,
         plot_histograms=False,
         bins=None,
@@ -108,6 +114,8 @@ def __init__(
             hide_unit_selector=hide_unit_selector,
             plot_legend=plot_legend,
             y_label="Amplitude",
+            y_lim=y_lim,
+            scatter_decimate=scatter_decimate,
         )
 
         BaseRasterWidget.__init__(self, **plot_data, backend=backend, **backend_kwargs)

src/spikeinterface/widgets/motion.py

@@ -6,6 +6,7 @@
 
 from spikeinterface.core import BaseRecording, SortingAnalyzer
 from spikeinterface.sortingcomponents.motion import Motion
+from .rasters import BaseRasterWidget
 
 
 class MotionWidget(BaseWidget):
@@ -97,7 +98,7 @@ def plot_matplotlib(self, data_plot, **backend_kwargs):
             ax.set_ylabel("Depth [um]")
 
 
-class DriftRasterMapWidget(BaseWidget):
+class DriftRasterMapWidget(BaseRasterWidget):
     """
     Plot the drift raster map from peaks or a SortingAnalyzer.
     The drift raster map is a scatter plot of the estimated peak depth vs time and it is
@@ -200,56 +201,14 @@ def __init__(
             if peak_amplitudes is not None:
                 peak_amplitudes = peak_amplitudes[peak_mask]
 
-        plot_data = dict(
-            peaks=peaks,
-            peak_locations=peak_locations,
-            peak_amplitudes=peak_amplitudes,
-            direction=direction,
-            sampling_frequency=sampling_frequency,
-            segment_index=segment_index,
-            depth_lim=depth_lim,
-            color_amplitude=color_amplitude,
-            color=color,
-            scatter_decimate=scatter_decimate,
-            cmap=cmap,
-            clim=clim,
-            alpha=alpha,
-            recording=recording,
-        )
-        BaseWidget.__init__(self, plot_data, backend=backend, **backend_kwargs)
-
-    def plot_matplotlib(self, data_plot, **backend_kwargs):
-        import matplotlib.pyplot as plt
-        from matplotlib.colors import Normalize
-        from .utils_matplotlib import make_mpl_figure
+        from matplotlib.pyplot import colormaps
 
-        from spikeinterface.sortingcomponents.motion import correct_motion_on_peaks
-
-        dp = to_attr(data_plot)
-
-        assert backend_kwargs["axes"] is None, "axes argument is not allowed in DriftRasterMapWidget. Use ax instead."
-
-        self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs)
-
-        if dp.recording is None:
-            peak_times = dp.peaks["sample_index"] / dp.sampling_frequency
-        else:
-            peak_times = dp.recording.sample_index_to_time(dp.peaks["sample_index"], segment_index=dp.segment_index)
-
-        peak_locs = dp.peak_locations[dp.direction]
-        if dp.scatter_decimate is not None:
-            peak_times = peak_times[:: dp.scatter_decimate]
-            peak_locs = peak_locs[:: dp.scatter_decimate]
-
-        if dp.color_amplitude:
-            amps = dp.peak_amplitudes
+        if color_amplitude:
+            amps = peak_amplitudes
             amps_abs = np.abs(amps)
             q_95 = np.quantile(amps_abs, 0.95)
-            if dp.scatter_decimate is not None:
-                amps = amps[:: dp.scatter_decimate]
-                amps_abs = amps_abs[:: dp.scatter_decimate]
-            cmap = plt.colormaps[dp.cmap]
-            if dp.clim is None:
+            cmap = colormaps[cmap]
+            if clim is None:
                 amps = amps_abs
                 amps /= q_95
                 c = cmap(amps)
@@ -259,17 +218,26 @@ def plot_matplotlib(self, data_plot, **backend_kwargs):
             color_kwargs = dict(
                 color=None,
                 c=c,
-                alpha=dp.alpha,
+                alpha=alpha,
             )
         else:
-            color_kwargs = dict(color=dp.color, c=None, alpha=dp.alpha)
-
-        self.ax.scatter(peak_times, peak_locs, s=1, **color_kwargs)
-        if dp.depth_lim is not None:
-            self.ax.set_ylim(*dp.depth_lim)
-        self.ax.set_title("Peak depth")
-        self.ax.set_xlabel("Times [s]")
-        self.ax.set_ylabel("Depth [$\\mu$m]")
+            color_kwargs = dict(color=color, c=None, alpha=alpha)
+
+        # convert data into format that `BaseRasterWidget` can take it in
+        spike_train_data = {0: peaks["sample_index"] / sampling_frequency}
+        y_axis_data = {0: peak_locations[direction]}
+
+        plot_data = dict(
+            spike_train_data=spike_train_data,
+            y_axis_data=y_axis_data,
+            y_lim=depth_lim,
+            color_kwargs=color_kwargs,
+            scatter_decimate=scatter_decimate,
+            title="Peak depth",
+            y_label="Depth [um]",
+        )
+
+        BaseRasterWidget.__init__(self, **plot_data, backend=backend, **backend_kwargs)
 
 
 class MotionInfoWidget(BaseWidget):

src/spikeinterface/widgets/rasters.py

@@ -4,7 +4,7 @@
 from warnings import warn
 
 from .base import BaseWidget, to_attr, default_backend_kwargs
-from . import get_some_colors
+from .utils import get_some_colors
 
 
 class BaseRasterWidget(BaseWidget):
@@ -188,7 +188,7 @@ def plot_ipywidgets(self, data_plot, **backend_kwargs):
 
         import ipywidgets.widgets as W
         from IPython.display import display
-        from spikeinterface.widgets.utils_ipywidgets import check_ipywidget_backend, UnitSelector
+        from .utils_ipywidgets import check_ipywidget_backend, UnitSelector
 
         check_ipywidget_backend()