Repository: mne-tools/mne-python

Issue:
raw.plot_psd cannot set picks if average is False
The following throws an exception if `average` is False, but not when it is True. 
The [documentation](https://martinos.org/mne/stable/generated/mne.io.Raw.html#mne.io.Raw.plot_psd) isn't super clear about what the expected behavior in this case is, so I'm not sure if I'm doing something unintended or not.

If `picks` can be specified when `average` is False, then this is definitely a bug. If not, then maybe updating the documentation for clarity or adding a warning message would be appropriate. I'm happy to help with either.

```python
import mne
from mne.datasets import sample

data_path = sample.data_path()
fname = data_path + "/MEG/sample/sample_audvis_raw.fif"
raw = mne.io.read_raw_fif(fname)
picks = mne.pick_types(raw.info, meg=True, eeg=False, stim=False)

# Changing average to True makes it work
raw.plot_psd(picks=picks, dB=False, show=True, average=False)
```

Results in stack trace
```
/miniconda3/envs/mne/lib/python2.7/site-packages/mne/viz/evoked.py:433: VisibleDeprecationWarning: boolean index did not match indexed array along dimension 0; dimension is 1 but corresponding boolean dimension is 2
  for ax in np.array(axes)[selectables]:
Traceback (most recent call last):
  File "psd_bug.py", line 10, in <module>
    ax = raw.plot_psd(picks=picks, dB=False, show=True, average=False)
  File "<string>", line 2, in plot_psd
  File "/home/erik/miniconda3/envs/mne/lib/python2.7/site-packages/mne/utils.py", line 728, in verbose
    return function(*args, **kwargs)
  File "/home/erik/miniconda3/envs/mne/lib/python2.7/site-packages/mne/io/base.py", line 1779, in plot_psd
    xscale=xscale, reject_by_annotation=reject_by_annotation)
  File "<string>", line 2, in plot_raw_psd
  File "/home/erik/miniconda3/envs/mne/lib/python2.7/site-packages/mne/utils.py", line 728, in verbose
    return function(*args, **kwargs)
  File "/home/erik/miniconda3/envs/mne/lib/python2.7/site-packages/mne/viz/raw.py", line 812, in plot_raw_psd
    line_alpha=line_alpha)
  File "/home/erik/miniconda3/envs/mne/lib/python2.7/site-packages/mne/viz/evoked.py", line 433, in _plot_lines
    for ax in np.array(axes)[selectables]:
IndexError: index 1 is out of bounds for axis 1 with size 1

```


Discussion:
sounds like an easy pb to solve for a new contributor?

Resolving patch:
diff --git a/mne/viz/evoked.py b/mne/viz/evoked.py
--- a/mne/viz/evoked.py
+++ b/mne/viz/evoked.py
@@ -307,6 +307,7 @@ def _plot_lines(data, info, picks, fig, axes, spatial_colors, unit, units,
     """Plot data as butterfly plot."""
     from matplotlib import patheffects
     from matplotlib.widgets import SpanSelector
+    assert len(axes) == len(ch_types_used)
     texts = list()
     idxs = list()
     lines = list()
@@ -2069,7 +2070,7 @@ def plot_compare_evokeds(evokeds, picks=None, gfp=False, colors=None,
                 raise TypeError("show_sensors must be numeric, str or bool, "
                                 "not " + str(type(show_sensors)))
             show_sensors = _check_loc_legal(show_sensors, "show_sensors")
-            _plot_legend(pos, ["k" for _ in picks], ax, list(), outlines,
+            _plot_legend(pos, ["k"] * len(picks), ax, list(), outlines,
                          show_sensors, size=25)
 
     # the condition legend
diff --git a/mne/viz/raw.py b/mne/viz/raw.py
--- a/mne/viz/raw.py
+++ b/mne/viz/raw.py
@@ -543,44 +543,41 @@ def _set_psd_plot_params(info, proj, picks, ax, area_mode):
     import matplotlib.pyplot as plt
     if area_mode not in [None, 'std', 'range']:
         raise ValueError('"area_mode" must be "std", "range", or None')
-    if picks is None:
-        # XXX this could be refactored more with e.g., plot_evoked
-        megs = ['mag', 'grad', False, False, False]
-        eegs = [False, False, True, False, False]
-        seegs = [False, False, False, True, False]
-        ecogs = [False, False, False, False, True]
-        names = ['mag', 'grad', 'eeg', 'seeg', 'ecog']
-        titles = _handle_default('titles', None)
-        units = _handle_default('units', None)
-        scalings = _handle_default('scalings', None)
-        picks_list = list()
-        titles_list = list()
-        units_list = list()
-        scalings_list = list()
-        for meg, eeg, seeg, ecog, name in zip(megs, eegs, seegs, ecogs, names):
-            picks = pick_types(info, meg=meg, eeg=eeg, seeg=seeg, ecog=ecog,
-                               ref_meg=False)
-            if len(picks) > 0:
-                picks_list.append(picks)
-                titles_list.append(titles[name])
-                units_list.append(units[name])
-                scalings_list.append(scalings[name])
-        if len(picks_list) == 0:
-            raise RuntimeError('No data channels found')
-        if ax is not None:
-            if isinstance(ax, plt.Axes):
-                ax = [ax]
-            if len(ax) != len(picks_list):
-                raise ValueError('For this dataset with picks=None %s axes '
-                                 'must be supplied, got %s'
-                                 % (len(picks_list), len(ax)))
-            ax_list = ax
-    else:
-        picks_list = [picks]
-        titles_list = ['Selected channels']
-        units_list = ['amplitude']
-        scalings_list = [1.]
-        ax_list = [ax]
+
+    # XXX this could be refactored more with e.g., plot_evoked
+    megs = ['mag', 'grad', False, False, False]
+    eegs = [False, False, True, False, False]
+    seegs = [False, False, False, True, False]
+    ecogs = [False, False, False, False, True]
+    names = ['mag', 'grad', 'eeg', 'seeg', 'ecog']
+    titles = _handle_default('titles', None)
+    units = _handle_default('units', None)
+    scalings = _handle_default('scalings', None)
+    picks_list = list()
+    titles_list = list()
+    units_list = list()
+    scalings_list = list()
+    for meg, eeg, seeg, ecog, name in zip(megs, eegs, seegs, ecogs, names):
+        these_picks = pick_types(info, meg=meg, eeg=eeg, seeg=seeg, ecog=ecog,
+                                 ref_meg=False)
+        if picks is not None:
+            these_picks = np.intersect1d(these_picks, picks)
+        if len(these_picks) > 0:
+            picks_list.append(these_picks)
+            titles_list.append(titles[name])
+            units_list.append(units[name])
+            scalings_list.append(scalings[name])
+    if len(picks_list) == 0:
+        raise RuntimeError('No data channels found')
+    if ax is not None:
+        if isinstance(ax, plt.Axes):
+            ax = [ax]
+        if len(ax) != len(picks_list):
+            raise ValueError('For this dataset with picks=None %s axes '
+                             'must be supplied, got %s'
+                             % (len(picks_list), len(ax)))
+        ax_list = ax
+    del picks
 
     make_label = False
     fig = None
@@ -827,6 +824,7 @@ def plot_raw_psd(raw, tmin=0., tmax=np.inf, fmin=0, fmax=np.inf, proj=False,
         for this_type in valid_channel_types:
             if this_type in types:
                 ch_types_used.append(this_type)
+        assert len(ch_types_used) == len(ax_list)
         unit = ''
         units = {t: yl for t, yl in zip(ch_types_used, ylabels)}
         titles = {c: t for c, t in zip(ch_types_used, titles_list)}


Test patch:
diff --git a/mne/viz/tests/test_raw.py b/mne/viz/tests/test_raw.py
--- a/mne/viz/tests/test_raw.py
+++ b/mne/viz/tests/test_raw.py
@@ -290,6 +290,13 @@ def test_plot_raw_psd():
     raw.plot_psd(reject_by_annotation=True)
     raw.plot_psd(reject_by_annotation=False)
 
+    # gh-5046
+    raw = read_raw_fif(raw_fname, preload=True).crop(0, 1)
+    picks = pick_types(raw.info)
+    raw.plot_psd(picks=picks, average=False)
+    raw.plot_psd(picks=picks, average=True)
+    plt.close('all')
+
 
 def test_plot_sensors():
     """Test plotting of sensor array."""