Repository: mne-tools/mne-python

Issue:
BUG: TFR.plot misc issues
I am experiencing several issues with the `.plot` method of TFR objects. I suspect that there are several things off, but for now I found this obvious bug:

These are valid inputs for the `title` parameter:

https://github.com/mne-tools/mne-python/blob/638847ea5539be32919b6ac4dd85ef8c4c3b79aa/mne/time_frequency/tfr.py#L1136-L1139

Yet in the code, we have this logic:

https://github.com/mne-tools/mne-python/blob/638847ea5539be32919b6ac4dd85ef8c4c3b79aa/mne/time_frequency/tfr.py#L1269-L1278

Hence, when I pass `title='auto'`, my figure will have the title `auto`




Discussion:
It appears that this is also not true:

https://github.com/mne-tools/mne-python/blob/638847ea5539be32919b6ac4dd85ef8c4c3b79aa/mne/time_frequency/tfr.py#L1183-L1185

if `combine=None` and `picks=[ch1, ch2]`, I still only get 1 figure, instead of the advertised 2 figures.
AFAIK these functions are not super widely used, so these are indeed likely bugs
yes that function is unusable for anything other than plotting a single TFR.
If you look at the code the `return` is inside the loop so it shouldn't be surprising that it doesn't work with multiple channels.
https://github.com/mne-tools/mne-python/blob/056c273f1bb96bb3229e60874122f830f946746c/mne/time_frequency/tfr.py#L1280-L1284

For what its worth it is relatively easy to write your own wrapper that will do what you need it to. I use something (https://gist.github.com/bloyl/a08a64899f2feffda5271890a5c6a677) based off the plotting in https://mne.tools/stable/auto_examples/time_frequency/plot_time_frequency_erds.html#sphx-glr-auto-examples-time-frequency-plot-time-frequency-erds-py to plot multi channel tfrs.


Thanks @bloyl that looks good. I'll give it a try tomorrow, but also attempted a fix of the plot method in #8019. I'd be happy if you could try and review if you have time.
the function also does not plot units for the colorbar label ... see #8026 
not sure, but perhaps this could be an "easy" issue for the sprint? @drammock? it would involve digging into the `.plot` method, finding all bugs and fixing them (without extending functionality), and adjusting docstrings so that they accurately reflect what is happening.
I could try some digging. 
That would be wonderful @eort 
oh boy, this is really quite a challenge to unpack :)

 a couple of design question: 

##### The `auto` value for `title`, number of picks, and `combine` are co-dependent

 `title='auto'` is supposed to concatenate channels names for up to 6 picks. This would only make sense when there are multiple picks *and* those picks are combined into a single figure (with the `combine` keyword). Otherwise, I don't see why I would want to have all picks in the title of a figure that only shows one pick. In case multiple picks are combined, I am not sure whether listing the channels is the best representation. At least in [viz/utils.py](https://github.com/mne-tools/mne-python/blob/0a61e1c7c271829bb3cc9480ea8fb08142c91e7b/mne/viz/utils.py#L1815) it seems to be to produce a title like `RMS/AVG of X electrodes` instead of just a list of channel names. 
Perhaps a combination of the two would be better? (e.g. `AVG of Ch1, Ch2, ...`, for fewer than 6 picks).

In any case, what `mne` promises `auto` will do, is only reasonable for a specific combination of parameters. What would be the best approach to deal with that? Explain it in the doc? Raising exceptions? Silently adjusting the output whenever it does not make sense (i.e. multiple picks, uncombined)?  

##### return separate figures or subplots?

Right now, the function promises to return multiple figures (even though it actually only returns a single figure). Is the idea to have a separate figure for each pick, or would it make sense to collect the tr plots in subplots of a figure and return that one? Or allow for both? If it should be possible to return multiple figures, they probably should be returned as a list of figures. 

Tomorrow during the sprint, I'll try to move further into this rabbit hole. Happy to chat then!
> At least in viz/utils.py it seems to be to produce a title like RMS/AVG of X electrodes instead of just a list of channel names.

Agreed, I like the `RMS/AVG of X electrodes` better. The combination you propose is also fine IMHO.

> what mne promises auto will do, is only reasonable for a specific combination of parameters. What would be the best approach to deal with that? Explain it in the doc? Raising exceptions? Silently adjusting the output whenever it does not make sense (i.e. multiple picks, uncombined)?

explaining it in the docstring might be the easiest choice here. If it's properly documented there, adjusting the title silently (or with a runtime warning) can be done if the parameters pushed by the user don't make sense for `title` ... 

> Right now, the function promises to return multiple figures (even though it actually only returns a single figure). Is the idea to have a separate figure for each pick, or would it make sense to collect the tr plots in subplots of a figure and return that one? Or allow for both? If it should be possible to return multiple figures, they probably should be returned as a list of figures.

I think allowing for both would be nice, returning a list of figures would be a good solution IMHO.


> Tomorrow during the sprint, I'll try to move further into this rabbit hole. Happy to chat then!

I am glad you are doing this :-) 

maybe after the sprint these lines within the mne codebase will finally spark joy again.

Resolving patch:
diff --git a/mne/time_frequency/tfr.py b/mne/time_frequency/tfr.py
--- a/mne/time_frequency/tfr.py
+++ b/mne/time_frequency/tfr.py
@@ -1184,7 +1184,7 @@ def plot(self, picks=None, baseline=None, mode='mean', tmin=None,
         mode : 'mean' | 'ratio' | 'logratio' | 'percent' | 'zscore' | 'zlogratio'
             Perform baseline correction by
 
-            - subtracting the mean of baseline values ('mean')
+            - subtracting the mean of baseline values ('mean') (default)
             - dividing by the mean of baseline values ('ratio')
             - dividing by the mean of baseline values and taking the log
               ('logratio')
@@ -1198,22 +1198,22 @@ def plot(self, picks=None, baseline=None, mode='mean', tmin=None,
 
         tmin : None | float
             The first time instant to display. If None the first time point
-            available is used.
+            available is used. Defaults to None.
         tmax : None | float
             The last time instant to display. If None the last time point
-            available is used.
+            available is used. Defaults to None.
         fmin : None | float
             The first frequency to display. If None the first frequency
-            available is used.
+            available is used. Defaults to None.
         fmax : None | float
             The last frequency to display. If None the last frequency
-            available is used.
+            available is used. Defaults to None.
         vmin : float | None
             The minimum value an the color scale. If vmin is None, the data
-            minimum value is used.
+            minimum value is used. Defaults to None.
         vmax : float | None
             The maximum value an the color scale. If vmax is None, the data
-            maximum value is used.
+            maximum value is used. Defaults to None.
         cmap : matplotlib colormap | 'interactive' | (colormap, bool)
             The colormap to use. If tuple, the first value indicates the
             colormap to use and the second value is a boolean defining
@@ -1230,19 +1230,23 @@ def plot(self, picks=None, baseline=None, mode='mean', tmin=None,
 
         dB : bool
             If True, 10*log10 is applied to the data to get dB.
+            Defaults to False.
         colorbar : bool
-            If true, colorbar will be added to the plot. For user defined axes,
-            the colorbar cannot be drawn. Defaults to True.
+            If true, colorbar will be added to the plot. Defaults to True.
         show : bool
-            Call pyplot.show() at the end.
+            Call pyplot.show() at the end. Defaults to True.
         title : str | 'auto' | None
-            String for title. Defaults to None (blank/no title). If 'auto',
-            automatically create a title that lists up to 6 of the channels
-            used in the figure.
+            String for ``title``. Defaults to None (blank/no title). If
+            'auto', and ``combine`` is None, the title for each figure
+            will be the channel name. If 'auto' and ``combine`` is not None,
+            ``title`` states how many channels were combined into that figure
+            and the method that was used for ``combine``. If str, that String
+            will be the title for each figure.
         axes : instance of Axes | list | None
             The axes to plot to. If list, the list must be a list of Axes of
-            the same length as the number of channels. If instance of Axes,
-            there must be only one channel plotted.
+            the same length as ``picks``. If instance of Axes, there must be
+            only one channel plotted. If ``combine`` is not None, ``axes``
+            must either be an instance of Axes, or a list of length 1.
         layout : Layout | None
             Layout instance specifying sensor positions. Used for interactive
             plotting of topographies on rectangle selection. If possible, the
@@ -1292,8 +1296,8 @@ def plot(self, picks=None, baseline=None, mode='mean', tmin=None,
 
         Returns
         -------
-        fig : matplotlib.figure.Figure
-            The figure containing the topography.
+        figs : list of instances of matplotlib.figure.Figure
+            A list of figures containing the time-frequency power.
         """  # noqa: E501
         return self._plot(picks=picks, baseline=baseline, mode=mode,
                           tmin=tmin, tmax=tmax, fmin=fmin, fmax=fmax,
@@ -1329,6 +1333,8 @@ def _plot(self, picks=None, baseline=None, mode='mean', tmin=None,
 
         data = tfr.data
         n_picks = len(tfr.ch_names) if combine is None else 1
+
+        # combine picks
         if combine == 'mean':
             data = data.mean(axis=0, keepdims=True)
         elif combine == 'rms':
@@ -1336,27 +1342,35 @@ def _plot(self, picks=None, baseline=None, mode='mean', tmin=None,
         elif combine is not None:
             raise ValueError('combine must be None, mean or rms.')
 
-        if isinstance(axes, list) or isinstance(axes, np.ndarray):
-            if len(axes) != n_picks:
-                raise RuntimeError('There must be an axes for each picked '
-                                   'channel.')
-
+        # figure overhead
+        # set plot dimension
         tmin, tmax = tfr.times[[0, -1]]
         if vmax is None:
             vmax = np.abs(data).max()
         if vmin is None:
             vmin = -np.abs(data).max()
-        if isinstance(axes, plt.Axes):
-            axes = [axes]
 
+        # set colorbar
         cmap = _setup_cmap(cmap)
-        for idx in range(len(data)):
-            if axes is None:
-                fig = plt.figure()
-                ax = fig.add_subplot(111)
-            else:
-                ax = axes[idx]
-                fig = ax.get_figure()
+
+        # make sure there are as many axes as there will be channels to plot
+        if isinstance(axes, list) or isinstance(axes, np.ndarray):
+            figs_and_axes = [(ax.get_figure(), ax) for ax in axes]
+        elif isinstance(axes, plt.Axes):
+            figs_and_axes = [(ax.get_figure(), ax) for ax in [axes]]
+        elif axes is None:
+            figs = [plt.figure() for i in range(n_picks)]
+            figs_and_axes = [(fig, fig.add_subplot(111)) for fig in figs]
+        else:
+            raise ValueError('axes must be None, plt.Axes, or list '
+                             'of plt.Axes.')
+        if len(figs_and_axes) != n_picks:
+            raise RuntimeError('There must be an axes for each picked '
+                               'channel.')
+
+        for idx in range(n_picks):
+            fig = figs_and_axes[idx][0]
+            ax = figs_and_axes[idx][1]
             onselect_callback = partial(
                 tfr._onselect, cmap=cmap, source_plot_joint=source_plot_joint,
                 topomap_args={k: v for k, v in topomap_args.items()
@@ -1368,22 +1382,19 @@ def _plot(self, picks=None, baseline=None, mode='mean', tmin=None,
                 yscale=yscale, mask=mask, mask_style=mask_style,
                 mask_cmap=mask_cmap, mask_alpha=mask_alpha)
 
-            if title is None:
-                if combine is None or len(tfr.info['ch_names']) == 1:
-                    title = tfr.info['ch_names'][0]
+            if title == 'auto':
+                if len(tfr.info['ch_names']) == 1 or combine is None:
+                    subtitle = tfr.info['ch_names'][idx]
                 else:
-                    title = _set_title_multiple_electrodes(
-                        title, combine, tfr.info["ch_names"], all=True,
+                    subtitle = _set_title_multiple_electrodes(
+                        None, combine, tfr.info["ch_names"], all=True,
                         ch_type=ch_type)
+            else:
+                subtitle = title
+            fig.suptitle(subtitle)
 
-            if title:
-                fig.suptitle(title)
-
-            plt_show(show)
-            # XXX This is inside the loop, guaranteeing a single iter!
-            # Also there is no None-contingent behavior here so the docstring
-            # was wrong (saying it would be collapsed)
-            return fig
+        plt_show(show)
+        return [fig for (fig, ax) in figs_and_axes]
 
     @verbose
     def plot_joint(self, timefreqs=None, picks=None, baseline=None,
@@ -1579,7 +1590,7 @@ def plot_joint(self, timefreqs=None, picks=None, baseline=None,
             colorbar=False, show=False, title=title, axes=tf_ax,
             yscale=yscale, combine=combine, exclude=None, copy=False,
             source_plot_joint=True, topomap_args=topomap_args_pass,
-            ch_type=ch_type, **image_args)
+            ch_type=ch_type, **image_args)[0]
 
         # set and check time and freq limits ...
         # can only do this after the tfr plot because it may change these


Test patch:
diff --git a/mne/time_frequency/tests/test_tfr.py b/mne/time_frequency/tests/test_tfr.py
--- a/mne/time_frequency/tests/test_tfr.py
+++ b/mne/time_frequency/tests/test_tfr.py
@@ -529,22 +529,55 @@ def test_plot():
                            ['mag', 'mag', 'mag'])
     tfr = AverageTFR(info, data=data, times=times, freqs=freqs,
                      nave=20, comment='test', method='crazy-tfr')
-    tfr.plot([1, 2], title='title', colorbar=False,
-             mask=np.ones(tfr.data.shape[1:], bool))
+
+    # test title=auto, combine=None, and correct length of figure list
+    picks = [1, 2]
+    figs = tfr.plot(picks, title='auto', colorbar=False,
+                    mask=np.ones(tfr.data.shape[1:], bool))
+    assert len(figs) == len(picks)
+    assert 'MEG' in figs[0].texts[0].get_text()
     plt.close('all')
-    ax = plt.subplot2grid((2, 2), (0, 0))
-    ax2 = plt.subplot2grid((2, 2), (1, 1))
-    ax3 = plt.subplot2grid((2, 2), (0, 1))
-    tfr.plot(picks=[0, 1, 2], axes=[ax, ax2, ax3])
+
+    # test combine and title keyword
+    figs = tfr.plot(picks, title='title', colorbar=False, combine='rms',
+                    mask=np.ones(tfr.data.shape[1:], bool))
+    assert len(figs) == 1
+    assert figs[0].texts[0].get_text() == 'title'
+    figs = tfr.plot(picks, title='auto', colorbar=False, combine='mean',
+                    mask=np.ones(tfr.data.shape[1:], bool))
+    assert len(figs) == 1
+    assert figs[0].texts[0].get_text() == 'Mean of 2 sensors'
+
+    with pytest.raises(ValueError, match='combine must be None'):
+        tfr.plot(picks, colorbar=False, combine='something',
+                 mask=np.ones(tfr.data.shape[1:], bool))
     plt.close('all')
 
-    tfr.plot([1, 2], title='title', colorbar=False, exclude='bads')
+    # test axes argument - first with list of axes
+    ax = plt.subplot2grid((2, 2), (0, 0))
+    ax2 = plt.subplot2grid((2, 2), (0, 1))
+    ax3 = plt.subplot2grid((2, 2), (1, 0))
+    figs = tfr.plot(picks=[0, 1, 2], axes=[ax, ax2, ax3])
+    assert len(figs) == len([ax, ax2, ax3])
+    # and as a single axes
+    figs = tfr.plot(picks=[0], axes=ax)
+    assert len(figs) == 1
     plt.close('all')
+    # and invalid inputs
+    with pytest.raises(ValueError, match='axes must be None'):
+        tfr.plot(picks, colorbar=False, axes={},
+                 mask=np.ones(tfr.data.shape[1:], bool))
+
+    # different number of axes and picks should throw a RuntimeError
+    with pytest.raises(RuntimeError, match='There must be an axes'):
+        tfr.plot(picks=[0], colorbar=False, axes=[ax, ax2],
+                 mask=np.ones(tfr.data.shape[1:], bool))
 
     tfr.plot_topo(picks=[1, 2])
     plt.close('all')
 
-    fig = tfr.plot(picks=[1], cmap='RdBu_r')  # interactive mode on by default
+    # interactive mode on by default
+    fig = tfr.plot(picks=[1], cmap='RdBu_r')[0]
     fig.canvas.key_press_event('up')
     fig.canvas.key_press_event(' ')
     fig.canvas.key_press_event('down')
@@ -585,7 +618,7 @@ def test_plot_joint():
 
     topomap_args = {'res': 8, 'contours': 0, 'sensors': False}
 
-    for combine in ('mean', 'rms', None):
+    for combine in ('mean', 'rms'):
         with catch_logging() as log:
             tfr.plot_joint(title='auto', colorbar=True,
                            combine=combine, topomap_args=topomap_args,
diff --git a/mne/viz/tests/test_topo.py b/mne/viz/tests/test_topo.py
--- a/mne/viz/tests/test_topo.py
+++ b/mne/viz/tests/test_topo.py
@@ -283,7 +283,7 @@ def test_plot_tfr_topo():
     freqs = np.logspace(*np.log10([3, 10]), num=3)
     tfr = AverageTFR(epochs.info, data, epochs.times, freqs, nave)
     fig = tfr.plot([4], baseline=(None, 0), mode='mean', vmax=14., show=False)
-    assert fig.axes[0].get_yaxis().get_scale() == 'log'
+    assert fig[0].axes[0].get_yaxis().get_scale() == 'log'
 
     # one timesample
     tfr = AverageTFR(epochs.info, data[:, :, [0]], epochs.times[[1]],