ENH: adds annotation filtering to raw and ica source figures

Author: DerAndereJohannes

doc/changes/dev/13460.newfeature.rst

@@ -0,0 +1 @@
+Added `annotation_regex` parameter to :func:`mne.viz.raw.plot_raw` and :func:`mne.viz.ica.plot_sources`, automatically hiding annotations not matching the regex, by `Johannes Herforth`_.

mne/io/base.py

@@ -1911,6 +1911,8 @@ def plot(
         color=None,
         bad_color="lightgray",
         event_color="cyan",
+        *,
+        annotation_regex=".*",
         scalings=None,
         remove_dc=True,
         order=None,
@@ -1934,7 +1936,6 @@ def plot(
         time_format="float",
         precompute=None,
         use_opengl=None,
-        *,
         picks=None,
         theme=None,
         overview_mode=None,
@@ -1951,22 +1952,23 @@ def plot(
             color,
             bad_color,
             event_color,
-            scalings,
-            remove_dc,
-            order,
-            show_options,
-            title,
-            show,
-            block,
-            highpass,
-            lowpass,
-            filtorder,
-            clipping,
-            show_first_samp,
-            proj,
-            group_by,
-            butterfly,
-            decim,
+            annotation_regex=annotation_regex,
+            scalings=scalings,
+            remove_dc=remove_dc,
+            order=order,
+            show_options=show_options,
+            title=title,
+            show=show,
+            block=block,
+            highpass=highpass,
+            lowpass=lowpass,
+            filtorder=filtorder,
+            clipping=clipping,
+            show_first_samp=show_first_samp,
+            proj=proj,
+            group_by=group_by,
+            butterfly=butterfly,
+            decim=decim,
             noise_cov=noise_cov,
             event_id=event_id,
             show_scrollbars=show_scrollbars,

mne/preprocessing/ica.py

@@ -2585,6 +2585,7 @@ def plot_sources(
         precompute=None,
         use_opengl=None,
         *,
+        annotation_regex=".*",
         psd_args=None,
         theme=None,
         overview_mode=None,
@@ -2599,6 +2600,7 @@ def plot_sources(
             title=title,
             show=show,
             block=block,
+            annotation_regex=annotation_regex,
             psd_args=psd_args,
             show_first_samp=show_first_samp,
             show_scrollbars=show_scrollbars,

mne/viz/_figure.py

@@ -6,6 +6,7 @@
 
 import importlib
 import inspect
+import re
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from contextlib import contextmanager
@@ -182,7 +183,10 @@ def _setup_annotation_colors(self):
                 segment_colors[key] = next(color_cycle)
         self.mne.annotation_segment_colors = segment_colors
         # init a couple other annotation-related variables
-        self.mne.visible_annotations = {label: True for label in labels}
+        annot_regex = re.compile(self.mne.annotation_regex)
+        self.mne.visible_annotations = {
+            label: True if annot_regex.findall(label) else False for label in labels
+        }
         self.mne.show_hide_annotation_checkboxes = None
 
     def _update_annotation_segments(self):

mne/viz/ica.py

@@ -50,6 +50,7 @@ def plot_ica_sources(
     precompute=None,
     use_opengl=None,
     *,
+    annotation_regex=".*",
     psd_args=None,
     theme=None,
     overview_mode=None,
@@ -92,6 +93,11 @@ def plot_ica_sources(
     %(time_format)s
     %(precompute)s
     %(use_opengl)s
+    annotation_regex : str
+        A regex pattern applied to each annotation's label.
+        Matching labels remain visible, non-matching labels are hidden.
+
+        .. versionadded:: 1.11
     psd_args : dict | None
         Dictionary of arguments to pass to :meth:`~mne.Epochs.compute_psd` in
         interactive  mode. Ignored if ``inst`` is not supplied. If ``None``,
@@ -140,6 +146,7 @@ def plot_ica_sources(
             show=show,
             title=title,
             block=block,
+            annotation_regex=annotation_regex,
             psd_args=psd_args,
             show_first_samp=show_first_samp,
             show_scrollbars=show_scrollbars,
@@ -1292,6 +1299,7 @@ def _plot_sources(
     precompute,
     use_opengl,
     *,
+    annotation_regex=".*",
     psd_args,
     theme=None,
     overview_mode=None,
@@ -1415,6 +1423,7 @@ def _plot_sources(
         decim=1,
         # events
         event_times=None if is_raw else event_times,
+        annotation_regex=annotation_regex,
         # preprocessing
         projs=list(),
         projs_on=np.array([], dtype=bool),

mne/viz/raw.py

@@ -37,6 +37,8 @@ def plot_raw(
     color=None,
     bad_color="lightgray",
     event_color="cyan",
+    *,
+    annotation_regex=".*",
     scalings=None,
     remove_dc=True,
     order=None,
@@ -61,7 +63,6 @@ def plot_raw(
     precompute=None,
     use_opengl=None,
     picks=None,
-    *,
     theme=None,
     overview_mode=None,
     splash=True,
@@ -99,6 +100,11 @@ def plot_raw(
         Color to make bad channels.
     %(event_color)s
         Defaults to ``'cyan'``.
+    annotation_regex : str
+        A regex pattern applied to each annotation's label.
+        Matching labels remain visible, non-matching labels are hidden.
+
+        .. versionadded:: 1.11
     %(scalings)s
     remove_dc : bool
         If True remove DC component when plotting data.
@@ -373,6 +379,7 @@ def plot_raw(
         event_times=event_times,
         event_nums=event_nums,
         event_id_rev=event_id_rev,
+        annotation_regex=annotation_regex,
         # preprocessing
         projs=projs,
         projs_on=projs_on,

mne/viz/tests/test_ica.py

@@ -375,6 +375,22 @@ def test_plot_ica_sources(raw_orig, browser_backend, monkeypatch):
     leg = ax.get_legend()
     assert len(leg.get_texts()) == len(ica.exclude) == 1
 
+    # Check if annotation filtering works - All annotations
+    annot = Annotations([0.1, 0.3], [0.1, 0.1], ["test", "test2"])
+    raw.set_annotations(annot)
+
+    fig = ica.plot_sources(raw)
+
+    assert fig.mne.visible_annotations["test"] and fig.mne.visible_annotations["test2"]
+
+    # Check if annotation filtering works - filtering annotations
+    # This should only make test2 visible and hide test
+    fig = ica.plot_sources(raw, annotation_regex="2$")
+
+    assert (
+        not fig.mne.visible_annotations["test"] and fig.mne.visible_annotations["test2"]
+    )
+
     # test passing psd_args argument
     ica.plot_sources(epochs, psd_args=dict(fmax=50))
 

mne/viz/tests/test_raw.py

@@ -809,6 +809,23 @@ def test_plot_annotations(raw, browser_backend):
     fig._toggle_single_channel_annotation(ch_pick, 0)
     assert fig.mne.inst.annotations.ch_names[0] == (ch_pick,)
 
+    # Check if annotation filtering works - All annotations
+    annot = Annotations([42, 50], [1, 1], ["test", "test2"], raw.info["meas_date"])
+    with pytest.warns(RuntimeWarning, match="expanding outside"):
+        raw.set_annotations(annot)
+
+    fig = raw.plot()
+
+    assert fig.mne.visible_annotations["test"] and fig.mne.visible_annotations["test2"]
+
+    # Check if annotation filtering works - filtering annotations
+    # This should only make test2 visible and hide test
+    fig = raw.plot(annotation_regex="2$")
+
+    assert (
+        not fig.mne.visible_annotations["test"] and fig.mne.visible_annotations["test2"]
+    )
+
 
 @pytest.mark.parametrize("active_annot_idx", (0, 1, 2))
 def test_overlapping_annotation_deletion(raw, browser_backend, active_annot_idx):