Topological distribution widget (#3142 followup)

pyproject.toml

@@ -179,6 +179,7 @@ widgets = [
     "matplotlib>=3.6",
     "ipympl",
     "ipywidgets",
+    "seaborn>=0.13.0",
     "figpack",
     "figpack-spike-sorting",
     "distinctipy",

src/spikeinterface/core/generate.py

@@ -8,7 +8,7 @@
 from .basesorting import SpikeVectorSortingSegment
 from .numpyextractors import NumpySorting
 
-from probeinterface import generate_linear_probe, generate_multi_columns_probe
+from probeinterface import Probe, generate_linear_probe, generate_multi_columns_probe
 
 from spikeinterface.core import BaseRecording, BaseRecordingSegment, BaseSorting, BaseSortingSegment
 from .core_tools import define_function_from_class, ms_to_samples, _ensure_seed

src/spikeinterface/generation/session_displacement_generator.py

@@ -196,7 +196,6 @@ def generate_session_displacement_recordings(
     amplitude_scaling_order = None
 
     for rec_idx, (shift, duration) in enumerate(zip(recording_shifts, recording_durations)):
-
         displacement_vector, displacement_unit_factor = _get_inter_session_displacements(
             shift,
             non_rigid_gradient[rec_idx] if isinstance(non_rigid_gradient, list) else non_rigid_gradient,
@@ -238,7 +237,6 @@ def generate_session_displacement_recordings(
         )
 
         if recording_amplitude_scalings is not None:
-
             if rec_idx == 0:
                 amplitude_scaling_order = _get_amplitude_scaling_order(
                     templates_array_moved,
@@ -399,13 +397,12 @@ def _check_generate_session_displacement_arguments(
         raise ValueError("Each record entry for `recording_shifts` must have two elements, the x and y shift.")
 
     if recording_amplitude_scalings is not None:
-
         keys = recording_amplitude_scalings.keys()
-        if not "method" in keys or not "scalings" in keys:
+        if "method" not in keys or "scalings" not in keys:
             raise ValueError("`recording_amplitude_scalings` must be a dict with keys `method` and `scalings`.")
 
         allowed_methods = ["by_passed_order", "by_amplitude_and_firing_rate", "by_firing_rate"]
-        if not recording_amplitude_scalings["method"] in allowed_methods:
+        if recording_amplitude_scalings["method"] not in allowed_methods:
             raise ValueError(f"`recording_amplitude_scalings` must be one of {allowed_methods}")
 
         rec_scalings = recording_amplitude_scalings["scalings"]
@@ -420,7 +417,7 @@ def _check_generate_session_displacement_arguments(
 
     if isinstance(non_rigid_gradient, list):
         if not len(non_rigid_gradient) == expected_num_recs:
-            raise ValueError("If `non_rigid_gradient` is a list, it must " "contain one value for each recording.")
+            raise ValueError("If `non_rigid_gradient` is a list, it must contain one value for each recording.")
 
 
 def _update_kwargs_for_extended_units(

src/spikeinterface/generation/tests/test_session_displacement_generator.py

@@ -76,7 +76,6 @@ def test_x_y_rigid_shifts_are_properly_set(self, options):
         locations_1 = extra_outputs["unit_locations"][0]
 
         for rec_idx in range(1, 4):
-
             shifts = recording_shifts[rec_idx]
 
             assert np.array_equal(
@@ -91,14 +90,12 @@ def test_x_y_rigid_shifts_are_properly_set(self, options):
         min_channel_loc = output_recordings[0].get_channel_locations()[0, 1]
         max_channel_loc = output_recordings[0].get_channel_locations()[-1, 1]
         for unit_idx in range(num_units):
-
             start_pos = self._get_peak_chan_loc_in_um(
                 extra_outputs["templates_array_moved"][0][unit_idx],
                 options["y_bin_um"],
             )
 
             for rec_idx in range(1, options["num_recs"]):
-
                 new_pos = self._get_peak_chan_loc_in_um(
                     extra_outputs["templates_array_moved"][rec_idx][unit_idx], options["y_bin_um"]
                 )
@@ -156,7 +153,6 @@ def test_spike_times_and_firing_rates_across_recordings(self, options):
 
         for unit_idx in range(options["kwargs"]["num_units"]):
             for rec_idx in range(1, options["num_recs"]):
-
                 # Exact spike times are not preserved when seed is None
                 assert np.array_equal(
                     output_sortings_same[0].get_unit_spike_train(str(unit_idx)),
@@ -222,7 +218,6 @@ def test_x_y_shift_non_rigid(self, options, dim_idx):
         # w.r.t the first recording), check the rigid and nonrigid shifts
         # are as expected.
         for rec_idx in range(1, options["num_recs"]):
-
             shift = options["kwargs"]["recording_shifts"][rec_idx][dim_idx]
 
             # Get the rigid shift between the first recording and this shifted recording
@@ -437,7 +432,6 @@ def test_shift_units_outside_probe(self, options):
         num_extended_units = num_units * 3
 
         for ses_idx in range(num_sessions):
-
             # There are 3x the number of units when new units are created
             # (one new set above, and one new set below the probe).
             for key in ["unit_locations", "templates_array_moved", "firing_rates"]:

src/spikeinterface/widgets/unit_spatial.py

@@ -0,0 +1,232 @@
+from typing import Any, Literal
+
+import numpy as np
+from numpy.typing import ArrayLike
+from probeinterface import Probe
+from spikeinterface.core import SortingAnalyzer
+from .base import BaseWidget, to_attr
+
+
+class UnitSpatialDistributionsWidget(BaseWidget):
+    """
+    Plot unit spatial distributions as a histogram or kernel density estimate.
+
+    Parameters
+    ----------
+    sorting_analyzer : SortingAnalyzer
+        Analyzer with the unit_locations extension already computed.
+    probe : Probe or None, default: None
+        Probe geometry to display. If None, use the analyzer's stored probe.
+    depth_axis : int, default: 1
+        Coordinate along the probe representing depth: 0 for x or 1 for y.
+    bins : int, tuple or None, default: None
+        Bin count for both axes, or a pair of counts or bin-edge arrays for x
+        and y. If None, use 75 um per depth bin.
+    cmap : str or Colormap, default: "viridis"
+        Colormap name or object for the spatial histogram or KDE.
+    kde : bool, default: False
+        If True, show a two-dimensional kernel density estimate instead of
+        the spatial histogram.
+    depth_hist : bool, default: True
+        Show a depth histogram to the right for y depth or below for x depth.
+    groups : array-like or None, default: None
+        Group label for each unit, in analyzer unit order. Colors the depth
+        histogram by group, e.g. ["good", "good", "bad", "bad"] for four units.
+        If None, use one group for all units.
+        histogram by group. If None, use one group for all units.
+    kde_kwargs : dict or None, default: None
+        Keyword arguments passed to seaborn.kdeplot, overriding the defaults.
+        Only used for KDE.
+    show_depth_hist_axis : bool, default: True
+        Show the count-axis on the depth histogram.
+        Only used when depth_hist is True.
+    """
+
+    def __init__(
+        self,
+        sorting_analyzer: SortingAnalyzer,
+        probe: Probe | None = None,
+        depth_axis: Literal[0, 1] = 1,
+        bins: int | tuple | None = None,
+        cmap="viridis",
+        kde: bool = False,
+        depth_hist: bool = True,
+        groups: ArrayLike | None = None,
+        kde_kwargs: dict[str, Any] | None = None,
+        backend: str | None = None,
+        show_depth_hist_axis: bool = True,
+        **backend_kwargs: Any,
+    ) -> None:
+        sorting_analyzer = self.ensure_sorting_analyzer(sorting_analyzer)
+
+        self.check_extensions(sorting_analyzer, "unit_locations")
+        ulc = sorting_analyzer.get_extension("unit_locations")
+        unit_locations = ulc.get_data(outputs="numpy")
+
+        if depth_axis not in (0, 1):
+            raise ValueError("depth_axis must be 0 (x) or 1 (y)")
+
+        x, y = unit_locations[:, 0], unit_locations[:, 1]
+
+        if probe is None:
+            probe = sorting_analyzer.get_probe()
+        elif not isinstance(probe, Probe):
+            raise TypeError("probe must be a Probe or None")
+
+        plot_data = dict(
+            probe=probe,
+            x=x,
+            y=y,
+            depth_axis=depth_axis,
+            bins=bins,
+            kde=kde,
+            cmap=cmap,
+            depth_hist=depth_hist,
+            show_depth_hist_axis=show_depth_hist_axis,
+            groups=groups,
+            kde_kwargs=kde_kwargs,
+        )
+
+        BaseWidget.__init__(self, plot_data, backend=backend, **backend_kwargs)
+
+    def plot_matplotlib(self, data_plot: dict[str, Any], **backend_kwargs: Any) -> None:
+        import matplotlib.patches as patches
+        import matplotlib.path as mpl_paths
+        from probeinterface.plotting import get_auto_lims
+        from seaborn import color_palette, kdeplot, histplot
+        from .utils_matplotlib import make_mpl_figure
+
+        dp = to_attr(data_plot)
+        xrange, yrange, _ = get_auto_lims(dp.probe, margin=0)
+        bins = dp.bins
+
+        if bins is None:
+            # if bins not supplied, divive the probe into 75 um bins
+            depth_range = (xrange, yrange)[dp.depth_axis]
+            bins = [1, 1]
+            bins[dp.depth_axis] = max(1, int(np.ceil(np.ptp(depth_range) / 75.0)))
+
+        # Explicitly create the depth bin edges so they can be
+        # shared between the spatial and marginal histograms.
+        bin_specs = (bins, bins) if np.isscalar(bins) else bins
+        bins = (
+            np.histogram_bin_edges(dp.x, bins=bin_specs[0], range=xrange),
+            np.histogram_bin_edges(dp.y, bins=bin_specs[1], range=yrange),
+        )
+        cmap = color_palette(dp.cmap, as_cmap=True) if type(dp.cmap) is str else dp.cmap
+
+        self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs)
+
+        ax = self.ax
+        if dp.depth_hist and backend_kwargs.get("ax") is None and backend_kwargs.get("axes") is None:
+            # Keep the marginal histogram inside the default figure bounds.
+            # Fixed layout proportions reserve space for the depth histogram.
+            position = ax.get_position()
+            if dp.depth_axis == 1:
+                ax.set_position([position.x0, position.y0, position.width / 2.5, position.height])
+            else:
+                height = position.height / 1.6
+                ax.set_position([position.x0, position.y0 + 0.6 * height, position.width, height])
+
+        # Set the clipping region for the spatial plot based on the probe dimensions
+        probe_outline = None
+        if dp.probe.probe_planar_contour is not None:
+            probe_outline_path = mpl_paths.Path(dp.probe.probe_planar_contour)
+            probe_outline = patches.PathPatch(probe_outline_path, facecolor="none", edgecolor="none")
+            ax.add_patch(probe_outline)
+
+        # Create the spatial plot
+        if dp.kde is not True:
+            hist, xedges, yedges = np.histogram2d(dp.x, dp.y, bins=bins, range=[xrange, yrange])
+            spatial_plot = ax.pcolormesh(xedges, yedges, hist.T, cmap=cmap)
+        else:
+            kde_kwargs = dict(levels=100, thresh=0, fill=True, bw_adjust=0.1)
+            if dp.kde_kwargs is not None:
+                kde_kwargs.update(dp.kde_kwargs)
+
+            data = dict(x=dp.x, y=dp.y)
+            kde_background = ax.add_patch(
+                patches.Rectangle(
+                    [xrange[0], yrange[0]],
+                    np.diff(xrange).squeeze(),
+                    np.diff(yrange).squeeze(),
+                    facecolor=cmap(0.0),
+                    fill=True,
+                )
+            )
+            if probe_outline is not None:
+                kde_background.set_clip_path(probe_outline)
+
+            kdeplot(data, x="x", y="y", clip=[xrange, yrange], cmap=cmap, ax=ax, **kde_kwargs)
+            spatial_plot = ax.collections[0]
+            ax.set_xlabel(None)
+            ax.set_ylabel(None)
+
+        if probe_outline is not None:
+            spatial_plot.set_clip_path(probe_outline)
+
+        # Create the histogram plot
+        xlim, ylim, _ = get_auto_lims(dp.probe, margin=10)
+        ax.set_xlim(*xlim)
+        ax.set_ylim(*ylim)
+        ax.spines["top"].set_visible(False)
+        ax.spines["bottom"].set_visible(False)
+        ax.spines["right"].set_visible(False)
+
+        if dp.depth_axis == 1:
+            ax.set_xticks([])
+            ax.set_xlabel("")
+            ax.set_ylabel("Depth (um)")
+        else:
+            ax.set_yticks([])
+            ax.set_xlabel("Depth (um)")
+            ax.set_ylabel("")
+
+        if dp.depth_hist is True:
+            bbox = ax.get_window_extent()
+            hist_height = 1.5 * bbox.width
+
+            if dp.depth_axis == 1:
+                ax_hist = ax.inset_axes([1, 0, hist_height / bbox.width, 1])
+                depth_key = "y"
+                depth_values = dp.y
+            else:
+                ax_hist = ax.inset_axes([0, -0.6, 1, 0.4])
+                depth_key = "x"
+                depth_values = dp.x
+
+            data = {depth_key: depth_values}
+            data["group"] = np.ones(dp.y.size) if dp.groups is None else dp.groups
+            palette = color_palette("bright", n_colors=1 if dp.groups is None else np.unique(dp.groups).size)
+
+            histplot(
+                data=data,
+                **{depth_key: depth_key},
+                hue="group",
+                bins=bins[dp.depth_axis],
+                binrange=(xrange, yrange)[dp.depth_axis],
+                palette=palette,
+                ax=ax_hist,
+                legend=False,
+            )
+
+            if dp.show_depth_hist_axis:
+                # Show axis for counts on histogram
+                if dp.depth_axis == 1:
+                    ax_hist.yaxis.set_visible(False)
+                    ax_hist.set_xlabel("Count")
+                    visible_spine = "bottom"
+                else:
+                    ax_hist.xaxis.set_visible(False)
+                    ax_hist.set_ylabel("Count")
+                    visible_spine = "left"
+
+                for name, spine in ax_hist.spines.items():
+                    spine.set_visible(name == visible_spine)
+            else:
+                ax_hist.axis("off")
+
+            if dp.depth_axis == 1:
+                ax_hist.set_ylim(*ylim)
+            else:
+                ax_hist.set_xlim(*xlim)

src/spikeinterface/widgets/widget_list.py

@@ -29,6 +29,7 @@
 from .unit_locations import UnitLocationsWidget
 from .unit_presence import UnitPresenceWidget
 from .unit_probe_map import UnitProbeMapWidget
+from .unit_spatial import UnitSpatialDistributionsWidget
 from .unit_summary import UnitSummaryWidget
 from .unit_templates import UnitTemplatesWidget
 from .unit_waveforms_density_map import UnitWaveformDensityMapWidget
@@ -75,6 +76,7 @@
     UnitLocationsWidget,
     UnitPresenceWidget,
     UnitProbeMapWidget,
+    UnitSpatialDistributionsWidget,
     UnitSummaryWidget,
     UnitTemplatesWidget,
     UnitWaveformDensityMapWidget,
@@ -161,6 +163,7 @@
 plot_unit_locations = UnitLocationsWidget
 plot_unit_presence = UnitPresenceWidget
 plot_unit_probe_map = UnitProbeMapWidget
+plot_unit_spatial_distribution = UnitSpatialDistributionsWidget
 plot_unit_summary = UnitSummaryWidget
 plot_unit_templates = UnitTemplatesWidget
 plot_unit_waveforms_density_map = UnitWaveformDensityMapWidget