ENH: vectorize annotation span conversion

Implements the get_annotation_spans() design suggested by @larsoner during review of #14240.

Author: timothyanderson096-ocdealcheck

doc/changes/dev/14240.newfeature.rst

@@ -1 +1 @@
-Add :meth:`mne.io.BaseRaw.get_annotation_span` to convert annotation spans to the time reference used by Raw data and plotting methods, by :newcontrib:`Tim Anderson`.
+Add :meth:`mne.io.BaseRaw.get_annotation_spans` to convert annotation spans to the time reference used by Raw data and plotting methods, by :newcontrib:`Tim Anderson`.

mne/io/base.py

@@ -740,20 +740,15 @@ def annotations(self):  # noqa: D401
         """:class:`~mne.Annotations` for marking segments of data."""
         return self._annotations
 
-    def get_annotation_span(self, index: int) -> tuple[float, float]:
-        """Get an annotation span relative to the first data sample.
-
-        Parameters
-        ----------
-        index : int
-            Index of the annotation in :attr:`annotations`.
+    def get_annotation_spans(self) -> tuple[np.ndarray, np.ndarray]:
+        """Get annotation spans relative to the first data sample.
 
         Returns
         -------
-        tmin : float
-            Annotation onset in seconds relative to the first data sample.
-        tmax : float
-            Annotation end in seconds relative to the first data sample.
+        tmin : ndarray, shape (n_annotations,)
+            Annotation onsets in seconds relative to the first data sample.
+        tmax : ndarray, shape (n_annotations,)
+            Annotation ends in seconds relative to the first data sample.
 
         Notes
         -----
@@ -763,9 +758,8 @@ def get_annotation_span(self, index: int) -> tuple[float, float]:
         use zero at the first available data sample. This method converts
         between those time references.
         """
-        _validate_type(index, "int-like", "index")
-        tmin = float(_sync_onset(self, self.annotations.onset[index]))
-        tmax = tmin + float(self.annotations.duration[index])
+        tmin = _sync_onset(self, self.annotations.onset)
+        tmax = tmin + self.annotations.duration
         return tmin, tmax
 
     @property

mne/io/tests/test_raw.py

@@ -691,7 +691,7 @@ def test_crop_by_annotations(meas_date, first_samp):
 
 @pytest.mark.parametrize("meas_date", [None, 0])
 @pytest.mark.parametrize("first_samp", [0, 50])
-def test_get_annotation_span(meas_date, first_samp):
+def test_get_annotation_spans(meas_date, first_samp):
     """Test converting annotation spans to the Raw time reference."""
     sfreq = 10.0
     data = np.arange(120.0)[np.newaxis]
@@ -702,21 +702,24 @@ def test_get_annotation_span(meas_date, first_samp):
         verbose="error",
     )
     raw.set_meas_date(meas_date)
-    raw.set_annotations(mne.Annotations(3.0, 1.0, "test"))
+    raw.set_annotations(mne.Annotations([3.0, 3.0], [1.0, 0.5], ["test", "test 2"]))
 
-    tmin, tmax = raw.get_annotation_span(0)
-    assert_allclose([tmin, tmax], [3.0, 4.0])
-    got, times = raw.get_data(tmin=tmin, tmax=tmax, return_times=True)
+    tmin, tmax = raw.get_annotation_spans()
+    assert_allclose(tmin, [3.0, 3.0])
+    assert_allclose(tmax, [3.5, 4.0])
+    got, times = raw.get_data(tmin=tmin[1], tmax=tmax[1], return_times=True)
     assert_array_equal(got, data[:, 30:40])
     assert_allclose(times, np.arange(30, 40) / sfreq)
 
     cropped = raw.copy().crop(2.0, 5.0)
-    tmin, tmax = cropped.get_annotation_span(0)
-    assert_allclose([tmin, tmax], [1.0, 2.0])
-    assert_array_equal(cropped.get_data(tmin=tmin, tmax=tmax), data[:, 30:40])
-
-    with pytest.raises(TypeError, match="index must be .* int"):
-        raw.get_annotation_span(0.0)
+    tmin, tmax = cropped.get_annotation_spans()
+    assert_allclose(tmin, [1.0, 1.0])
+    assert_allclose(tmax, [1.5, 2.0])
+    assert_array_equal(cropped.get_data(tmin=tmin[1], tmax=tmax[1]), data[:, 30:40])
+
+    raw.set_annotations(None)
+    tmin, tmax = raw.get_annotation_spans()
+    assert tmin.shape == tmax.shape == (0,)
 
 
 @pytest.mark.parametrize(

mne/viz/tests/test_raw.py

@@ -795,8 +795,8 @@ def test_plot_annotation_span(browser_backend):
     )
     raw.set_annotations(Annotations(3.0, 1.0, "test"))
 
-    tmin, tmax = raw.get_annotation_span(0)
-    fig = raw.plot(start=tmin, duration=tmax - tmin, show=False)
+    tmin, tmax = raw.get_annotation_spans()
+    fig = raw.plot(start=tmin[0], duration=tmax[0] - tmin[0], show=False)
     assert fig._get_start_stop() == (30, 40)
     browser_backend._close_all()