diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py index f37b24046d..2de9dc22d2 100644 --- a/src/spikeinterface/core/sorting_tools.py +++ b/src/spikeinterface/core/sorting_tools.py @@ -447,7 +447,7 @@ def random_spikes_selection( if maximum_rate is None: raise ValueError(f"maximum_rate must be defined") - t_duration = np.sum(get_segment_durations(sorting)) + t_duration = np.sum(get_segment_durations(sorting)[0]) rng_size = min(int(t_duration * maximum_rate), max_spikes_per_unit, all_unit_indices.size) selected_unit_indices = rng.choice(all_unit_indices, size=rng_size, replace=False, shuffle=False) diff --git a/src/spikeinterface/widgets/amplitudes.py b/src/spikeinterface/widgets/amplitudes.py index 374135aeff..2846be56d4 100644 --- a/src/spikeinterface/widgets/amplitudes.py +++ b/src/spikeinterface/widgets/amplitudes.py @@ -121,15 +121,15 @@ def __init__( if plot_histograms and bins is None: bins = 100 - # Calculate durations for all segments for x-axis limits - durations = get_segment_durations(sorting, segment_indices) + _, segment_start_stop_times = get_segment_durations(sorting, segment_indices) # Build the plot data with the full dict of dicts structure plot_data = dict( unit_colors=unit_colors, plot_histograms=plot_histograms, bins=bins, - durations=durations, + segment_start_stop_times=segment_start_stop_times, + sampling_frequency=sorting.get_sampling_frequency(), unit_ids=unit_ids, hide_unit_selector=hide_unit_selector, plot_legend=plot_legend, @@ -180,9 +180,14 @@ def plot_figpack(self, data_plot, **backend_kwargs): for u in unit_ids ] + segment_times = dp.segment_start_stop_times.values() + + start_time_sec = min(start for start, _ in segment_times) + end_time_sec = max(stop for _, stop in segment_times) + self.view = vv_views.SpikeAmplitudes( - start_time_sec=0, - end_time_sec=np.sum(dp.durations), + start_time_sec=start_time_sec, + end_time_sec=end_time_sec, plots=sa_items, # hide_unit_selector=dp.hide_unit_selector, ) diff --git a/src/spikeinterface/widgets/motion.py b/src/spikeinterface/widgets/motion.py index 030f11b276..58b0da7111 100644 --- a/src/spikeinterface/widgets/motion.py +++ b/src/spikeinterface/widgets/motion.py @@ -4,6 +4,7 @@ from spikeinterface.core import BaseRecording, SortingAnalyzer from .rasters import BaseRasterWidget +from .utils import estimate_segment_durations_from_spike_vector, get_segment_start_stop_times from spikeinterface.core.motion import Motion @@ -155,7 +156,6 @@ def __init__( color: str = "Gray", clim: tuple[float, float] | None = None, alpha: float = 1, - segment_index: int | list[int] | None = None, backend: str | None = None, **backend_kwargs, ): @@ -229,7 +229,10 @@ def __init__( segment_locations = filtered_locations[segment_mask] # Convert peak times to seconds - spike_times = segment_peaks["sample_index"] / sampling_frequency + if recording is None: + spike_times = segment_peaks["sample_index"] / sampling_frequency + else: + spike_times = recording.sample_index_to_time(segment_peaks["sample_index"], segment_index=seg_idx) # Store in dict of dicts format (using 0 as the "unit" id) spike_train_data[seg_idx] = {0: spike_times} @@ -257,19 +260,14 @@ def __init__( else: color_kwargs = dict(color=color, c=None, alpha=alpha) - # Calculate segment durations for x-axis limits + # Calculate segment start/stop times for x-axis limits if recording is not None: - durations = [recording.get_duration(seg_idx) for seg_idx in segment_indices] + segment_start_stop_times = get_segment_start_stop_times(recording, segment_indices) else: # Find boundaries between segments using searchsorted - segment_boundaries = [ - np.searchsorted(filtered_peaks["segment_index"], [seg_idx, seg_idx + 1]) for seg_idx in segment_indices - ] - - # Calculate durations from max sample in each segment - durations = [ - (filtered_peaks["sample_index"][end - 1] + 1) / sampling_frequency for (_, end) in segment_boundaries - ] + _, segment_start_stop_times = estimate_segment_durations_from_spike_vector( + filtered_peaks, segment_indices, sampling_frequency + ) plot_data = dict( spike_train_data=spike_train_data, @@ -280,7 +278,8 @@ def __init__( scatter_decimate=scatter_decimate, title="Peak depth", y_label="Depth [um]", - durations=durations, + segment_start_stop_times=segment_start_stop_times, + sampling_frequency=sampling_frequency, ) BaseRasterWidget.__init__(self, **plot_data, backend=backend, **backend_kwargs) diff --git a/src/spikeinterface/widgets/rasters.py b/src/spikeinterface/widgets/rasters.py index 91000be7c1..2cc054f152 100644 --- a/src/spikeinterface/widgets/rasters.py +++ b/src/spikeinterface/widgets/rasters.py @@ -3,7 +3,11 @@ from spikeinterface.core import SortingAnalyzer, BaseSorting from .base import BaseWidget, to_attr, default_backend_kwargs -from .utils import get_some_colors, validate_segment_indices, get_segment_durations +from .utils import ( + get_some_colors, + validate_segment_indices, + get_segment_durations, +) class BaseRasterWidget(BaseWidget): @@ -25,8 +29,10 @@ class BaseRasterWidget(BaseWidget): segment_indices : list | None, default: None For multi-segment data, specifies which segment(s) to plot. If None, uses all available segments. For single-segment data, this parameter is ignored. - durations : list | None, default: None + segment_start_stop_times : dict | None, default: None List of durations per segment of spike_train_data in seconds. + sampling_frequency : float | None, default: None + Sampling frequency in Hz, required to space multiple segments. plot_histograms : bool, default: False Plot histogram of y-axis data in another subplot bins : int | None, default: None @@ -67,7 +73,7 @@ def __init__( sort_by_depth: bool = False, unit_ids: list | None = None, segment_indices: list | None = None, - durations: list | None = None, + segment_start_stop_times: dict[int, tuple[float, float]] | None = None, plot_histograms: bool = False, bins: int | None = None, scatter_decimate: int = 1, @@ -82,9 +88,9 @@ def __init__( hide_unit_selector: bool = True, segment_boundary_kwargs: dict | None = None, backend: str | None = None, + sampling_frequency: float | None = None, **backend_kwargs, ): - # Set default segment boundary kwargs if not provided if segment_boundary_kwargs is None: segment_boundary_kwargs = {"color": "gray", "linestyle": "--", "alpha": 0.7} @@ -113,16 +119,27 @@ def __init__( all_units.update(spike_train_data[seg_idx].keys()) unit_ids = list(all_units) - # Calculate cumulative durations for segment boundaries - segment_boundaries = np.cumsum(durations) - cumulative_durations = np.concatenate([[0], segment_boundaries]) + segment_time_offsets = None + # Drop segment times we don't want to use + if segment_start_stop_times is not None: + segment_start_stop_times = {seg: segment_start_stop_times[seg] for seg in segments_to_use} + + if len(segment_start_stop_times) > 1: + if sampling_frequency is None: + raise ValueError("sampling_frequency is required when plotting multiple segments") + # Align the segment times such that consecutive segments appear directly next to each other + # on the plot, rather than at their true times (if true times are used, large inter-segment + # gaps can appear on the plot) + starts, stops = np.array(list(segment_start_stop_times.values())).T + ts = 1 / sampling_frequency + segment_time_offsets = np.r_[0, np.cumsum(starts[1:] - (stops[:-1] + ts))] # Concatenate data across segments with proper time offsets concatenated_spike_trains = {unit_id: np.array([]) for unit_id in unit_ids} concatenated_y_axis = {unit_id: np.array([]) for unit_id in unit_ids} - for offset, spike_train_segment, y_axis_segment in zip( - cumulative_durations, + for segment_index, spike_train_segment, y_axis_segment in zip( + range(len(segments_to_use)), [spike_train_data[idx] for idx in segments_to_use], [y_axis_data[idx] for idx in segments_to_use], ): @@ -134,13 +151,11 @@ def __init__( # Get y-axis values for this unit y_values = y_axis_segment[unit_id] - # Apply offset to spike times - adjusted_times = spike_times + offset + if segment_time_offsets is not None: + spike_times = spike_times - segment_time_offsets[segment_index] # Add to concatenated data - concatenated_spike_trains[unit_id] = np.concatenate( - [concatenated_spike_trains[unit_id], adjusted_times] - ) + concatenated_spike_trains[unit_id] = np.concatenate([concatenated_spike_trains[unit_id], spike_times]) concatenated_y_axis[unit_id] = np.concatenate([concatenated_y_axis[unit_id], y_values]) plot_data = dict( @@ -157,12 +172,12 @@ def __init__( unit_colors=unit_colors, y_label=y_label, title=title, - durations=durations, + segment_start_stop_times=segment_start_stop_times, + segment_time_offsets=segment_time_offsets, plot_legend=plot_legend, bins=bins, y_ticks=y_ticks, hide_unit_selector=hide_unit_selector, - segment_boundaries=segment_boundaries, segment_boundary_kwargs=segment_boundary_kwargs, ) @@ -233,11 +248,6 @@ def plot_matplotlib(self, data_plot, **backend_kwargs): count, bins = np.histogram(unit_y_data, bins=bins) ax_hist.plot(count, bins[:-1], color=unit_colors[unit_id], alpha=0.8) - # Add segment boundary lines if provided - if getattr(dp, "segment_boundaries", None) is not None: - for boundary in dp.segment_boundaries: - scatter_ax.axvline(boundary, **dp.segment_boundary_kwargs) - if dp.plot_histograms: ax_hist = self.axes.flatten()[1] ax_hist.set_ylim(scatter_ax.get_ylim()) @@ -253,9 +263,23 @@ def plot_matplotlib(self, data_plot, **backend_kwargs): if dp.y_lim is not None: scatter_ax.set_ylim(*dp.y_lim) x_lim = dp.x_lim - if x_lim is None: - x_lim = [0, np.sum(dp.durations)] - scatter_ax.set_xlim(x_lim) + + # Add segment boundary lines if provided and handle x limits + if dp.segment_start_stop_times is not None: + + segment_times = np.array(list(dp.segment_start_stop_times.values())) + if len(segment_times) > 1: + + segment_times = segment_times - dp.segment_time_offsets[:, None] + self._set_segment_time_ticks(scatter_ax, segment_times, dp.segment_time_offsets) + + for boundary in segment_times[1:, 0]: + scatter_ax.axvline(boundary, **dp.segment_boundary_kwargs) + + if x_lim is None: + x_lim = (segment_times[0, 0], segment_times[-1, 1]) + if x_lim is not None: + scatter_ax.set_xlim(x_lim) if dp.sort_by_depth and dp.depth_dict is not None: scatter_ax.set_yticks(ticks=list(range(len(dp.depth_dict))), labels=list(dp.depth_dict.keys())) @@ -381,6 +405,24 @@ def _update_plot(self, change=None): self.figure.canvas.draw() self.figure.canvas.flush_events() + @staticmethod + def _set_segment_time_ticks(ax, segment_times, segment_offsets): + """Label concatenated segments with original times, using default tick placement. + + ``segment_times`` contains displayed bounds; ``segment_offsets`` restores + each segment's original timestamps after gaps have been removed. + """ + from matplotlib.ticker import ScalarFormatter + + time_formatter = ScalarFormatter() + time_formatter.set_axis(ax.xaxis) + + def format_original_time(x, pos): + segment = np.searchsorted(segment_times[1:, 0], x, side="right") + return time_formatter.format_data_short(x + segment_offsets[segment]) + + ax.xaxis.set_major_formatter(format_original_time) + class RasterWidget(BaseRasterWidget): """ @@ -443,20 +485,16 @@ def __init__( # Create a lookup dictionary for unit indices unit_indices_map = {unit_id: i for i, unit_id in enumerate(unit_ids)} - # Estimate segment duration from max spike time in each segment - durations = get_segment_durations(sorting, segment_indices) - # Extract spike data for all segments and units at once spike_train_data = {seg_idx: {} for seg_idx in segment_indices} y_axis_data = {seg_idx: {} for seg_idx in segment_indices} + _, segment_start_stop_times = get_segment_durations(sorting, segment_indices) + for seg_idx in segment_indices: for unit_id in unit_ids: # Get spikes for this segment and unit - spike_times = ( - sorting.get_unit_spike_train(unit_id=unit_id, segment_index=seg_idx) / sorting.sampling_frequency - ) - + spike_times = sorting.get_unit_spike_train_in_seconds(unit_id=unit_id, segment_index=seg_idx) # Store data spike_train_data[seg_idx][unit_id] = spike_times y_axis_data[seg_idx][unit_id] = unit_indices_map[unit_id] * np.ones(len(spike_times)) @@ -486,7 +524,8 @@ def __init__( unit_colors=unit_colors, plot_histograms=None, y_ticks=y_ticks, - durations=durations, + segment_start_stop_times=segment_start_stop_times, + sampling_frequency=sorting.get_sampling_frequency(), ) BaseRasterWidget.__init__(self, **plot_data, backend=backend, **backend_kwargs) diff --git a/src/spikeinterface/widgets/tests/test_widgets.py b/src/spikeinterface/widgets/tests/test_widgets.py index c811c386e2..8ece826de9 100644 --- a/src/spikeinterface/widgets/tests/test_widgets.py +++ b/src/spikeinterface/widgets/tests/test_widgets.py @@ -632,6 +632,11 @@ def test_plot_rasters(self): if backend not in self.skip_backends: sw.plot_rasters(self.sorting) + shifted_sorting = self.sorting.clone() + shifted_sorting.shift_times(5.0) + widget = sw.plot_rasters(shifted_sorting, backend="matplotlib") + assert widget.ax.get_xlim()[0] >= 5.0 + def test_plot_unit_probe_map(self): possible_backends = list(sw.UnitProbeMapWidget.get_possible_backends()) for backend in possible_backends: @@ -691,6 +696,16 @@ def test_drift_raster_map(self): sw.plot_drift_raster_map( peaks=peaks, peak_locations=peak_locations, recording=recording, color_amplitude=True ) + if backend == "matplotlib": + shifted_recording = recording.clone() + shifted_recording.shift_times(5.0) + widget = sw.plot_drift_raster_map( + peaks=peaks, peak_locations=peak_locations, recording=shifted_recording + ) + np.testing.assert_allclose( + widget.ax.get_xlim(), + [shifted_recording.get_start_time(), shifted_recording.get_end_time()], + ) # without recording sw.plot_drift_raster_map( peaks=peaks, diff --git a/src/spikeinterface/widgets/tests/test_widgets_utils.py b/src/spikeinterface/widgets/tests/test_widgets_utils.py index 3adf31c189..92b5fd2e64 100644 --- a/src/spikeinterface/widgets/tests/test_widgets_utils.py +++ b/src/spikeinterface/widgets/tests/test_widgets_utils.py @@ -1,7 +1,12 @@ import pytest -from spikeinterface import generate_sorting -from spikeinterface.widgets.utils import get_some_colors, validate_segment_indices, get_segment_durations +from spikeinterface import NumpyRecording, generate_sorting +from spikeinterface.widgets.utils import ( + get_some_colors, + validate_segment_indices, + get_segment_durations, +) +import numpy as np def test_get_some_colors(): @@ -51,22 +56,31 @@ def test_validate_segment_indices(): validate_segment_indices([5], sorting_multiple) -def test_get_segment_durations(): +def test_get_segment_durations_estimate(): + """Test a sortings get_segment_durations when estimated from the spike time""" from spikeinterface import generate_sorting # Test with a normal multi-segment sorting durations = [5.0, 10.0, 15.0] + t_starts = [10, 20, 30] # Create sorting with high fr to ensure spikes near the end segments sorting = generate_sorting( durations=durations, firing_rates=15.0, + t_starts=t_starts, ) segment_indices = list(range(sorting.get_num_segments())) # Calculate durations - calculated_durations = get_segment_durations(sorting, segment_indices) + calculated_durations, segment_start_stop_times = get_segment_durations(sorting, segment_indices) + + # Check results + expected_start_stop_times = np.array([(10, 15), (20, 30), (30, 45)], dtype=float) + assert list(segment_start_stop_times) == segment_indices + start_stop_times = np.array(list(segment_start_stop_times.values())) + assert np.allclose(expected_start_stop_times, start_stop_times, rtol=0, atol=0.1) # Check results assert len(calculated_durations) == len(durations) @@ -82,16 +96,47 @@ def test_get_segment_durations(): sorting_single = generate_sorting( durations=[7.0], firing_rates=15.0, + t_starts=[4], ) - single_duration = get_segment_durations(sorting_single, [0])[0] + single_duration, segment_start_stop_times = get_segment_durations(sorting_single, [0]) + + expected_start_stop_times = [(4, 11)] + assert list(segment_start_stop_times) == [0] + start_stop_times = np.array(list(segment_start_stop_times.values())) + assert np.allclose(expected_start_stop_times, start_stop_times, rtol=0, atol=0.1) # Test that the calculated duration is reasonable assert single_duration <= 7.0 assert 7.0 - single_duration < 0.7 # Within 10% +def test_get_segment_durations_with_recording(): + """Test a sortings get_segment_durations when a recording is + registered, meaning times should be exact and not estimated from spike times""" + durations = [5.0, 10.0, 15.0] + t_starts = [10, 20, 30] + sampling_frequency = 1024 + + sorting = generate_sorting(durations=durations, sampling_frequency=sampling_frequency, seed=42) + + recording = NumpyRecording( + [np.zeros((int(duration * sampling_frequency), 1), dtype="float32") for duration in durations], + sampling_frequency=sampling_frequency, + t_starts=t_starts, + ) + sorting.register_recording(recording) + + actual_durations, bounds = get_segment_durations(sorting) + np.testing.assert_array_equal(actual_durations, durations) + assert bounds == { + seg: (start, start + duration - 1 / sampling_frequency) + for seg, (start, duration) in enumerate(zip(t_starts, durations)) + } + + if __name__ == "__main__": test_get_some_colors() test_validate_segment_indices() - test_get_segment_durations() + test_get_segment_durations_estimate() + test_get_segment_durations_with_recording() diff --git a/src/spikeinterface/widgets/utils.py b/src/spikeinterface/widgets/utils.py index 40da8fd0d0..9d7be46cd4 100644 --- a/src/spikeinterface/widgets/utils.py +++ b/src/spikeinterface/widgets/utils.py @@ -400,9 +400,22 @@ def validate_segment_indices(segment_indices: list[int] | None, sorting: BaseSor return segment_indices -def get_segment_durations(sorting: BaseSorting, segment_indices: list[int] = None) -> list[float]: +def get_segment_start_stop_times(recording_or_sorting, segment_indices): + """Convenience function to return the start and stop times of + a recording or sorting in the `segment_start_stop_times` format. """ - Calculate the duration of each segment in a sorting object. + return { + seg: (recording_or_sorting.get_start_time(seg), recording_or_sorting.get_end_time(seg)) + for seg in segment_indices + } + + +def get_segment_durations(sorting: BaseSorting, segment_indices: list[int] = None): + """ + Calculate the duration and start, stop times of each segment in a sorting object. + + Use first/last sample timestamps when a recording is available; otherwise + estimate from spikes. Parameters ---------- @@ -413,18 +426,65 @@ def get_segment_durations(sorting: BaseSorting, segment_indices: list[int] = Non Returns ------- - list[float] - List of segment durations in seconds + durations : np.ndarray + Segment durations in seconds, ordered like ``segment_indices``. + segment_start_stop_times : dict[int, tuple[float, float]] + Start and stop times of each segment, where the dict key is segment index. """ if segment_indices is None: segment_indices = range(sorting.get_num_segments()) + if sorting.has_recording(): + segment_start_stop_times = get_segment_start_stop_times(sorting, segment_indices) + ts = 1 / sorting.get_sampling_frequency() + durations = np.array([stop - start + ts for start, stop in segment_start_stop_times.values()]) + + return durations, segment_start_stop_times + spikes = sorting.to_spike_vector() - segment_boundaries = [ - np.searchsorted(spikes["segment_index"], [seg_idx, seg_idx + 1]) for seg_idx in segment_indices - ] + segment_t_starts = [sorting.get_start_time(seg_idx) for seg_idx in segment_indices] + + return estimate_segment_durations_from_spike_vector( + spikes, segment_indices, sorting.get_sampling_frequency(), segment_t_starts + ) + + +def estimate_segment_durations_from_spike_vector( + spike_vector: np.ndarray, + segment_indices: list[int], + sampling_frequency: float, + segment_t_starts: list[float] | None = None, +): + """ + When we don't have the true segment times, we can estimate the + start and stop times for each segment from the spike times. + + Starts are supplied or default to zero. Durations include the final sample interval. + Empty segments have zero duration and placeholder bounds from their start to one sample interval later. + """ + segment_start_stop_times = {} + durations = [] + ts = 1 / sampling_frequency + for i, seg_idx in enumerate(segment_indices): + + if segment_t_starts is None: + segment_start = 0.0 + else: + segment_start = segment_t_starts[i] + + seg_end_idx = np.searchsorted(spike_vector["segment_index"], seg_idx + 1) + + if seg_end_idx == 0 or spike_vector["segment_index"][seg_end_idx - 1] != seg_idx: + segment_start_stop_times[seg_idx] = (segment_start, segment_start + ts) + durations.append(0.0) + continue + + segment_end = spike_vector["sample_index"][seg_end_idx - 1] / sampling_frequency + segment_start + + segment_start_stop_times[seg_idx] = (segment_start, segment_end) + durations.append(segment_end - segment_start + ts) - durations = [(spikes["sample_index"][end - 1] + 1) / sorting.sampling_frequency for (_, end) in segment_boundaries] + durations = np.array(durations) - return durations + return durations, segment_start_stop_times