[med-svn] [python-mne] 44/353: ENH : adding pick_events function

[Yaroslav Halchenko]

This is an automated email from the git hooks/post-receive script.

yoh pushed a commit to tag 0.4
in repository python-mne.

commit 3b273f6ab90e5d33db6d6a22b114c3d1c26673d6
Author: Alexandre Gramfort <alexandre.gramfort at inria.fr>
Date:   Fri Jan 6 08:41:47 2012 +0100

    ENH : adding pick_events function
---
 examples/read_events.py |  5 ++++-
 mne/__init__.py         |  3 ++-
 mne/event.py            | 55 +++++++++++++++++++++++++++++++++++++++++++++++--
 mne/tests/test_event.py | 10 ++++++++-
 4 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/examples/read_events.py b/examples/read_events.py
index 2e3fbb3..9c3c2b6 100644
--- a/examples/read_events.py
+++ b/examples/read_events.py
@@ -16,7 +16,10 @@ data_path = sample.data_path('.')
 fname = data_path + '/MEG/sample/sample_audvis_raw-eve.fif'
 
 # Reading events
-events = mne.read_events(fname)
+# events = mne.read_events(fname)  # all
+events = mne.read_events(fname, include=1)  # restricted to event 1
+events = mne.read_events(fname, include=[1, 2])  # restricted to event 1 or 2
+events = mne.read_events(fname, exclude=[4, 32])  # keep all but 4 and 32
 
 # Writing events
 mne.write_events('events.fif', events)
diff --git a/mne/__init__.py b/mne/__init__.py
index 9679ff5..07e8088 100644
--- a/mne/__init__.py
+++ b/mne/__init__.py
@@ -2,7 +2,8 @@ __version__ = '0.2.git'
 
 from .cov import read_cov, write_cov, write_cov_file, Covariance, \
                  compute_raw_data_covariance, compute_covariance
-from .event import read_events, write_events, find_events, merge_events
+from .event import read_events, write_events, find_events, merge_events, \
+                   pick_events
 from .forward import read_forward_solution
 from .source_estimate import read_stc, write_stc, read_w, write_w, \
                              SourceEstimate, morph_data, \
diff --git a/mne/event.py b/mne/event.py
index fefd59a..a2d67ed 100644
--- a/mne/event.py
+++ b/mne/event.py
@@ -16,13 +16,61 @@ from .fiff.write import write_int, start_block, start_file, end_block, end_file
 from .fiff.pick import pick_channels
 
 
-def read_events(filename):
+def pick_events(events, include=None, exclude=None):
+    """Select some events
+
+    Parameters
+    ----------
+    include: int | list | None
+        A event id to include or a list of them.
+        If None all events are included.
+    exclude: int | list | None
+        A event id to exclude or a list of them.
+        If None no event is excluded. If include is not None
+        the exclude parameter is ignored.
+
+    Returns
+    -------
+    events: array, shape (n_events, 3)
+        The list of events
+    """
+    if include is not None:
+        if not isinstance(include, list):
+            include = [include]
+        mask = np.zeros(len(events), dtype=np.bool)
+        for e in include:
+            mask = np.logical_or(mask, events[:, 2] == e)
+        events = events[mask]
+    elif exclude is not None:
+        if not isinstance(exclude, list):
+            exclude = [exclude]
+        mask = np.ones(len(events), dtype=np.bool)
+        for e in exclude:
+            mask = np.logical_and(mask, events[:, 2] != e)
+        events = events[mask]
+    else:
+        events = np.copy(events)
+
+    if len(events) == 0:
+        raise RuntimeError("No events found")
+
+    return events
+
+
+def read_events(filename, include=None, exclude=None):
     """Reads events from fif file
 
     Parameters
     ----------
     filename: string
         name of the fif file
+    include: int | list | None
+        A event id to include or a list of them.
+        If None all events are included.
+    exclude: int | list | None
+        A event id to exclude or a list of them.
+        If None no event is excluded. If include is not None
+        the exclude parameter is ignored.
 
     Returns
     -------
@@ -54,6 +102,9 @@ def read_events(filename):
         raise ValueError('Could not find any events')
 
     event_list = event_list.reshape(len(event_list) / 3, 3)
+
+    event_list = pick_events(event_list, include, exclude)
+
     return event_list
 
 
@@ -129,4 +180,4 @@ def merge_events(events, ids, new_id):
     events_numbers = events[:, 2]
     for i in ids:
         events_numbers[events_numbers == i] = new_id
-    return events
\ No newline at end of file
+    return events
diff --git a/mne/tests/test_event.py b/mne/tests/test_event.py
index 965fb61..ccb21bc 100644
--- a/mne/tests/test_event.py
+++ b/mne/tests/test_event.py
@@ -1,6 +1,6 @@
 import os.path as op
 
-from numpy.testing import assert_array_almost_equal
+from numpy.testing import assert_array_almost_equal, assert_array_equal
 
 from .. import read_events, write_events, find_events
 from .. import fiff
@@ -21,6 +21,14 @@ def test_io_events():
     events2 = read_events('events.fif')
     assert_array_almost_equal(events, events2)
 
+    a = read_events('events.fif', include=1)
+    b = read_events('events.fif', include=[1])
+    c = read_events('events.fif', exclude=[2, 3, 4, 5, 32])
+    d = read_events('events.fif', include=1, exclude=[2, 3])
+    assert_array_equal(a, b)
+    assert_array_equal(a, c)
+    assert_array_equal(a, d)
+
 
 def test_find_events():
     """Test find events in raw file

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-med/python-mne.git