[MRG] Add support for indexing/slicing Annotations objects (#5800)

* Add support for indexing and slicing Annotations

* FIX: Avoid colateral effects

* [wip] documents slicing

* xx

* remove plt close

* more on slicning / iter

* fix

* use dict

* clean up

* wip: illustrate __iter__

* return a OrderedDict when iterating

* fix test to reflect iterator returns dictionary

* illustrate the fact that we have no access to Annotations attr.

* remove copy

* update tutorial

* Use slicing that returns a single Annotation

* TST: break slicing

* fix tuple

* update tests

* typo

* cosmit

* this should break if slicing returns a view and not a copy

* return copy on slice

* nitpick

* remove copy

* add iterating to whats new

* better test

* remove the test

Author: massich

doc/whats_new.rst

@@ -19,6 +19,8 @@ Current
 Changelog
 ~~~~~~~~~
 
+- Add support for indexing, slicing, and iterating :class:`mne.Annotations` by `Joan Massich`_
+
 - :meth:`mne.io.Raw.plot` now uses the lesser of ``n_channels`` and ``raw.ch_names``, by `Joan Massich`_
 
 - Add ``chunk_duration`` parameter to :func:`mne.events_from_annotations` to allow multiple events from a single annotation by `Joan Massich`_

mne/annotations.py

@@ -8,7 +8,7 @@
 import re
 from copy import deepcopy
 from itertools import takewhile
-
+from collections import OrderedDict
 
 import numpy as np
 
@@ -214,6 +214,25 @@ def __iadd__(self, other):
                                                  other.orig_time))
         return self.append(other.onset, other.duration, other.description)
 
+    def __iter__(self):
+        """Iterate over the annotations."""
+        for idx in range(len(self.onset)):
+            yield self.__getitem__(idx)
+
+    def __getitem__(self, key):
+        """Propagate indexing and slicing to the underlying numpy structure."""
+        if isinstance(key, int):
+            out_keys = ('onset', 'duration', 'description', 'orig_time')
+            out_vals = (self.onset[key], self.duration[key],
+                        self.description[key], self.orig_time)
+            return OrderedDict(zip(out_keys, out_vals))
+        else:
+            key = list(key) if isinstance(key, tuple) else key
+            return Annotations(onset=self.onset[key],
+                               duration=self.duration[key],
+                               description=self.description[key],
+                               orig_time=self.orig_time)
+
     def append(self, onset, duration, description):
         """Add an annotated segment. Operates inplace.
 
@@ -795,20 +814,17 @@ def events_from_annotations(raw, event_id=None, regexp=None, use_rounding=True,
         inds = inds[event_sel]
     else:
         inds = values = np.array([]).astype(int)
-        iterator = list(zip(annotations.onset[event_sel],
-                            annotations.duration[event_sel],
-                            annotations.description[event_sel]))
-
-        for onset, duration, description in iterator:
-            _onsets = np.arange(start=onset, stop=(onset + duration),
+        for annot in annotations[event_sel]:
+            _onsets = np.arange(start=annot['onset'],
+                                stop=(annot['onset'] + annot['duration']),
                                 step=chunk_duration)
             _inds = raw.time_as_index(_onsets,
                                       use_rounding=use_rounding,
                                       origin=annotations.orig_time)
             _inds += raw.first_samp
             inds = np.append(inds, _inds)
             _values = np.full(shape=len(_inds),
-                              fill_value=event_id_[description],
+                              fill_value=event_id_[annot['description']],
                               dtype=int)
             values = np.append(values, _values)
 

mne/tests/test_annotations.py

@@ -4,6 +4,7 @@
 
 from datetime import datetime
 from itertools import repeat
+from collections import OrderedDict
 
 import os.path as op
 
@@ -794,4 +795,71 @@ def test_read_annotation_txt_orig_time(
     assert_array_equal(annot.description, ['AA', 'BB'])
 
 
+def test_annotations_simple_iteration():
+    """Test indexing Annotations."""
+    NUM_ANNOT = 5
+    EXPECTED_ELEMENTS_TYPE = (np.float64, np.float64, np.str_)
+    EXPECTED_ONSETS = EXPECTED_DURATIONS = [x for x in range(NUM_ANNOT)]
+    EXPECTED_DESCS = [x.__repr__() for x in range(NUM_ANNOT)]
+
+    annot = Annotations(onset=EXPECTED_ONSETS,
+                        duration=EXPECTED_DURATIONS,
+                        description=EXPECTED_DESCS,
+                        orig_time=None)
+
+    for ii, elements in enumerate(annot[:2]):
+        assert isinstance(elements, OrderedDict)
+        expected_values = (ii, ii, str(ii))
+        for elem, expected_type, expected_value in zip(elements.values(),
+                                                       EXPECTED_ELEMENTS_TYPE,
+                                                       expected_values):
+            assert np.isscalar(elem)
+            assert type(elem) == expected_type
+            assert elem == expected_value
+
+
+@requires_version('numpy', '1.12')
+def test_annotations_slices():
+    """Test indexing Annotations."""
+    NUM_ANNOT = 5
+    EXPECTED_ONSETS = EXPECTED_DURATIONS = [x for x in range(NUM_ANNOT)]
+    EXPECTED_DESCS = [x.__repr__() for x in range(NUM_ANNOT)]
+
+    annot = Annotations(onset=EXPECTED_ONSETS,
+                        duration=EXPECTED_DURATIONS,
+                        description=EXPECTED_DESCS,
+                        orig_time=None)
+
+    # Indexing returns a copy. So this has no effect in annot
+    annot[0]['onset'] = 42
+    annot[0]['duration'] = 3.14
+    annot[0]['description'] = 'foobar'
+
+    annot[:1].onset[0] = 42
+    annot[:1].duration[0] = 3.14
+    annot[:1].description[0] = 'foobar'
+
+    # Slicing with single element returns a dictionary
+    for ii in EXPECTED_ONSETS:
+        assert annot[ii] == dict(zip(['onset', 'duration',
+                                      'description', 'orig_time'],
+                                     [ii, ii, str(ii), None]))
+
+    # Slices should give back Annotations
+    for current in (annot[slice(0, None, 2)],
+                    annot[[bool(ii % 2) for ii in range(len(annot))]],
+                    annot[:1],
+                    annot[[0, 2, 2]],
+                    annot[(0, 2, 2)],
+                    annot[np.array([0, 2, 2])],
+                    annot[1::2],
+                    ):
+        assert isinstance(current, Annotations)
+        assert len(current) != len(annot)
+
+    for bad_ii in [len(EXPECTED_ONSETS), 42, 'foo']:
+        with pytest.raises(IndexError):
+            annot[bad_ii]
+
+
 run_tests_if_main()

tutorials/plot_object_annotations.py

@@ -1,27 +1,12 @@
 """
-The **events** and :class:`~mne.Annotations` data structures
+The :term:`Events <events>` and :class:`~mne.Annotations` data structures
 =========================================================================
 
-Events and :class:`~mne.Annotations` are quite similar.
+:term:`Events <events>` and :term:`annotations` are quite similar.
 This tutorial highlights their differences and similarities, and tries to shed
 some light on which one is preferred to use in different situations when using
 MNE.
 
-Here are the definitions from the :ref:`glossary`.
-
-    events
-        Events correspond to specific time points in raw data; e.g., triggers,
-        experimental condition events, etc. MNE represents events with integers
-        that are stored in numpy arrays of shape (n_events, 3). Such arrays are
-        classically obtained from a trigger channel, also referred to as stim
-        channel.
-
-    annotations
-        An annotation is defined by an onset, a duration, and a string
-        description. It can contain information about the experiment, but
-        also details on signals marked by a human: bad data segments,
-        sleep scores, sleep events (spindles, K-complex) etc.
-
 Both events and :class:`~mne.Annotations` can be seen as triplets
 where the first element answers to **when** something happens and the last
 element refers to **what** it is.
@@ -101,8 +86,8 @@
 
 
 ###############################################################################
-# Working with Annotations
-# ------------------------
+# Add :term:`annotations` to :term:`raw` objects
+# ----------------------------------------------
 #
 # An important element of :class:`~mne.Annotations` is
 # ``orig_time`` which is the time reference for the ``onset``.
@@ -179,6 +164,12 @@
 print('raw_a.annotations.onset[0] is {}'.format(raw_a.annotations.onset[0]))
 
 ###############################################################################
+# Valid operations in :class:`mne.Annotations`
+# --------------------------------------------
+#
+# Concatenate
+# ~~~~~~~~~~~
+#
 # It is possible to concatenate two annotations with the + operator (just like
 # lists) if both share the same ``orig_time``
 
@@ -189,6 +180,55 @@
 print(annot)
 
 ###############################################################################
+# Iterating, Indexing and Slicing :class:`mne.Annotations`
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# :class:`~mne.Annotations` supports iterating, indexing and slicing.
+# Iterating over :class:`~mne.Annotations` and indexing with an integer returns
+# a dictionary. While slicing returns a new :class:`~mne.Annotations` instance.
+#
+# See the following examples and usages:
+
+# difference between indexing and slicing a single element
+print(annot[0])  # indexing
+print(annot[:1])  # slicing
+
+###############################################################################
+# How about iterations?
+
+for key, val in annot[0].items():  # iterate on one element which is dictionary
+    print(key, val)
+
+###############################################################################
+
+for idx, my_annot in enumerate(annot):  # iterate on the Annotations object
+    print('annot #{0}: onset={1}'.format(idx, my_annot['onset']))
+    print('annot #{0}: duration={1}'.format(idx, my_annot['duration']))
+    print('annot #{0}: description={1}'.format(idx, my_annot['description']))
+
+###############################################################################
+
+for idx, my_annot in enumerate(annot[:1]):
+    for key, val in my_annot.items():
+        print('annot #{0}: {1} = {2}'.format(idx, key, val))
+
+###############################################################################
+# Iterating, indexing and slicing return a copy. This has implications like the
+# fact that changes are not kept.
+
+# this change is not kept
+annot[0]['onset'] = 42
+print(annot[0])
+
+# this change is kept
+annot.onset[0] = 42
+print(annot[0])
+
+
+###############################################################################
+# Save
+# ~~~~
+#
 # Note that you can also save annotations to disk in FIF format::
 #
 #     >>> annot.save('my-annot.fif')