Merge branch 'main' into fix_chris_bug

Author: h-mayorquin

src/spikeinterface/core/node_pipeline.py

@@ -376,6 +376,7 @@ def __init__(
         parents: Optional[list[PipelineNode]] = None,
         return_output: bool = False,
         radius_um: float = 100.0,
+        sparsity_mask: np.ndarray = None,
     ):
         """
         Extract sparse waveforms from a recording. The strategy in this specific node is to reshape the waveforms
@@ -400,6 +401,11 @@ def __init__(
             Pass parents nodes to perform a previous computation
         return_output : bool, default: False
             Whether or not the output of the node is returned by the pipeline
+        radius_um : float, default: 100.0
+            The radius to determine the neighborhood of channels to extract waveforms from.
+        sparsity_mask : np.ndarray, default: None
+            Optional mask to specify the sparsity of the waveforms. If provided, it should be a boolean array of shape
+            (num_channels, num_channels) where True indicates that the channel is active in the neighborhood.
         """
         WaveformsNode.__init__(
             self,
@@ -410,10 +416,15 @@ def __init__(
             return_output=return_output,
         )
 
-        self.radius_um = radius_um
         self.contact_locations = recording.get_channel_locations()
         self.channel_distance = get_channel_distances(recording)
-        self.neighbours_mask = self.channel_distance <= radius_um
+
+        if sparsity_mask is not None:
+            self.neighbours_mask = sparsity_mask
+            self.radius_um = None
+        else:
+            self.radius_um = radius_um
+            self.neighbours_mask = self.channel_distance <= radius_um
         self.max_num_chans = np.max(np.sum(self.neighbours_mask, axis=1))
 
     def get_trace_margin(self):

src/spikeinterface/postprocessing/spike_amplitudes.py

@@ -31,7 +31,6 @@ class ComputeSpikeAmplitudes(AnalyzerExtension):
     spike_retriver_kwargs : dict
         A dictionary to control the behavior for getting the maximum channel for each spike
         This dictionary contains:
-
           * channel_from_template: bool, default: True
               For each spike is the maximum channel computed from template or re estimated at every spikes
               channel_from_template = True is old behavior but less acurate

src/spikeinterface/sortingcomponents/clustering/method_list.py

@@ -24,3 +24,11 @@
     "circus": CircusClustering,
     "tdc_clustering": TdcClustering,
 }
+
+try:
+    # Kilosort licence (GPL 3) is forcing us to make and use an external package
+    from spikeinterface_kilosort_components import KiloSortClustering
+
+    clustering_methods["kilosort_clustering"] = KiloSortClustering
+except ImportError:
+    pass

src/spikeinterface/widgets/all_amplitudes_distributions.py

@@ -19,8 +19,9 @@ class AllAmplitudesDistributionsWidget(BaseWidget):
         The SortingAnalyzer
     unit_ids : list
         List of unit ids, default None
-    unit_colors : None or dict
-        Dict of colors with key : unit, value : color, default None
+    unit_colors : dict | None, default: None
+        Dict of colors with unit ids as keys and colors as values. Colors can be any type accepted
+        by matplotlib. If None, default colors are chosen using the `get_some_colors` function.
     """
 
     def __init__(

src/spikeinterface/widgets/amplitudes.py

@@ -3,13 +3,16 @@
 import numpy as np
 from warnings import warn
 
+from .rasters import BaseRasterWidget
 from .base import BaseWidget, to_attr
 from .utils import get_some_colors
 
 from ..core.sortinganalyzer import SortingAnalyzer
 
+from spikeinterface.core import SortingAnalyzer
 
-class AmplitudesWidget(BaseWidget):
+
+class AmplitudesWidget(BaseRasterWidget):
     """
     Plots spike amplitudes
 
@@ -19,10 +22,17 @@ class AmplitudesWidget(BaseWidget):
         The input waveform extractor
     unit_ids : list or None, default: None
         List of unit ids
+    unit_colors : dict | None, default: None
+        Dict of colors with unit ids as keys and colors as values. Colors can be any type accepted
+        by matplotlib. If None, default colors are chosen using the `get_some_colors` function.
     segment_index : int or None, default: None
         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 equal to n, each nth spike is kept for plotting.
     hide_unit_selector : bool, default: False
         If True the unit selector is not displayed
         (sortingview backend)
@@ -31,7 +41,7 @@ class AmplitudesWidget(BaseWidget):
         (matplotlib backend)
     bins : int or None, default: None
         If plot_histogram is True, the number of bins for the amplitude histogram.
-        If None this is automatically adjusted
+        If None, uses 100 bins.
     plot_legend : bool, default: True
         True includes legend in plot
     """
@@ -43,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,
@@ -61,25 +73,21 @@ def __init__(
         if unit_ids is None:
             unit_ids = sorting.unit_ids
 
-        if unit_colors is None:
-            unit_colors = get_some_colors(sorting.unit_ids)
-
         if sorting.get_num_segments() > 1:
             if segment_index is None:
-                warn("More than one segment available! Using segment_index 0")
+                warn("More than one segment available! Using `segment_index = 0`.")
                 segment_index = 0
         else:
             segment_index = 0
+
         amplitudes_segment = amplitudes[segment_index]
         total_duration = sorting_analyzer.get_num_samples(segment_index) / sorting_analyzer.sampling_frequency
 
-        spiketrains_segment = {}
-        for i, unit_id in enumerate(sorting.unit_ids):
-            times = sorting.get_unit_spike_train(unit_id, segment_index=segment_index)
-            times = times / sorting.get_sampling_frequency()
-            spiketrains_segment[unit_id] = times
+        all_spiketrains = {
+            unit_id: sorting.get_unit_spike_train(unit_id, segment_index=segment_index, return_times=True)
+            for unit_id in sorting.unit_ids
+        }
 
-        all_spiketrains = spiketrains_segment
         all_amplitudes = amplitudes_segment
         if max_spikes_per_unit is not None:
             spiketrains_to_plot = dict()
@@ -101,163 +109,21 @@ def __init__(
             bins = 100
 
         plot_data = dict(
-            sorting_analyzer=sorting_analyzer,
-            amplitudes=amplitudes_to_plot,
-            unit_ids=unit_ids,
+            spike_train_data=spiketrains_to_plot,
+            y_axis_data=amplitudes_to_plot,
             unit_colors=unit_colors,
-            spiketrains=spiketrains_to_plot,
-            total_duration=total_duration,
             plot_histograms=plot_histograms,
             bins=bins,
+            total_duration=total_duration,
+            unit_ids=unit_ids,
             hide_unit_selector=hide_unit_selector,
             plot_legend=plot_legend,
+            y_label="Amplitude",
+            y_lim=y_lim,
+            scatter_decimate=scatter_decimate,
         )
 
-        BaseWidget.__init__(self, plot_data, backend=backend, **backend_kwargs)
-
-    def plot_matplotlib(self, data_plot, **backend_kwargs):
-        import matplotlib.pyplot as plt
-        from .utils_matplotlib import make_mpl_figure
-
-        dp = to_attr(data_plot)
-
-        if backend_kwargs["axes"] is not None:
-            axes = backend_kwargs["axes"]
-            if dp.plot_histograms:
-                assert np.asarray(axes).size == 2
-            else:
-                assert np.asarray(axes).size == 1
-        elif backend_kwargs["ax"] is not None:
-            assert not dp.plot_histograms
-        else:
-            if dp.plot_histograms:
-                backend_kwargs["num_axes"] = 2
-                backend_kwargs["ncols"] = 2
-            else:
-                backend_kwargs["num_axes"] = None
-
-        self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs)
-
-        scatter_ax = self.axes.flatten()[0]
-
-        for unit_id in dp.unit_ids:
-            spiketrains = dp.spiketrains[unit_id]
-            amps = dp.amplitudes[unit_id]
-            scatter_ax.scatter(spiketrains, amps, color=dp.unit_colors[unit_id], s=3, alpha=1, label=unit_id)
-
-            if dp.plot_histograms:
-                if dp.bins is None:
-                    bins = int(len(spiketrains) / 30)
-                else:
-                    bins = dp.bins
-                ax_hist = self.axes.flatten()[1]
-                # this is super slow, using plot and np.histogram is really much faster (and nicer!)
-                # ax_hist.hist(amps, bins=bins, orientation="horizontal", color=dp.unit_colors[unit_id], alpha=0.8)
-                count, bins = np.histogram(amps, bins=bins)
-                ax_hist.plot(count, bins[:-1], color=dp.unit_colors[unit_id], alpha=0.8)
-
-        if dp.plot_histograms:
-            ax_hist = self.axes.flatten()[1]
-            ax_hist.set_ylim(scatter_ax.get_ylim())
-            ax_hist.axis("off")
-            # self.figure.tight_layout()
-
-        if dp.plot_legend:
-            if hasattr(self, "legend") and self.legend is not None:
-                self.legend.remove()
-            self.legend = self.figure.legend(
-                loc="upper center", bbox_to_anchor=(0.5, 1.0), ncol=5, fancybox=True, shadow=True
-            )
-
-        scatter_ax.set_xlim(0, dp.total_duration)
-        scatter_ax.set_xlabel("Times [s]")
-        scatter_ax.set_ylabel(f"Amplitude")
-        scatter_ax.spines["top"].set_visible(False)
-        scatter_ax.spines["right"].set_visible(False)
-        self.figure.subplots_adjust(bottom=0.1, top=0.9, left=0.1)
-
-    def plot_ipywidgets(self, data_plot, **backend_kwargs):
-        import matplotlib.pyplot as plt
-
-        # import ipywidgets.widgets as widgets
-        import ipywidgets.widgets as W
-        from IPython.display import display
-        from .utils_ipywidgets import check_ipywidget_backend, UnitSelector
-
-        check_ipywidget_backend()
-
-        self.next_data_plot = data_plot.copy()
-
-        cm = 1 / 2.54
-        analyzer = data_plot["sorting_analyzer"]
-
-        width_cm = backend_kwargs["width_cm"]
-        height_cm = backend_kwargs["height_cm"]
-
-        ratios = [0.15, 0.85]
-
-        with plt.ioff():
-            output = W.Output()
-            with output:
-                self.figure = plt.figure(figsize=((ratios[1] * width_cm) * cm, height_cm * cm))
-                plt.show()
-
-        self.unit_selector = UnitSelector(analyzer.unit_ids)
-        self.unit_selector.value = list(analyzer.unit_ids)[:1]
-
-        self.checkbox_histograms = W.Checkbox(
-            value=data_plot["plot_histograms"],
-            description="hist",
-        )
-
-        left_sidebar = W.VBox(
-            children=[
-                self.unit_selector,
-                self.checkbox_histograms,
-            ],
-            layout=W.Layout(align_items="center", width="100%", height="100%"),
-        )
-
-        self.widget = W.AppLayout(
-            center=self.figure.canvas,
-            left_sidebar=left_sidebar,
-            pane_widths=ratios + [0],
-        )
-
-        # a first update
-        self._full_update_plot()
-
-        self.unit_selector.observe(self._update_plot, names="value", type="change")
-        self.checkbox_histograms.observe(self._full_update_plot, names="value", type="change")
-
-        if backend_kwargs["display"]:
-            display(self.widget)
-
-    def _full_update_plot(self, change=None):
-        self.figure.clear()
-        data_plot = self.next_data_plot
-        data_plot["unit_ids"] = self.unit_selector.value
-        data_plot["plot_histograms"] = self.checkbox_histograms.value
-        data_plot["plot_legend"] = False
-
-        backend_kwargs = dict(figure=self.figure, axes=None, ax=None)
-        self.plot_matplotlib(data_plot, **backend_kwargs)
-        self._update_plot()
-
-    def _update_plot(self, change=None):
-        for ax in self.axes.flatten():
-            ax.clear()
-
-        data_plot = self.next_data_plot
-        data_plot["unit_ids"] = self.unit_selector.value
-        data_plot["plot_histograms"] = self.checkbox_histograms.value
-        data_plot["plot_legend"] = False
-
-        backend_kwargs = dict(figure=None, axes=self.axes, ax=None)
-        self.plot_matplotlib(data_plot, **backend_kwargs)
-
-        self.figure.canvas.draw()
-        self.figure.canvas.flush_events()
+        BaseRasterWidget.__init__(self, **plot_data, backend=backend, **backend_kwargs)
 
     def plot_sortingview(self, data_plot, **backend_kwargs):
         import sortingview.views as vv
@@ -270,8 +136,8 @@ def plot_sortingview(self, data_plot, **backend_kwargs):
         sa_items = [
             vv.SpikeAmplitudesItem(
                 unit_id=u,
-                spike_times_sec=dp.spiketrains[u].astype("float32"),
-                spike_amplitudes=dp.amplitudes[u].astype("float32"),
+                spike_times_sec=dp.spike_train_data[u].astype("float32"),
+                spike_amplitudes=dp.y_axis_data[u].astype("float32"),
             )
             for u in unit_ids
         ]

src/spikeinterface/widgets/crosscorrelograms.py

@@ -31,8 +31,9 @@ class CrossCorrelogramsWidget(BaseWidget):
         this argument is ignored
     hide_unit_selector : bool, default: False
         For sortingview backend, if True the unit selector is not displayed
-    unit_colors : dict or None, default: None
-        If given, a dictionary with unit ids as keys and colors as values
+    unit_colors : dict | None, default: None
+        Dict of colors with unit ids as keys and colors as values. Colors can be any type accepted
+        by matplotlib. If None, default colors are chosen using the `get_some_colors` function.
     """
 
     def __init__(

src/spikeinterface/widgets/metrics.py

@@ -24,8 +24,9 @@ class MetricsBaseWidget(BaseWidget):
         If given, a list of quality metrics to skip, default: None
     include_metrics: list or None, default: None
         If given, a list of quality metrics to include, default: None
-    unit_colors :  dict or None, default: None
-        If given, a dictionary with unit ids as keys and colors as values
+    unit_colors : dict | None, default: None
+        Dict of colors with unit ids as keys and colors as values. Colors can be any type accepted
+        by matplotlib. If None, default colors are chosen using the `get_some_colors` function.
     hide_unit_selector : bool, default: False
         For sortingview backend, if True the unit selector is not displayed
     include_metrics_data : bool, default: True