ENH: dict features for Epochs objects

mne/epochs.py

@@ -43,8 +43,11 @@ class Epochs(object):
     events : array, of shape [n_events, 3]
         Returned by the read_events function
 
-    event_id : int | None
-        The id of the event to consider. If None all events are used.
+    event_id : int | dict
+        The id of the event to consider. If dict,
+        the keys can later be used to acces associated events. Example:
+        dict(auditory=1, visual=3). If int, a dict will be created with
+        the str value resulting from the name parameter as key.
 
     tmin : float
         Start time before event
@@ -105,6 +108,9 @@ class Epochs(object):
     info: dict
         Measurement info
 
+    event_ids : dict
+        Labels mapped to event_ids.
+
     ch_names: list of string
         List of channels' names
 
@@ -159,10 +165,15 @@ def __init__(self, raw, events, event_id, tmin, tmax, baseline=(None, 0),
 
         self.raw = raw
         self.verbose = raw.verbose if verbose is None else verbose
-        self.event_id = event_id
+        self.name = name
+        if isinstance(event_id, dict):
+            if not all([isinstance(v, int) for v in event_id.values()]):
+                raise ValueError('Events must be of type integer')
+            self.event_ids = event_id
+        else:
+            self.event_ids = {name: event_id}
         self.tmin = tmin
         self.tmax = tmax
-        self.name = name
         self.keep_comp = keep_comp
         self.dest_comp = dest_comp
         self.baseline = baseline
@@ -209,10 +220,11 @@ def __init__(self, raw, events, event_id, tmin, tmax, baseline=(None, 0),
 
         #    Select the desired events
         self.events = events
+        self._events = {}
         if event_id is not None:
-            selected = np.logical_and(events[:, 1] == 0,
-                                      events[:, 2] == event_id)
-            self.events = self.events[selected]
+            overlap = np.in1d(events[:, 2], event_id.values())
+            selected = np.logical_and(events[:, 1] == 0, overlap)
+            self.events = events[selected]
 
         n_events = len(self.events)
 
@@ -463,26 +475,36 @@ def __repr__(self):
     def __getitem__(self, key):
         """Return an Epochs object with a subset of epochs
         """
-        if not self._bad_dropped:
-            warnings.warn("Bad epochs have not been dropped, indexing will be "
-                          "inaccurate. Use drop_bad_epochs() or preload=True")
-
-        epochs = cp.copy(self)  # XXX : should use deepcopy but breaks ...
-        epochs.events = np.atleast_2d(self.events[key])
+        epochs = self.copy()
+        if not isinstance(key, str):
+            if not self._bad_dropped:
+                warnings.warn("Bad epochs have not been dropped, indexing will"
+                              " be inaccurate. Use drop_bad_epochs() or"
+                              " preload=True")
+
+            epochs.events = np.atleast_2d(self.events[key])
+
+            if self.preload:
+                if isinstance(key, slice):
+                    epochs._data = self._data[key]
+                else:
+                    key = np.atleast_1d(key)
+                    epochs._data = self._data[key]
+        elif key not in self.event_ids:
+            raise KeyError('Event "%s" is not in Epochs.' % key)
+        else:
+            epochs.events = self._get_events(key)
+            if self.preload:
+                epochs._data[epochs.events[:, 0]]
 
-        if self.preload:
-            if isinstance(key, slice):
-                epochs._data = self._data[key]
-            else:
-                key = np.atleast_1d(key)
-                epochs._data = self._data[key]
         return epochs
 
     def average(self, picks=None):
         """Compute average of epochs
 
         Parameters
         ----------
+
         picks: None | array of int
             If None only MEG and EEG channels are kept
             otherwise the channels indices in picks are kept.
@@ -492,6 +514,7 @@ def average(self, picks=None):
         evoked : Evoked instance
             The averaged epochs
         """
+
         return self._compute_mean_or_stderr(picks, 'ave')
 
     def standard_error(self, picks=None):
@@ -510,12 +533,18 @@ def standard_error(self, picks=None):
         """
         return self._compute_mean_or_stderr(picks, 'stderr')
 
+    def _get_events(self, event_name):
+        """Aux function: get event ids"""
+        ids = None
+        if event_name in self.event_ids:
+            ids = self.events[self.events[0:, 2] == self.event_ids[event_name]]
+
+        return ids
+
     def _compute_mean_or_stderr(self, picks, mode='ave'):
         """Compute the mean or std over epochs and return Evoked"""
-        if mode == 'stderr':
-            _do_std = True
-        else:
-            _do_std = False
+
+        _do_std = True if mode == 'stderr' else False
         evoked = Evoked(None)
         evoked.info = cp.deepcopy(self.info)
         n_channels = len(self.ch_names)
@@ -643,9 +672,8 @@ def resample(self, sfreq, npad=100, window='boxcar'):
     def copy(self):
         """ Return copy of Epochs instance
         """
-        raw = self.raw.copy()
         new = deepcopy(self)
-        new.raw = raw
+        new.raw = self.raw.copy()
         return new
 
     def save(self, fname):
@@ -704,7 +732,6 @@ def save(self, fname):
         end_block(fid, FIFF.FIFFB_MEAS)
         end_file(fid)
 
-
     def as_data_frame(self, frame=True):
         """Get the epochs as Pandas panel of data frames