diff --git a/pyproject.toml b/pyproject.toml index e81fba5912..e5aa5765de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,6 +179,7 @@ widgets = [ "matplotlib>=3.6", "ipympl", "ipywidgets", + "seaborn>=0.13.0", "figpack", "figpack-spike-sorting", "distinctipy", diff --git a/src/spikeinterface/core/generate.py b/src/spikeinterface/core/generate.py index 724d92023d..af005a8031 100644 --- a/src/spikeinterface/core/generate.py +++ b/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 diff --git a/src/spikeinterface/generation/session_displacement_generator.py b/src/spikeinterface/generation/session_displacement_generator.py index ce81c2d4f9..d6421b9886 100644 --- a/src/spikeinterface/generation/session_displacement_generator.py +++ b/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( diff --git a/src/spikeinterface/generation/tests/test_session_displacement_generator.py b/src/spikeinterface/generation/tests/test_session_displacement_generator.py index 05d6db4c8e..b1d8cffbfa 100644 --- a/src/spikeinterface/generation/tests/test_session_displacement_generator.py +++ b/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"]: diff --git a/src/spikeinterface/widgets/unit_spatial.py b/src/spikeinterface/widgets/unit_spatial.py new file mode 100644 index 0000000000..49521d5be2 --- /dev/null +++ b/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) diff --git a/src/spikeinterface/widgets/widget_list.py b/src/spikeinterface/widgets/widget_list.py index 988019fbac..a0ef344d12 100644 --- a/src/spikeinterface/widgets/widget_list.py +++ b/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