diff --git a/playing.py b/playing.py new file mode 100644 index 0000000000..2894b60265 --- /dev/null +++ b/playing.py @@ -0,0 +1,92 @@ +""" +Generate a multi-segment recording whose segments start at non-zero times, +then plot a raster of the segments. +""" + +import matplotlib + +matplotlib.use("QtAgg") +import matplotlib.pyplot as plt + +from spikeinterface.core import append_recordings, append_sortings, create_sorting_analyzer, generate_ground_truth_recording +from spikeinterface.generation import generate_drifting_recording +from spikeinterface.preprocessing import compute_motion +import spikeinterface.widgets as sw + + +def make_multi_segment_recording(num_segments=3, seed=2205): + """Build a multi-segment recording and sorting with non-zero segment start times.""" + + common_kwargs = dict( + durations=[10.0], + sampling_frequency=30000.0, + num_channels=64, + num_units=60, + ) + + recordings = [] + sortings = [] + # Offsets so segments do NOT start at time zero. + start_time_offsets = [5.0, 120.0, 260.0] + + for seg in range(num_segments): + recording, sorting = generate_ground_truth_recording(seed=seed + seg, **common_kwargs) + # Give this segment a non-zero start time. + recording.shift_times(shift=start_time_offsets[seg % len(start_time_offsets)]) + recordings.append(recording) + sortings.append(sorting) + + multi_segment_recording = append_recordings(recordings) + multi_segment_sorting = append_sortings(sortings) + return multi_segment_recording, multi_segment_sorting + + +def plot_amplitudes_output(recording, sorting): + """Plot spike amplitudes for a few units across all recording segments.""" + + analyzer = create_sorting_analyzer(sorting=sorting, recording=recording, format="memory") + analyzer.compute(["random_spikes", "waveforms", "templates", "spike_amplitudes"]) + + amplitude_widget = sw.plot_amplitudes( + analyzer, + unit_ids=sorting.unit_ids[:6], + segment_indices=list(range(sorting.get_num_segments())), + plot_histograms=True, + ) + amplitude_widget.figure.canvas.manager.set_window_title("Amplitude output") + amplitude_widget.axes.flatten()[0].set_title("Spike amplitudes") + + +def main(): + recording, sorting = make_multi_segment_recording() + sorting.register_recording(recording) + + for seg in range(recording.get_num_segments()): + print(f" segment {seg}: start={recording.get_start_time(seg):.1f}s end={recording.get_end_time(seg):.1f}s") + + # Plot a raster of all segments. + segment_indices = list(range(sorting.get_num_segments())) + # sw.plot_rasters(sorting, segment_indices=segment_indices) + # plot_amplitudes_output(recording, sorting) + + + motion_recording = recording.select_segments([0]) + motion, motion_info = compute_motion( + motion_recording, + preset="rigid_fast", + estimate_motion_kwargs=dict(method="decentralized", conv_engine="numpy"), + n_jobs=-1, + progress_bar=True, + output_motion_info=True, + ) + print(motion) + motion_widget = sw.plot_motion_info(motion_info, recording=motion_recording) + motion_widget.figure.canvas.manager.set_window_title("Motion output") + motion_widget.figure.canvas.draw() + motion_widget.figure.show() + + plt.show(block=True) + + +if __name__ == "__main__": + main() 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..5a129821e8 100644 --- a/src/spikeinterface/widgets/amplitudes.py +++ b/src/spikeinterface/widgets/amplitudes.py @@ -122,14 +122,14 @@ def __init__( 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, 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..36f91a0a0c 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 compute_segment_durations_from_spike_vector 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,17 @@ 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 = { + seg_idx: (recording.get_start_time(seg_idx), recording.get_end_time(seg_idx)) + for seg_idx in 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 = compute_segment_durations_from_spike_vector( + filtered_peaks, segment_indices, sampling_frequency + ) plot_data = dict( spike_train_data=spike_train_data, @@ -280,7 +281,7 @@ def __init__( scatter_decimate=scatter_decimate, title="Peak depth", y_label="Depth [um]", - durations=durations, + segment_start_stop_times=segment_start_stop_times, ) BaseRasterWidget.__init__(self, **plot_data, backend=backend, **backend_kwargs) diff --git a/src/spikeinterface/widgets/rasters.py b/src/spikeinterface/widgets/rasters.py index 2ec8f42777..5e70777e06 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): @@ -67,7 +71,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, @@ -84,7 +88,6 @@ def __init__( backend: str | 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 +116,15 @@ 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]) + # 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} # 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 spike_train_segment, y_axis_segment in zip( [spike_train_data[idx] for idx in segments_to_use], [y_axis_data[idx] for idx in segments_to_use], ): @@ -134,13 +136,8 @@ 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 - # 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 +154,11 @@ def __init__( unit_colors=unit_colors, y_label=y_label, title=title, - durations=durations, + segment_start_stop_times=segment_start_stop_times, 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 +229,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 +244,20 @@ 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: + + # We only plot boundaries for sequential segments. + segment_boundaries = np.array(list(dp.segment_start_stop_times.values())).ravel() + + if np.all(np.diff(segment_boundaries) > 0): + for boundary in segment_boundaries: + scatter_ax.axvline(boundary, **dp.segment_boundary_kwargs) + + if x_lim is None: + x_lim = (segment_boundaries.min(), segment_boundaries.max()) + 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())) @@ -455,24 +457,39 @@ 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 the spikes for this segment. Also, if the sorting does not have a recording, + # build the segment start / stop times from the spike times. This is not done + # with get_segment_durations to avoid the conversion to spike_vector. for seg_idx in segment_indices: + + min_spiketime = np.inf + max_spiketime = -np.inf + 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)) + if not sorting.has_recording() and len(spike_times) > 0: + min_spiketime = min(min_spiketime, spike_times[0]) + max_spiketime = max(max_spiketime, spike_times[-1]) + + if sorting.has_recording(): + segment_start_stop_times[seg_idx] = ( + sorting.get_start_time(seg_idx), + sorting.get_end_time(seg_idx), + ) + else: + segment_start_stop_times[seg_idx] = (min_spiketime, max_spiketime) + # Apply time range filtering if specified if time_range is not None: assert len(time_range) == 2, "'time_range' should be a list with start and end time in seconds" @@ -498,7 +515,7 @@ def __init__( unit_colors=unit_colors, plot_histograms=None, y_ticks=y_ticks, - durations=durations, + segment_start_stop_times=segment_start_stop_times, ) 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..1a3868acea 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.widgets.utils import ( + get_some_colors, + validate_segment_indices, + get_segment_durations, +) +import numpy as np def test_get_some_colors(): @@ -56,17 +61,25 @@ def test_get_segment_durations(): # 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,9 +95,15 @@ 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 diff --git a/src/spikeinterface/widgets/utils.py b/src/spikeinterface/widgets/utils.py index 40da8fd0d0..3602f9f07c 100644 --- a/src/spikeinterface/widgets/utils.py +++ b/src/spikeinterface/widgets/utils.py @@ -400,7 +400,7 @@ 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_durations(sorting: BaseSorting, segment_indices: list[int] = None): """ Calculate the duration of each segment in a sorting object. @@ -413,18 +413,49 @@ 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 keyed by segment index. """ if segment_indices is None: segment_indices = range(sorting.get_num_segments()) spikes = sorting.to_spike_vector() + segment_t_starts = [sorting.get_start_time(seg_idx) for seg_idx in segment_indices] + + return compute_segment_durations_from_spike_vector( + spikes, segment_indices, sorting.get_sampling_frequency(), segment_t_starts + ) + + +def compute_segment_durations_from_spike_vector( + spike_vector: np.ndarray, + segment_indices: list[int], + sampling_frequency: float, + segment_t_starts: list[float] | None = None, +): + """ + If segment_t_starts is `None` then assume 0 + """ + if segment_t_starts is None: + segment_t_starts = [0] * len(segment_indices) + segment_boundaries = [ - np.searchsorted(spikes["segment_index"], [seg_idx, seg_idx + 1]) for seg_idx in segment_indices + np.searchsorted(spike_vector["segment_index"], [seg_idx, seg_idx + 1]) for seg_idx in segment_indices ] - durations = [(spikes["sample_index"][end - 1] + 1) / sorting.sampling_frequency for (_, end) in segment_boundaries] + segment_start_stop_times = {} + for seg_idx, segment_t_start, (start, end) in zip(segment_indices, segment_t_starts, segment_boundaries): + + segment_start = spike_vector["sample_index"][start] / sampling_frequency + segment_t_start + segment_end = (spike_vector["sample_index"][end - 1] + 1) / sampling_frequency + segment_t_start + + segment_start_stop_times[seg_idx] = (segment_start, segment_end) + + durations = np.array( + [segment_start_stop_times[seg_idx][1] - segment_start_stop_times[seg_idx][0] for seg_idx in segment_indices] + ) - return durations + return durations, segment_start_stop_times