Repository: mne-tools/mne-python

Issue:
ENH: Annotations.sort
Sort annotations based on onset. Should return `self` for consistency.


Discussion:
If we have a sorting method, we should used in places like this (there are more):
https://github.com/mne-tools/mne-python/blob/e96d706b65b6c6a21821d3dcbc06e10903e9c05b/mne/viz/utils.py#L1913-L1914

The other option it would be to actually keep annotations always sorted by onset, and use `numpy.searchsorted` in `append` and `_combine_annotations`
I'd be +1 for always sorting, but I'm also ok with `annotations.sort()`.
+1 for always sorting

>

If we do this, then we should make sure there is extend in addition to append. I automate annotation of infant data, marking segments of fast movement as bad. There can be thousands of tiny segments like this per file, and so calling sort over and over again will incur a worsening penalty as the number grows. (Though using searchsorted with assume_sorted=True might be reasonable.)

We could also add a `sort=True` kwarg but this seems less clean than just always sorting and adding an extend method, if we don't have it.
+1 for always sorting and adding an extend method. As tie breakers I suggest that we use description and then duration.
It looks like `append` already has `extend`-like behavior, so we can just update the docstring.

Resolving patch:
diff --git a/mne/annotations.py b/mne/annotations.py
--- a/mne/annotations.py
+++ b/mne/annotations.py
@@ -24,6 +24,35 @@
 from .io.tag import read_tag
 
 
+def _check_o_d_s(onset, duration, description):
+    onset = np.atleast_1d(np.array(onset, dtype=float))
+    if onset.ndim != 1:
+        raise ValueError('Onset must be a one dimensional array, got %s '
+                         '(shape %s).'
+                         % (onset.ndim, onset.shape))
+    duration = np.array(duration, dtype=float)
+    if duration.ndim == 0 or duration.shape == (1,):
+        duration = np.repeat(duration, len(onset))
+    if duration.ndim != 1:
+        raise ValueError('Duration must be a one dimensional array, '
+                         'got %d.' % (duration.ndim,))
+
+    description = np.array(description, dtype=str)
+    if description.ndim == 0 or description.shape == (1,):
+        description = np.repeat(description, len(onset))
+    if description.ndim != 1:
+        raise ValueError('Description must be a one dimensional array, '
+                         'got %d.' % (description.ndim,))
+    if any([';' in desc for desc in description]):
+        raise ValueError('Semicolons in descriptions not supported.')
+
+    if not (len(onset) == len(duration) == len(description)):
+        raise ValueError('Onset, duration and description must be '
+                         'equal in sizes, got %s, %s, and %s.'
+                         % (len(onset), len(duration), len(description)))
+    return onset, duration, description
+
+
 class Annotations(object):
     """Annotation object for annotating segments of raw data.
 
@@ -154,26 +183,9 @@ def __init__(self, onset, duration, description,
         if orig_time is not None:
             orig_time = _handle_meas_date(orig_time)
         self.orig_time = orig_time
-
-        onset = np.array(onset, dtype=float)
-        if onset.ndim != 1:
-            raise ValueError('Onset must be a one dimensional array, got %s '
-                             '(shape %s).'
-                             % (onset.ndim, onset.shape))
-        duration = np.array(duration, dtype=float)
-        if isinstance(description, str):
-            description = np.repeat(description, len(onset))
-        if duration.ndim != 1:
-            raise ValueError('Duration must be a one dimensional array.')
-        if not (len(onset) == len(duration) == len(description)):
-            raise ValueError('Onset, duration and description must be '
-                             'equal in sizes.')
-        if any([';' in desc for desc in description]):
-            raise ValueError('Semicolons in descriptions not supported.')
-
-        self.onset = onset
-        self.duration = duration
-        self.description = np.array(description, dtype=str)
+        self.onset, self.duration, self.description = _check_o_d_s(
+            onset, duration, description)
+        self._sort()  # ensure we're sorted
 
     def __repr__(self):
         """Show the representation."""
@@ -236,12 +248,12 @@ def append(self, onset, duration, description):
 
         Parameters
         ----------
-        onset : float
+        onset : float | array-like
             Annotation time onset from the beginning of the recording in
             seconds.
-        duration : float
+        duration : float | array-like
             Duration of the annotation in seconds.
-        description : str
+        description : str | array-like
             Description for the annotation. To reject epochs, use description
             starting with keyword 'bad'
 
@@ -249,10 +261,19 @@ def append(self, onset, duration, description):
         -------
         self : mne.Annotations
             The modified Annotations object.
-        """
+
+        Notes
+        -----
+        The array-like support for arguments allows this to be used similarly
+        to not only ``list.append``, but also
+        `list.extend <https://docs.python.org/3/library/stdtypes.html#mutable-sequence-types>`__.
+        """  # noqa: E501
+        onset, duration, description = _check_o_d_s(
+            onset, duration, description)
         self.onset = np.append(self.onset, onset)
         self.duration = np.append(self.duration, duration)
         self.description = np.append(self.description, description)
+        self._sort()
         return self
 
     def copy(self):
@@ -264,8 +285,9 @@ def delete(self, idx):
 
         Parameters
         ----------
-        idx : int | list of int
-            Index of the annotation to remove.
+        idx : int | array-like of int
+            Index of the annotation to remove. Can be array-like to
+            remove multiple indices.
         """
         self.onset = np.delete(self.onset, idx)
         self.duration = np.delete(self.duration, idx)
@@ -295,6 +317,16 @@ def save(self, fname):
             with start_file(fname) as fid:
                 _write_annotations(fid, self)
 
+    def _sort(self):
+        """Sort in place."""
+        # instead of argsort here we use sorted so that it gives us
+        # the onset-then-duration hierarchy
+        vals = sorted(zip(self.onset, self.duration, range(len(self))))
+        order = list(list(zip(*vals))[-1]) if len(vals) else []
+        self.onset = self.onset[order]
+        self.duration = self.duration[order]
+        self.description = self.description[order]
+
     def crop(self, tmin=None, tmax=None, emit_warning=False):
         """Remove all annotation that are outside of [tmin, tmax].
 
@@ -450,12 +482,12 @@ def _annotations_starts_stops(raw, kinds, name='unknown', invert=False):
         idxs = [idx for idx, desc in enumerate(raw.annotations.description)
                 if any(desc.upper().startswith(kind.upper())
                        for kind in kinds)]
+        # onsets are already sorted
         onsets = raw.annotations.onset[idxs]
         onsets = _sync_onset(raw, onsets)
         ends = onsets + raw.annotations.duration[idxs]
-        order = np.argsort(onsets)
-        onsets = raw.time_as_index(onsets[order], use_rounding=True)
-        ends = raw.time_as_index(ends[order], use_rounding=True)
+        onsets = raw.time_as_index(onsets, use_rounding=True)
+        ends = raw.time_as_index(ends, use_rounding=True)
     assert (onsets <= ends).all()  # all durations >= 0
     if invert:
         # We need to eliminate overlaps here, otherwise wacky things happen,
diff --git a/mne/viz/raw.py b/mne/viz/raw.py
--- a/mne/viz/raw.py
+++ b/mne/viz/raw.py
@@ -1079,7 +1079,7 @@ def _plot_raw_traces(params, color, bad_color, event_lines=None,
                 continue
             start = max(segment[0], times[0] + params['first_time'])
             end = min(times[-1] + params['first_time'], segment[1])
-            dscr = params['annot_description'][idx]
+            dscr = params['raw'].annotations.description[idx]
             segment_color = params['segment_colors'][dscr]
             params['ax'].fill_betweenx(ylim, start, end, color=segment_color,
                                        alpha=0.3)
diff --git a/mne/viz/utils.py b/mne/viz/utils.py
--- a/mne/viz/utils.py
+++ b/mne/viz/utils.py
@@ -1908,23 +1908,17 @@ def _plot_annotations(raw, params):
     while len(params['ax_hscroll'].collections) > 0:
         params['ax_hscroll'].collections.pop()
     segments = list()
-    # sort the segments by start time
-    ann_order = raw.annotations.onset.argsort(axis=0)
-    descriptions = raw.annotations.description[ann_order]
-
     _setup_annotation_colors(params)
-    for idx, onset in enumerate(raw.annotations.onset[ann_order]):
-        annot_start = _sync_onset(raw, onset) + params['first_time']
-        annot_end = annot_start + raw.annotations.duration[ann_order][idx]
+    for idx, annot in enumerate(raw.annotations):
+        annot_start = _sync_onset(raw, annot['onset']) + params['first_time']
+        annot_end = annot_start + annot['duration']
         segments.append([annot_start, annot_end])
-        dscr = descriptions[idx]
         params['ax_hscroll'].fill_betweenx(
             (0., 1.), annot_start, annot_end, alpha=0.3,
-            color=params['segment_colors'][dscr])
+            color=params['segment_colors'][annot['description']])
     # Do not adjust half a sample backward (even though this would make it
     # clearer what is included) because this breaks click-drag functionality
     params['segments'] = np.array(segments)
-    params['annot_description'] = descriptions
 
 
 def _get_color_list(annotations=False):


Test patch:
diff --git a/mne/io/fiff/tests/test_raw_fiff.py b/mne/io/fiff/tests/test_raw_fiff.py
--- a/mne/io/fiff/tests/test_raw_fiff.py
+++ b/mne/io/fiff/tests/test_raw_fiff.py
@@ -25,6 +25,7 @@
                  pick_info)
 from mne.utils import (_TempDir, requires_pandas, object_diff,
                        requires_mne, run_subprocess, run_tests_if_main)
+from mne.utils.testing import assert_and_remove_boundary_annot
 from mne.annotations import Annotations
 
 testing_path = testing.data_path(download=False)
@@ -1235,6 +1236,7 @@ def test_annotation_crop():
     r2 = raw.copy().crop(12.5, 17.5)
     r3 = raw.copy().crop(10., 12.)
     raw = concatenate_raws([r1, r2, r3])  # segments reordered
+    assert_and_remove_boundary_annot(raw, 2)
     onsets = raw.annotations.onset
     durations = raw.annotations.duration
     # 2*5s clips combined with annotations at 2.5s + 2s clip, annotation at 1s
diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py
--- a/mne/io/tests/test_raw.py
+++ b/mne/io/tests/test_raw.py
@@ -302,9 +302,9 @@ def test_5839():
     #    latency  .         0    0    1    1    2    2    3
     #             .              5    0    5    0    5    0
     #
-    EXPECTED_ONSET = [1.5, 2.5, 2., 2.]
-    EXPECTED_DURATION = [0.2, 0.2, 0., 0.]
-    EXPECTED_DESCRIPTION = ['dummy', 'dummy', 'BAD boundary', 'EDGE boundary']
+    EXPECTED_ONSET = [1.5, 2., 2., 2.5]
+    EXPECTED_DURATION = [0.2, 0., 0., 0.2]
+    EXPECTED_DESCRIPTION = ['dummy', 'BAD boundary', 'EDGE boundary', 'dummy']
 
     def raw_factory(meas_date):
         raw = RawArray(data=np.empty((10, 10)),
diff --git a/mne/tests/test_annotations.py b/mne/tests/test_annotations.py
--- a/mne/tests/test_annotations.py
+++ b/mne/tests/test_annotations.py
@@ -20,6 +20,7 @@
 from mne import Epochs, Annotations
 from mne.utils import (run_tests_if_main, _TempDir, requires_version,
                        catch_logging)
+from mne.utils.testing import assert_and_remove_boundary_annot
 from mne.io import read_raw_fif, RawArray, concatenate_raws
 from mne.io.tests.test_raw import _raw_annot
 from mne.annotations import _sync_onset, _handle_meas_date
@@ -62,9 +63,7 @@ def test_basics():
     assert_array_equal(raw2.annotations.onset, onset + offset)
     assert id(raw2.annotations) != id(annot)
     concatenate_raws([raw, raw2])
-    raw.annotations.delete(-1)  # remove boundary annotations
-    raw.annotations.delete(-1)
-
+    assert_and_remove_boundary_annot(raw)
     assert_allclose(onset + offset + delta, raw.annotations.onset, rtol=1e-5)
     assert_array_equal(annot.duration, raw.annotations.duration)
     assert_array_equal(raw.annotations.description, np.repeat('test', 10))
@@ -85,23 +84,19 @@ def test_basics():
     raw.set_annotations(Annotations([1.], [.5], 'x', None))
     raws.append(raw)
     raw = concatenate_raws(raws, verbose='debug')
-    boundary_idx = np.where(raw.annotations.description == 'BAD boundary')[0]
-    assert len(boundary_idx) == 3
-    raw.annotations.delete(boundary_idx)
-    boundary_idx = np.where(raw.annotations.description == 'EDGE boundary')[0]
-    assert len(boundary_idx) == 3
-    raw.annotations.delete(boundary_idx)
+    assert_and_remove_boundary_annot(raw, 3)
     assert_array_equal(raw.annotations.onset, [124., 125., 134., 135.,
                                                144., 145., 154.])
     raw.annotations.delete(2)
     assert_array_equal(raw.annotations.onset, [124., 125., 135., 144.,
                                                145., 154.])
     raw.annotations.append(5, 1.5, 'y')
-    assert_array_equal(raw.annotations.onset, [124., 125., 135., 144.,
-                                               145., 154.,   5.])
-    assert_array_equal(raw.annotations.duration, [.5, .5, .5, .5, .5, .5, 1.5])
-    assert_array_equal(raw.annotations.description, ['x', 'x', 'x', 'x', 'x',
-                                                     'x', 'y'])
+    assert_array_equal(raw.annotations.onset,
+                       [5., 124., 125., 135., 144., 145., 154.])
+    assert_array_equal(raw.annotations.duration,
+                       [1.5, .5, .5, .5, .5, .5, .5])
+    assert_array_equal(raw.annotations.description,
+                       ['y', 'x', 'x', 'x', 'x', 'x', 'x'])
 
 
 def test_crop():
@@ -135,9 +130,7 @@ def test_crop():
                                       verbose='debug')
     assert_allclose(raw_concat.times, raw.times)
     assert_allclose(raw_concat[:][0], raw[:][0], atol=1e-20)
-    # Get rid of the boundary events
-    raw_concat.annotations.delete(-1)
-    raw_concat.annotations.delete(-1)
+    assert_and_remove_boundary_annot(raw_concat)
     # Ensure we annotations survive round-trip crop->concat
     assert_array_equal(raw_concat.annotations.description,
                        raw.annotations.description)
@@ -154,8 +147,7 @@ def test_crop():
     raw2.set_annotations(Annotations([2.], [3], 'BAD', None))
     expected_onset = [45., 2. + raw._last_time]
     raw = concatenate_raws([raw, raw2])
-    raw.annotations.delete(-1)  # remove boundary annotations
-    raw.annotations.delete(-1)
+    assert_and_remove_boundary_annot(raw)
     assert_array_almost_equal(raw.annotations.onset, expected_onset, decimal=2)
 
     # Test IO
@@ -226,14 +218,7 @@ def test_crop_more():
     assert_allclose(raw_concat.times, raw.times)
     assert_allclose(raw_concat[:][0], raw[:][0])
     assert raw_concat.first_samp == raw.first_samp
-    boundary_idx = np.where(
-        raw_concat.annotations.description == 'BAD boundary')[0]
-    assert len(boundary_idx) == 2
-    raw_concat.annotations.delete(boundary_idx)
-    boundary_idx = np.where(
-        raw_concat.annotations.description == 'EDGE boundary')[0]
-    assert len(boundary_idx) == 2
-    raw_concat.annotations.delete(boundary_idx)
+    assert_and_remove_boundary_annot(raw_concat, 2)
     assert len(raw_concat.annotations) == 4
     assert_array_equal(raw_concat.annotations.description,
                        raw.annotations.description)
@@ -872,4 +857,23 @@ def test_annotations_slices():
             annot[bad_ii]
 
 
+def test_sorting():
+    """Test annotation sorting."""
+    annot = Annotations([10, 20, 30], [1, 2, 3], 'BAD')
+    # assert_array_equal(annot.onset, [0, 5, 10])
+    annot.append([5, 15, 25, 35], 0.5, 'BAD')
+    onset = list(range(5, 36, 5))
+    duration = list(annot.duration)
+    assert_array_equal(annot.onset, onset)
+    assert_array_equal(annot.duration, duration)
+    annot.append([10, 10], [0.1, 9], 'BAD')  # 0.1 should be before, 9 after
+    want_before = onset.index(10)
+    duration.insert(want_before, 0.1)
+    duration.insert(want_before + 2, 9)
+    onset.insert(want_before, 10)
+    onset.insert(want_before, 10)
+    assert_array_equal(annot.onset, onset)
+    assert_array_equal(annot.duration, duration)
+
+
 run_tests_if_main()
diff --git a/mne/utils/testing.py b/mne/utils/testing.py
--- a/mne/utils/testing.py
+++ b/mne/utils/testing.py
@@ -402,3 +402,14 @@ def dec(*args, **kwargs):
                 raise SkipTest(msg)
             raise
     return dec
+
+
+def assert_and_remove_boundary_annot(annotations, n=1):
+    """Assert that there are boundary annotations and remove them."""
+    from ..io.base import BaseRaw
+    if isinstance(annotations, BaseRaw):  # allow either input
+        annotations = annotations.annotations
+    for key in ('EDGE', 'BAD'):
+        idx = np.where(annotations.description == '%s boundary' % key)[0]
+        assert len(idx) == n
+        annotations.delete(idx)