Customize annotation colors (#13838)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: cbrnr

doc/changes/dev/13838.newfeature.rst

@@ -0,0 +1 @@
+Add ``annotation_colors`` parameter to :meth:`mne.io.Raw.plot` and :meth:`mne.Epochs.plot` to allow users to specify custom colors for annotations by passing a dict mapping annotation description strings to colors (for example, ``annotation_colors=dict(bad_segment="orange")``), by `Clemens Brunner`_.

mne/epochs.py

@@ -1321,6 +1321,7 @@ def plot(
         theme=None,
         overview_mode=None,
         splash=True,
+        annotation_colors=None,
     ):
         return plot_epochs(
             self,
@@ -1347,6 +1348,7 @@ def plot(
             theme=theme,
             overview_mode=overview_mode,
             splash=splash,
+            annotation_colors=annotation_colors,
         )
 
     @copy_function_doc_to_method_doc(plot_topo_image_epochs)

mne/io/base.py

@@ -1964,6 +1964,7 @@ def plot(
         bad_color="lightgray",
         event_color="cyan",
         *,
+        annotation_colors=None,
         annotation_regex=".*",
         scalings=None,
         remove_dc=True,
@@ -2004,6 +2005,7 @@ def plot(
             color,
             bad_color,
             event_color,
+            annotation_colors=annotation_colors,
             annotation_regex=annotation_regex,
             scalings=scalings,
             remove_dc=remove_dc,

mne/viz/_figure.py

@@ -167,16 +167,24 @@ def _get_annotation_labels(self):
 
     def _setup_annotation_colors(self):
         """Set up colors for annotations; init some annotation vars."""
+        from matplotlib.colors import to_hex
+
         segment_colors = getattr(self.mne, "annotation_segment_colors", dict())
         labels = self._get_annotation_labels()
+        user_colors = {
+            k: to_hex(v)
+            for k, v in (getattr(self.mne, "annotation_colors", None) or {}).items()
+        }
         red = "#ff0000"
         colors = _get_color_list(remove=("#fa8174", "#d62728", "#ff0000"))
         color_cycle = cycle(colors)
         for key, color in segment_colors.items():
-            if color != red and key in labels:
+            if color != red and key in labels and key not in user_colors:
                 next(color_cycle)
         for idx, key in enumerate(labels):
-            if key.lower().startswith("bad") or key.lower().startswith("edge"):
+            if key in user_colors:
+                segment_colors[key] = user_colors[key]
+            elif key.lower().startswith("bad") or key.lower().startswith("edge"):
                 segment_colors[key] = red
             elif key in segment_colors:
                 continue

mne/viz/epochs.py

@@ -29,6 +29,7 @@
     _handle_precompute,
     _make_combine_callable,
     _make_event_color_dict,
+    _normalize_annotation_colors,
     _set_title_multiple_electrodes,
     _set_window_title,
     _setup_cmap,
@@ -763,6 +764,7 @@ def plot_epochs(
     theme=None,
     overview_mode=None,
     splash=True,
+    annotation_colors=None,
 ):
     """Visualize epochs.
 
@@ -865,6 +867,14 @@ def plot_epochs(
     %(splash)s
 
         .. versionadded:: 1.6
+    annotation_colors : dict | None
+        A dictionary mapping annotation description strings to colors. Use this to
+        override the default color assigned to specific annotation types (e.g.,
+        ``dict(bad_segment='orange')``). Colors can be any valid Matplotlib color
+        specification. Keys that do not match any annotation description in the data
+        will trigger a warning. If ``None`` (default), automatic colors are used.
+
+        .. versionadded:: 1.13
 
     Returns
     -------
@@ -1014,6 +1024,13 @@ def plot_epochs(
         raise TypeError(f"title must be None or a string, got a {type(title)}")
 
     precompute = _handle_precompute(precompute)
+
+    # handle annotation_colors
+    if annotation_colors is not None:
+        annotation_colors = _normalize_annotation_colors(
+            annotation_colors, epochs.annotations
+        )
+
     params = dict(
         inst=epochs,
         info=info,
@@ -1058,6 +1075,7 @@ def plot_epochs(
         ch_color_dict=color,
         epoch_color_bad=(1, 0, 0),
         epoch_colors=epoch_colors,
+        annotation_colors=annotation_colors,
         # display
         butterfly=butterfly,
         clipping=None,

mne/viz/raw.py

@@ -11,7 +11,13 @@
 from .._fiff.pick import _picks_to_idx, pick_channels, pick_types
 from ..defaults import _handle_default
 from ..filter import create_filter
-from ..utils import _check_option, _get_stim_channel, _validate_type, legacy, verbose
+from ..utils import (
+    _check_option,
+    _get_stim_channel,
+    _validate_type,
+    legacy,
+    verbose,
+)
 from ..utils.spectrum import _split_psd_kwargs
 from .utils import (
     _check_cov,
@@ -20,6 +26,7 @@
     _handle_decim,
     _handle_precompute,
     _make_event_color_dict,
+    _normalize_annotation_colors,
     _shorten_path_from_middle,
 )
 
@@ -38,6 +45,7 @@ def plot_raw(
     bad_color="lightgray",
     event_color="cyan",
     *,
+    annotation_colors=None,
     annotation_regex=".*",
     scalings=None,
     remove_dc=True,
@@ -104,6 +112,14 @@ def plot_raw(
         Color to make bad channels.
     %(event_color)s
         Defaults to ``'cyan'``.
+    annotation_colors : dict | None
+        A dictionary mapping annotation description strings to colors. Use this to
+        override the default color assigned to specific annotation types (e.g.,
+        ``dict(bad_segment='orange')``). Colors can be any valid Matplotlib color
+        specification. Keys that do not match any annotation description in the data
+        will trigger a warning. If ``None`` (default), automatic colors are used.
+
+        .. versionadded:: 1.13
     annotation_regex : str
         A regex pattern applied to each annotation's label.
         Matching labels remain visible, non-matching labels are hidden.
@@ -335,6 +351,12 @@ def plot_raw(
     if order.size == 0:
         raise RuntimeError("No channels found to plot")
 
+    # handle annotation_colors
+    if annotation_colors is not None:
+        annotation_colors = _normalize_annotation_colors(
+            annotation_colors, raw.annotations
+        )
+
     # handle event colors
     event_color_dict = _make_event_color_dict(event_color, events, event_id)
 
@@ -399,6 +421,7 @@ def plot_raw(
         # colors
         ch_color_bad=bad_color,
         ch_color_dict=color,
+        annotation_colors=annotation_colors,
         # display
         butterfly=butterfly,
         clipping=clipping,

mne/viz/tests/test_raw.py

@@ -865,6 +865,44 @@ def test_plot_annotations(raw, browser_backend):
         assert "A" in raw.annotations.description
 
 
+def test_annotation_colors(raw, browser_backend):
+    """Test that annotation_colors overrides default colors."""
+    from matplotlib.colors import to_hex
+
+    with raw.info._unlock():
+        raw.info["lowpass"] = 10.0
+
+    raw.set_annotations(
+        Annotations(
+            onset=[1, 3, 5],
+            duration=[1, 1, 1],
+            description=["BAD_test", "BAD_other", "stimulus"],
+        )
+    )
+
+    # User-provided colors override defaults (including bad* → red rule).
+    # BAD_other has no override and should remain red.
+    fig = raw.plot(
+        annotation_colors={"BAD_test": "orange", "stimulus": "#00ff00"},
+    )
+    colors = fig.mne.annotation_segment_colors
+    assert colors["BAD_test"] == to_hex("orange"), (
+        "User color for BAD_test should override red default"
+    )
+    assert colors["stimulus"] == "#00ff00"
+    assert colors["BAD_other"] == "#ff0000", (
+        "BAD_other has no user override and should remain red"
+    )
+
+    # Unknown label key triggers a warning
+    with pytest.warns(RuntimeWarning, match="do not match"):
+        fig = raw.plot(annotation_colors={"nonexistent_label": "blue"})
+
+    # Invalid color value raises ValueError
+    with pytest.raises(ValueError, match="not a valid matplotlib color"):
+        raw.plot(annotation_colors={"BAD_test": "not_a_color"}, show=False)
+
+
 @pytest.mark.parametrize("active_annot_idx", (0, 1, 2))
 def test_overlapping_annotation_deletion(raw, browser_backend, active_annot_idx):
     """Test deletion of annotations via right-click."""

mne/viz/utils.py

@@ -2866,3 +2866,34 @@ def _get_plot_ch_type(inst, ch_type, allow_ref_meg=False):
                 f"No plottable channel types found. Allowed types are: {allowed_types}"
             )
     return ch_type
+
+
+def _normalize_annotation_colors(annotation_colors, annotations):
+    """Normalize annotation_colors and check that keys match annotation descriptions.
+
+    Parameters
+    ----------
+    annotation_colors : dict[str, color]
+        The annotation colors to normalize (``color`` can be any valid Matplotlib color
+        specification).
+    annotations : mne.Annotations
+        The Annotations object to check against.
+    """
+    from matplotlib.colors import to_hex
+
+    _validate_type(annotation_colors, dict, "annotation_colors")
+    normalized = {}
+    for k, v in annotation_colors.items():
+        try:
+            normalized[k] = to_hex(v)
+        except ValueError:
+            raise ValueError(
+                f"annotation_colors[{k!r}] is not a valid matplotlib color: {v!r}"
+            ) from None
+    unknown = set(normalized) - set(annotations.description)
+    if unknown:
+        warn(
+            "The following annotation_colors keys do not match any annotation "
+            f"description in the data: {sorted(unknown)}"
+        )
+    return normalized