Fix raster plot segment

[JoeZiminski]

This PR updates the raster widgets (and consequentially two other

Before

After

Test code

Test code

"""
Generate a multi-segment recording whose segments start at non-zero times,
then plot a raster and motion-correction diagnostics.
"""

import matplotlib.pyplot as plt

import spikeinterface.widgets as sw
from spikeinterface.core import (
    append_recordings,
    append_sortings,
    create_sorting_analyzer,
    generate_ground_truth_recording,
)
from spikeinterface.preprocessing import correct_motion


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, 25, 50]

    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, recordings[0]


def main():
    recording, sorting, motion_recording = 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)

    sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory")
    sorting_analyzer.compute("random_spikes")
    sorting_analyzer.compute("waveforms")
    sorting_analyzer.compute("templates")
    sorting_analyzer.compute("spike_amplitudes")
    sw.plot_amplitudes(
        sorting_analyzer,
        segment_indices=segment_indices,
        plot_histograms=True,
        plot_legend=False,
    )

    _, motion_info = correct_motion(
        motion_recording,
        preset="kilosort_like",
        output_motion_info=True,
        n_jobs=1,
    )
    sw.plot_motion_info(motion_info, recording=motion_recording, motion_lim=1.0)

    plt.show()


if __name__ == "__main__":
    main()