Repository: mne-tools/mne-python

Issue:
visualization of autoreject results in MNE
Currently, the `autoreject` code hacks into MNE by copying the epochs plotting code and adding [two parameters](https://github.com/autoreject/autoreject/blob/master/autoreject/viz.py#L123). Unfortunately this is a [bit unwieldy to maintain](https://github.com/autoreject/autoreject/issues/153) on the `autoreject` side as we're using a ton of private functions in MNE and duplicating almost 1000 lines of code ... 

are people open to adding these parameters to `epochs.plot` in MNE, or alternatively some code to interpret the `drop_log` in a special way if it starts with the keyword `AUTOREJECT`. I can make a PR with the existing diff and iterate the API from there if people are okay with this.


Discussion:
no objection from me.  I can work on this with you at the upcoming sprint if you want...  it may have been my recent changes to the plotting code that caused this in the first place
- the first param, `bad_epoch_idx` -- I'm not sure why we'd need it, can't you just do `epochs[bad_epoch_idx].plot()`? If not, can you describe what it's trying to do?
-  the second param `log_labels` -- it seems like its job is to color specific parts of the `(n_channels, n_epochs)` plotting grid. So maybe the thing to do is make a new param called `epoch_colors` that defaults to `'auto'`, but can be any `ndarray` of `str` that broadcasts to shape `(n_channels, n_epochs)`.
    - Each entry in the array is expected to be an `mpl-like` color string. (This is general enough to support arbitrary colors via the mpl hex strings, even if it's a bit cumbersome).
   - if they want to color squares in the plane (like you do) they pass an array of shape `(n_channels, n_epochs)`.
   - if people want to color specific channels they can pass `(n_channels,1)` array
   - if they want to color specific epochs the can pass `(n_epochs,)` array

I think this would be more general than just what's required for `autoreject` without adding too much more work.
> can't you just do epochs[bad_epoch_idx].plot()

you could but I think it's nicer for people to see them next to the good epochs in red as you see them when you mark them interactively

I'm fine making `log_labels` more general than `autoreject`. I guess we could just make it so that the horizontal scrollbar colors the epoch if a column in `epoch_colors` have the same repeated value. Wdyt?
> you could but I think it's nicer for people to see them next to the good epochs in red as you see them when you mark them interactively

What does `bad_epoch_idx` actually do, then? If all it does is change the color, then that would already be handled by the proposed `epoch_colors` arg.

> I guess we could just make it so that the horizontal scrollbar colors the epoch if a column in epoch_colors have the same repeated value. Wdyt?

Whatever ends up looking reasonable. An average over the colors for the given epoch might be fine, too.
> What does bad_epoch_idx actually do, then? If all it does is change the color, then that would already be handled by the proposed epoch_colors arg.

yep I think so

Resolving patch:
diff --git a/mne/epochs.py b/mne/epochs.py
--- a/mne/epochs.py
+++ b/mne/epochs.py
@@ -1042,14 +1042,15 @@ def ch_names(self):
     def plot(self, picks=None, scalings=None, n_epochs=20, n_channels=20,
              title=None, events=None, event_colors=None, order=None,
              show=True, block=False, decim='auto', noise_cov=None,
-             butterfly=False, show_scrollbars=True):
+             butterfly=False, show_scrollbars=True, epoch_colors=None):
         return plot_epochs(self, picks=picks, scalings=scalings,
                            n_epochs=n_epochs, n_channels=n_channels,
                            title=title, events=events,
                            event_colors=event_colors, order=order,
                            show=show, block=block, decim=decim,
                            noise_cov=noise_cov, butterfly=butterfly,
-                           show_scrollbars=show_scrollbars)
+                           show_scrollbars=show_scrollbars,
+                           epoch_colors=epoch_colors)
 
     @copy_function_doc_to_method_doc(plot_epochs_psd)
     def plot_psd(self, fmin=0, fmax=np.inf, tmin=None, tmax=None,
diff --git a/mne/viz/epochs.py b/mne/viz/epochs.py
--- a/mne/viz/epochs.py
+++ b/mne/viz/epochs.py
@@ -18,9 +18,10 @@
 import numpy as np
 
 from ..defaults import _handle_default
-from ..utils import (verbose, logger, warn, fill_doc, check_version,
-                     _validate_type)
-from ..io.meas_info import create_info
+
+from ..utils import verbose, logger, warn, fill_doc, check_version
+from ..io.meas_info import create_info, _validate_type
+
 from ..io.pick import (pick_types, channel_type, _get_channel_types,
                        _picks_to_idx, _DATA_CH_TYPES_SPLIT,
                        _DATA_CH_TYPES_ORDER_DEFAULT, _VALID_CHANNEL_TYPES)
@@ -723,7 +724,7 @@ def _epochs_axes_onclick(event, params):
 def plot_epochs(epochs, picks=None, scalings=None, n_epochs=20, n_channels=20,
                 title=None, events=None, event_colors=None, order=None,
                 show=True, block=False, decim='auto', noise_cov=None,
-                butterfly=False, show_scrollbars=True):
+                butterfly=False, show_scrollbars=True, epoch_colors=None):
     """Visualize epochs.
 
     Bad epochs can be marked with a left click on top of the epoch. Bad
@@ -808,6 +809,8 @@ def plot_epochs(epochs, picks=None, scalings=None, n_epochs=20, n_channels=20,
 
         .. versionadded:: 0.18.0
     %(show_scrollbars)s
+    epoch_colors : list of (n_epochs) list (of n_channels) | None
+        Colors to use for individual epochs. If None, use default colors.
 
     Returns
     -------
@@ -841,7 +844,8 @@ def plot_epochs(epochs, picks=None, scalings=None, n_epochs=20, n_channels=20,
                   bad_color=(0.8, 0.8, 0.8), histogram=None, decim=decim,
                   data_picks=data_picks, noise_cov=noise_cov,
                   use_noise_cov=noise_cov is not None,
-                  show_scrollbars=show_scrollbars)
+                  show_scrollbars=show_scrollbars,
+                  epoch_colors=epoch_colors)
     params['label_click_fun'] = partial(_pick_bad_channels, params=params)
     _prepare_mne_browse_epochs(params, projs, n_channels, n_epochs, scalings,
                                title, picks, events=events, order=order,
@@ -978,6 +982,22 @@ def _prepare_mne_browse_epochs(params, projs, n_channels, n_epochs, scalings,
         raise RuntimeError('Some channels not classified. Please'
                            ' check your picks')
     ch_names = [params['info']['ch_names'][idx] for idx in inds]
+    _validate_type(params['epoch_colors'], (list, None), 'epoch_colors')
+    if params['epoch_colors'] is not None:
+        if len(params['epoch_colors']) != len(params['epochs'].events):
+            raise ValueError('epoch_colors must be list of len(epochs.events).'
+                             ' Got %s' % len(params['epoch_colors']))
+        for epoch_idx in range(len(params['epoch_colors'])):
+            these_colors = params['epoch_colors'][epoch_idx]
+            _validate_type(these_colors, list,
+                           'epoch_colors[%s]' % (epoch_idx,))
+            if len(these_colors) != len(params['epochs'].ch_names):
+                raise ValueError('epoch_colors for the %dth epoch '
+                                 'has length %d, expected %d.'
+                                 % (epoch_idx, len(these_colors),
+                                    len(params['epochs'].ch_names)))
+            params['epoch_colors'][epoch_idx] = \
+                [these_colors[idx] for idx in inds]
 
     # set up plotting
     n_epochs = min(n_epochs, len(epochs_events))
@@ -1139,6 +1159,21 @@ def _prepare_mne_browse_epochs(params, projs, n_channels, n_epochs, scalings,
 
     params['plot_fun'] = partial(_plot_traces, params=params)
 
+    # Plot epoch_colors
+    if params['epoch_colors'] is not None:
+        for epoch_idx, epoch_color in enumerate(params['epoch_colors']):
+            for ch_idx in range(len(params['ch_names'])):
+                if epoch_color[ch_idx] is not None:
+                    params['colors'][ch_idx][epoch_idx] = \
+                        colorConverter.to_rgba(epoch_color[ch_idx])
+
+            # plot on horizontal patch if all colors are same
+            if epoch_color.count(epoch_color[0]) == len(epoch_color):
+                params['ax_hscroll'].patches[epoch_idx].set_color(
+                    epoch_color[0])
+                params['ax_hscroll'].patches[epoch_idx].set_zorder(3)
+                params['ax_hscroll'].patches[epoch_idx].set_edgecolor('w')
+
     # callbacks
     callback_scroll = partial(_plot_onscroll, params=params)
     params['fig'].canvas.mpl_connect('scroll_event', callback_scroll)
diff --git a/mne/viz/ica.py b/mne/viz/ica.py
--- a/mne/viz/ica.py
+++ b/mne/viz/ica.py
@@ -990,7 +990,8 @@ def _plot_sources_epochs(ica, epochs, picks, exclude, start, stop, show,
                   bads=list(), bad_color=(1., 0., 0.),
                   t_start=start * len(epochs.times),
                   data_picks=list(), decim=1, whitened_ch_names=(),
-                  noise_cov=None, show_scrollbars=show_scrollbars)
+                  noise_cov=None, show_scrollbars=show_scrollbars,
+                  epoch_colors=None)
     params['label_click_fun'] = partial(_label_clicked, params=params)
     # changing the order to 'misc' before 'eog' and 'ecg'
     order = list(_DATA_CH_TYPES_ORDER_DEFAULT)


Test patch:
diff --git a/mne/viz/tests/test_epochs.py b/mne/viz/tests/test_epochs.py
--- a/mne/viz/tests/test_epochs.py
+++ b/mne/viz/tests/test_epochs.py
@@ -83,6 +83,21 @@ def test_plot_epochs_basic(capsys):
     with pytest.warns(RuntimeWarning, match='projection'):
         epochs.plot(noise_cov=cov)
     plt.close('all')
+    # check the epochs color plotting
+    epoch_colors = [['r'] * len(epochs.ch_names) for _ in
+                    range(len(epochs.events))]
+    epochs.plot(epoch_colors=epoch_colors)
+    with pytest.raises(TypeError, match='must be an instance of'):
+        epochs.plot(epoch_colors='r')
+    with pytest.raises(ValueError, match='must be list of len'):
+        epochs.plot(epoch_colors=['r'] * len(epochs.ch_names))
+    epoch_colors[0] = None
+    with pytest.raises(TypeError, match='must be an instance of'):
+        epochs.plot(epoch_colors=epoch_colors)
+    epoch_colors[0] = ['r'] * 5
+    with pytest.raises(ValueError, match='epoch_colors for the 0th epoch has'
+                       ' length'):
+        epochs.plot(epoch_colors=epoch_colors)
     # other options
     fig = epochs[0].plot(picks=[0, 2, 3], scalings=None)
     fig.canvas.key_press_event('escape')