Repository: mne-tools/mne-python

Issue:
mne.viz.plot_sensors accepts ch_groups in the type list (of lists) but this is undocumented 
### Proposed documentation enhancement

The [documentation](https://mne.tools/stable/generated/mne.viz.plot_sensors.html)  of `mne.viz.plot_sensors` says that the `ch_groups` argument accepts ‘position’ | `array` of shape (n_ch_groups, n_picks) | `None`, which can be especially misleading when the user wants to create channel groups containing unequal numbers of channels. Following the documentation which specifies a fixed size of `(n_ch_groups, n_picks)`, I believe the only way to do so currently is to pass an ndarray with invalid channel indices filling up the group(s) to match `n_picks`.

I initially edited the function so that it also works with an `ndarray` of `dtype=object` with unequal numbers of items in each row, but it turns out that the function already works with a list containing uneven-sized lists e.g. `[[1, 2], [3, 4, 5, 6]]` and therefore already supports channel groups with unequal sizes. This should be made clearer in the doc. 


Discussion:
Hello! 👋 Thanks for opening your first issue here! ❤️ We will try to get back to you soon. 🚴🏽‍♂️

can you share a full code snippet? it will be clearer for me to understand

thanks
Sure. For example, if I want to use `plot_sensors` to plot the channels 1,2,3 with a certain color and 4,5,6 with another color, I would see that the documentation accepts an `ndarray` of `(n_ch_groups, n_picks)`, so I would do this: 
```
import mne
from mne.channels import make_standard_montage
from mne.viz import plot_sensors
import numpy as np

info = create_info(ch_names, sfreq=256, ch_types="eeg")
info.set_montage(montage, on_missing='ignore')
fig = plot_sensors(info, kind='3d', show_names=show_names, show=True,
                        title=title, sphere=None, ch_groups=np.array([[1, 2, 3], [4, 5, 6]]))
```
which raises a `FutureWarning` but works as intended. 

Now if I wanted to color channels 1, 2, 3 one color and 4, 5, 6, 7, 8 another color, following the documentation, one might try passing an ndarray like this:
```
fig = plot_sensors(info, kind='3d', show_names=False, show=True, sphere=None, ch_groups=np.array([[1, 2, 3], [4, 5, 6, 7, 8]], dtype=object))
```
which returns a `ValueError` due to insufficient type-checking within `plot_sensors`. Instead, using a list of list(s) works, despite list not being an accepted type in the documentation: 

```
fig = plot_sensors(info, kind='3d', show_names=False, show=True, sphere=None, ch_groups=[[1, 2, 3],[4, 5, 6, 7, 8]])
```
or, for a single group,
```
fig = plot_sensors(info, kind='3d', show_names=False, show=True, sphere=None, ch_groups=[[1, 2, 3]])
```

Therefore the [documentation](https://mne.tools/stable/generated/mne.viz.plot_sensors.html) mentioning `array` of shape (n_ch_groups, n_picks) should be changed to `array_like` instead. 

P.S. Both the `FutureWarning` and `ValueError` can be fixed by changing line 1027 from `if ch_groups in ['position', 'selection']:` to `if isinstance(ch_groups, str) and ch_groups in ['position', 'selection']:`

agree that this seems bug-like (at least, the documentation doesn't correctly specify how to do something that is already partially supported and quite reasonable to want to do). I agree that `array-like` is probably the right type spec, and we should also add something like "accepts list-of-lists or an object array to allow channel groups to have different sizes".  @mcvain do you have time to make a PR to fix this? It might be a while before one of us gets around to it.
Yes I will soon. 

Resolving patch:
diff --git a/mne/viz/utils.py b/mne/viz/utils.py
--- a/mne/viz/utils.py
+++ b/mne/viz/utils.py
@@ -924,11 +924,13 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
     show_names : bool | array of str
         Whether to display all channel names. If an array, only the channel
         names in the array are shown. Defaults to False.
-    ch_groups : 'position' | array of shape (n_ch_groups, n_picks) | None
+    ch_groups : 'position' | list of list | None
         Channel groups for coloring the sensors. If None (default), default
         coloring scheme is used. If 'position', the sensors are divided
         into 8 regions. See ``order`` kwarg of :func:`mne.viz.plot_raw`. If
-        array, the channels are divided by picks given in the array.
+        array, the channels are divided by picks given in the array. Also
+        accepts a list of lists to allow channel groups of the same or
+        different sizes.
 
         .. versionadded:: 0.13.0
     to_sphere : bool
@@ -1019,12 +1021,16 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
 
     ch_names = np.array([ch['ch_name'] for ch in chs])
     bads = [idx for idx, name in enumerate(ch_names) if name in info['bads']]
+    _validate_type(ch_groups, (list, np.ndarray, str, None), 'ch_groups')
     if ch_groups is None:
         def_colors = _handle_default('color')
         colors = ['red' if i in bads else def_colors[channel_type(info, pick)]
                   for i, pick in enumerate(picks)]
     else:
-        if ch_groups in ['position', 'selection']:
+        if isinstance(ch_groups, str):
+            _check_option(
+                'ch_groups', ch_groups, ['position', 'selection'],
+                extra='when str')
             # Avoid circular import
             from ..channels import (read_vectorview_selection, _SELECTIONS,
                                     _EEG_SELECTIONS, _divide_to_regions)
@@ -1048,13 +1054,10 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
                 x, y, z = pos[color_picks].T
                 color = np.mean(_rgb(x, y, z), axis=0)
                 color_vals[idx, :3] = color  # mean of spatial color
-        else:
+        else:  # array-like
             import matplotlib.pyplot as plt
             colors = np.linspace(0, 1, len(ch_groups))
             color_vals = [plt.cm.jet(colors[i]) for i in range(len(ch_groups))]
-        if not isinstance(ch_groups, (np.ndarray, list)):
-            raise ValueError("ch_groups must be None, 'position', "
-                             "'selection', or an array. Got %s." % ch_groups)
         colors = np.zeros((len(picks), 4))
         for pick_idx, pick in enumerate(picks):
             for ind, value in enumerate(ch_groups):


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
@@ -855,6 +855,8 @@ def test_plot_sensors(raw):
     raw.plot_sensors(ch_groups='position', axes=ax)
     raw.plot_sensors(ch_groups='selection', to_sphere=False)
     raw.plot_sensors(ch_groups=[[0, 1, 2], [3, 4]])
+    raw.plot_sensors(ch_groups=np.array([[0, 1, 2], [3, 4, 5]]))
+    raw.plot_sensors(ch_groups=np.array([[0, 1, 2], [3, 4]], dtype=object))
     pytest.raises(ValueError, raw.plot_sensors, ch_groups='asd')
     pytest.raises(TypeError, plot_sensors, raw)  # needs to be info
     pytest.raises(ValueError, plot_sensors, raw.info, kind='sasaasd')