finalise all tests and time handling.

Author: JoeZiminski

src/spikeinterface/widgets/peaks_on_probe.py

@@ -104,7 +104,7 @@ def plot_matplotlib(self, data_plot, **backend_kwargs):
             ax = self.axes[ax_idx]
             plot_probe_map(dp.recording, ax=ax)
 
-            time_mask = self._get_peaks_time_mask(dp.recording, time_range, fs, peaks_to_plot)
+            time_mask = self._get_peaks_time_mask(dp.recording, time_range, peaks_to_plot)
 
             if dp.segment_index is not None:
                 segment_mask = peaks_to_plot["segment_index"] == dp.segment_index
@@ -142,7 +142,7 @@ def plot_matplotlib(self, data_plot, **backend_kwargs):
 
         self.figure.suptitle(f"Peaks on Probe Plot")
 
-    def _get_peaks_time_mask(self, recording, time_range, fs, peaks_to_plot):
+    def _get_peaks_time_mask(self, recording, time_range, peaks_to_plot):
         """
         Return a mask of `True` where the peak is within the given time range
         and `False` otherwise.
@@ -162,10 +162,12 @@ def _get_peaks_time_mask(self, recording, time_range, fs, peaks_to_plot):
 
             seg_mask = peaks_to_plot["segment_index"] == seg_idx
 
-            time_mask[seg_mask] = (t_start_sample <= peaks_to_plot[seg_mask]["sample_index"]) & (
-                peaks_to_plot[seg_mask]["sample_index"] <= t_stop_sample
+            time_mask[seg_mask] = (t_start_sample < peaks_to_plot[seg_mask]["sample_index"]) & (
+                peaks_to_plot[seg_mask]["sample_index"] < t_stop_sample
             )
 
+        return time_mask
+
     def _get_min_and_max_times_in_recording(self, recording):
         """
         Find the maximum and minimum time across all segments in the recording.
@@ -180,7 +182,7 @@ def _get_min_and_max_times_in_recording(self, recording):
 
             t_starts.append(segment.sample_index_to_time(0))
 
-            t_stops.append(segment.sample_index_to_time(segment.get_num_samples()))
+            t_stops.append(segment.sample_index_to_time(segment.get_num_samples() - 1))
 
         time_range = (np.min(t_starts), np.max(t_stops))
 

src/spikeinterface/widgets/tests/test_peaks_on_probe.py

@@ -14,7 +14,7 @@ def peak_info(self):
         Fixture (created only once per test run) of a small
         ground truth recording with peaks and peak locations calculated.
         """
-        recording, sorting = generate_ground_truth_recording(num_units=5, num_channels=16, durations=[20, 9], seed=0)
+        recording, _ = generate_ground_truth_recording(num_units=5, num_channels=16, durations=[20, 9], seed=0)
         peaks = detect_peaks(recording)
 
         peak_locations = localize_peaks(
@@ -25,7 +25,7 @@ def peak_info(self):
             method="center_of_mass",
         )
 
-        return (recording, sorting, peaks, peak_locations)
+        return (recording, peaks, peak_locations)
 
     def data_from_widget(self, widget, axes_idx):
         """
@@ -40,7 +40,7 @@ def test_peaks_on_probe_main(self, peak_info):
         Plot all peaks, and check every peak is plot.
         Check the labels are corect.
         """
-        recording, sorting, peaks, peak_locations = peak_info
+        recording, peaks, peak_locations = peak_info
 
         widget = plot_peaks_on_probe(recording, peaks, peak_locations, decimate=1)
 
@@ -58,7 +58,7 @@ def test_segment_selection(self, peak_info, segment_index):
         from a sepecific segment, that only peaks
         from that segment are plot.
         """
-        recording, sorting, peaks, peak_locations = peak_info
+        recording, peaks, peak_locations = peak_info
 
         widget = plot_peaks_on_probe(
             recording,
@@ -82,7 +82,7 @@ def test_multiple_inputs(self, peak_info):
         Check that these separate peaks / peak locations
         are plot on different axes.
         """
-        recording, sorting, peaks, peak_locations = peak_info
+        recording, peaks, peak_locations = peak_info
 
         half_num_peaks = int(peaks.shape[0] / 2)
 
@@ -107,16 +107,14 @@ def test_multiple_inputs(self, peak_info):
 
         assert np.array_equal(np.sort(locs_change["y"]), np.sort(ax_1_y_data))
 
-    def test_times(self, peak_info):
+    def test_times_all(self, peak_info):
         """
         Check that when the times of peaks to plot is restricted,
         only peaks within the given time range are plot. Set the
         limits just before and after the second peak, and check only
         that peak is plot.
         """
-        recording, sorting, peaks, peak_locations = peak_info
-
-        peak_times_ms = recording.sample_index_to_time(peaks["sample_index"], segment_index=0) * 1000
+        recording, peaks, peak_locations = peak_info
 
         peak_idx = 1
         peak_cutoff_low = peaks["sample_index"][peak_idx] - 1
@@ -137,13 +135,104 @@ def test_times(self, peak_info):
 
         assert np.array_equal([peak_locations[peak_idx]["y"]], ax_y_data)
 
+    def test_times_per_segment(self, peak_info):
+        """
+        Test that the time bounds for multi-segment recordings
+        with different times are handled properly. The time bounds
+        given must respect the times for each segment. Here, we build
+        two segments with times 0-100s and 100-200s. We set the
+        time limits for peaks to plot as 50-150 i.e. all peaks
+        from the second half of the first segment, and the first half
+        of the second segment, should be plotted.
+
+        Recompute peaks here for completeness even though this does
+        duplicate the fixture.
+        """
+        recording, _, _ = peak_info
+
+        first_seg_times = np.linspace(0, 100, recording.get_num_samples(0))
+        second_seg_times = np.linspace(100, 200, recording.get_num_samples(1))
+
+        recording.set_times(first_seg_times, segment_index=0)
+        recording.set_times(second_seg_times, segment_index=1)
+
+        # After setting the peak times above, re-detect peaks and plot
+        # with a time range 50-150 s
+        peaks = detect_peaks(recording)
+
+        peak_locations = localize_peaks(
+            recording,
+            peaks,
+            ms_before=0.3,
+            ms_after=0.6,
+            method="center_of_mass",
+        )
+
+        widget = plot_peaks_on_probe(
+            recording,
+            peaks,
+            peak_locations,
+            decimate=1,
+            time_range=(
+                50,
+                150,
+            ),
+        )
+
+        # Find the peaks that are expected to be plot given the time
+        # restriction (second half of first segment, first half of
+        # second segment) and check that indeed the expected locations
+        # are displayed.
+        seg_one_num_samples = recording.get_num_samples(0)
+        seg_two_num_samples = recording.get_num_samples(1)
+
+        okay_peaks_one = np.logical_and(
+            peaks["segment_index"] == 0, peaks["sample_index"] > int(seg_one_num_samples / 2)
+        )
+        okay_peaks_two = np.logical_and(
+            peaks["segment_index"] == 1, peaks["sample_index"] < int(seg_two_num_samples / 2)
+        )
+        okay_peaks = np.logical_or(okay_peaks_one, okay_peaks_two)
+
+        ax_y_data = self.data_from_widget(widget, 0)[:, 1]
+
+        assert any(okay_peaks), "someting went wrong in test generation, no peaks within the set time bounds detected"
+
+        assert np.array_equal(np.sort(ax_y_data), np.sort(peak_locations[okay_peaks]["y"]))
+
+    def test_get_min_and_max_times_in_recording(self, peak_info):
+        """
+        Check that the function which finds the minimum and maximum times
+        across all segments in the recording returns correctly. First
+        set times of the segments such that the earliest time is 50s and
+        latest 200s. Check the function returns (50, 200).
+        """
+        recording, peaks, peak_locations = peak_info
+
+        first_seg_times = np.linspace(50, 100, recording.get_num_samples(0))
+        second_seg_times = np.linspace(100, 200, recording.get_num_samples(1))
+
+        recording.set_times(first_seg_times, segment_index=0)
+        recording.set_times(second_seg_times, segment_index=1)
+
+        widget = plot_peaks_on_probe(
+            recording,
+            peaks,
+            peak_locations,
+            decimate=1,
+        )
+
+        min_max_times = widget._get_min_and_max_times_in_recording(recording)
+
+        assert min_max_times == (50, 200)
+
     def test_ylim(self, peak_info):
         """
         Specify some y-axis limits (which is the probe height
         to show) and check that the plot is restricted to
         these limits.
         """
-        recording, sorting, peaks, peak_locations = peak_info
+        recording, peaks, peak_locations = peak_info
 
         widget = plot_peaks_on_probe(
             recording,
@@ -163,7 +252,7 @@ def test_decimate(self, peak_info):
         checks the decimate argument, to ensure peaks that are
         plot are correctly decimated.
         """
-        recording, sorting, peaks, peak_locations = peak_info
+        recording, peaks, peak_locations = peak_info
 
         decimate = 5
 
@@ -184,7 +273,7 @@ def test_errors(self, peak_info):
         Test all validation errors are raised when data in
         incorrect form is passed to the plotting function.
         """
-        recording, sorting, peaks, peak_locations = peak_info
+        recording, peaks, peak_locations = peak_info
 
         # All lists must be same length
         with pytest.raises(ValueError) as e: