[MRG] Forward argument axes from plot_sensors to DigMontage.plot

doc/changes/latest.inc

@@ -26,6 +26,7 @@ Enhancements
 - Adjusted the algorithm used in :class:`mne.decoding.SSD` to support non-full rank data (:gh:`11458` by :newcontrib:`Thomas Binns`)
 - Changed suggested type for ``ch_groups``` in `mne.viz.plot_sensors` from array to list of list(s) (arrays are still supported). (:gh:`11465` by `Hyonyoung Shin`_)
 - Add support for UCL/FIL OPM data using :func:`mne.io.read_raw_fil` (:gh:`11366` by :newcontrib:`George O'Neill` and `Robert Seymour`_)
+- Forward argument ``axes`` from `mne.viz.plot_sensors` to `mne.channels.DigMontage.plot` (:gh:`11470` by :newcontrib:`Jan Ebert` and `Mathieu Scheltienne`_)
 - Added ability to read stimulus durations from SNIRF files when using :func:`mne.io.read_raw_snirf` (:gh:`11397` by `Robert Luke`_)
 - Add :meth:`mne.Info.save` to save an :class:`mne.Info` object to a fif file (:gh:`11401` by `Alex Rockhill`_)
 - Improved error message when downloads are corrupted for :func:`mne.datasets.sample.data_path` and related functions (:gh:`11407` by `Eric Larson`_)

doc/changes/names.inc

@@ -116,6 +116,8 @@
 
 .. _Denis Engemann: http://denis-engemann.de
 
+.. _Dinara Issagaliyeva: https://github.com/dissagaliyeva
+
 .. _Dirk Gütlin: https://github.com/DiGyt
 
 .. _Dmitrii Altukhov: https://github.com/dmalt
@@ -178,6 +180,8 @@
 
 .. _Hamid Maymandi: https://github.com/HamidMandi
 
+.. _Hakimeh Pourakbari: https://github.com/Hpakbari
+
 .. _Hari Bharadwaj: http://www.haribharadwaj.com
 
 .. _Henrich Kolkhorst: https://github.com/hekolk
@@ -200,6 +204,8 @@
 
 .. _Jair Montoya Martinez: https://github.com/jmontoyam
 
+.. _Jan Ebert: https://www.jan-ebert.com/
+
 .. _Jan Sedivy: https://github.com/honzaseda
 
 .. _Jan Sosulski: https://jan-sosulski.de
@@ -214,6 +220,8 @@
 
 .. _Jeff Stout: https://megcore.nih.gov/index.php/Staff
 
+.. _Jennifer Behnke: https://github.com/JKBehnke
+
 .. _Jeroen Van Der Donckt: https://github.com/jvdd
 
 .. _Jesper Duemose Nielsen: https://github.com/jdue
@@ -378,6 +386,8 @@
 
 .. _Paul Roujansky: https://github.com/paulroujansky
 
+.. _Pavel Navratil: https://github.com/navrpa13
+
 .. _Peter Molfese: https://github.com/pmolfese
 
 .. _Phillip Alday: https://palday.bitbucket.io
@@ -492,6 +502,8 @@
 
 .. _Theodore Papadopoulo: https://github.com/papadop
 
+.. _Thomas Binns: https://github.com/tsbinns
+
 .. _Thomas Hartmann: https://github.com/thht
 
 .. _Thomas Radman: https://github.com/tradman
@@ -500,6 +512,8 @@
 
 .. _Tod Flak: https://github.com/todflak
 
+.. _Tom Ma: https://github.com/myd7349
+
 .. _Tommy Clausner: https://github.com/TommyClausner
 
 .. _Toomas Erik Anijärv: http://www.toomaserikanijarv.com/
@@ -523,15 +537,3 @@
 .. _Yu-Han Luo: https://github.com/yh-luo
 
 .. _Zhi Zhang: https://github.com/tczhangzhi/
-
-.. _Dinara Issagaliyeva: https://github.com/dissagaliyeva
-
-.. _Jennifer Behnke: https://github.com/JKBehnke
-
-.. _Hakimeh Pourakbari: https://github.com/Hpakbari
-
-.. _Pavel Navratil: https://github.com/navrpa13
-
-.. _Tom Ma: https://github.com/myd7349
-
-.. _Thomas Binns: https://github.com/tsbinns

mne/channels/montage.py

@@ -328,10 +328,10 @@ def __repr__(self):
 
     @copy_function_doc_to_method_doc(plot_montage)
     def plot(self, scale_factor=20, show_names=True, kind='topomap', show=True,
-             sphere=None, verbose=None):
+             sphere=None, *, axes=None, verbose=None):
         return plot_montage(self, scale_factor=scale_factor,
                             show_names=show_names, kind=kind, show=show,
-                            sphere=sphere)
+                            sphere=sphere, axes=axes)
 
     @fill_doc
     def rename_channels(self, mapping, allow_duplicates=False):

mne/channels/tests/test_montage.py

@@ -1644,6 +1644,17 @@ def test_plot_montage():
     montage.plot()
     plt.close('all')
 
+    f, ax = plt.subplots(1, 1)
+    montage.plot(axes=ax)
+    plt.close("all")
+
+    with pytest.raises(TypeError, match="must be an instance of Axes"):
+        montage.plot(axes=101)
+    with pytest.raises(TypeError, match="when 'kind' is '3d'"):
+        montage.plot(axes=ax, kind="3d")
+    with pytest.raises(TypeError, match="when 'kind' is '3d'"):
+        montage.plot(axes=101, kind="3d")
+
 
 def test_montage_equality():
     """Test montage equality."""

mne/utils/docs.py

@@ -280,6 +280,10 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75):
 docdict['axes_evoked_plot_topomap'] = _axes_list.format(
     'axes',
     'match the number of ``times`` provided (unless ``times`` is ``None``)')
+docdict['axes_montage'] = """
+axes : instance of Axes | instance of Axes3D | None
+    Axes to draw the sensors to. If ``kind='3d'``, axes must be an instance
+    of Axes3D. If None (default), a new axes will be created."""
 docdict['axes_plot_projs_topomap'] = _axes_list.format(
     'axes', 'match the number of projectors')
 docdict['axes_plot_topomap'] = _axes_base.format('axes', '', '', '')

mne/viz/montage.py

@@ -8,7 +8,7 @@
 
 @verbose
 def plot_montage(montage, scale_factor=20, show_names=True, kind='topomap',
-                 show=True, sphere=None, verbose=None):
+                 show=True, sphere=None, *, axes=None, verbose=None):
     """Plot a montage.
 
     Parameters
@@ -25,6 +25,9 @@ def plot_montage(montage, scale_factor=20, show_names=True, kind='topomap',
     show : bool
         Show figure if True.
     %(sphere_topomap_auto)s
+    %(axes_montage)s
+
+        .. versionadded:: 1.4
     %(verbose)s
 
     Returns
@@ -69,7 +72,7 @@ def plot_montage(montage, scale_factor=20, show_names=True, kind='topomap',
     info = create_info(ch_names, sfreq=256, ch_types="eeg")
     info.set_montage(montage, on_missing='ignore')
     fig = plot_sensors(info, kind=kind, show_names=show_names, show=show,
-                       title=title, sphere=sphere)
+                       title=title, sphere=sphere, axes=axes)
     collection = fig.axes[0].collections[0]
     collection.set_sizes([scale_factor])
     return fig

mne/viz/utils.py

@@ -908,16 +908,17 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
     %(info_not_none)s
     kind : str
         Whether to plot the sensors as 3d, topomap or as an interactive
-        sensor selection dialog. Available options 'topomap', '3d', 'select'.
-        If 'select', a set of channels can be selected interactively by using
-        lasso selector or clicking while holding control key. The selected
-        channels are returned along with the figure instance. Defaults to
-        'topomap'.
+        sensor selection dialog. Available options ``'topomap'``, ``'3d'``,
+        ``'select'``. If ``'select'``, a set of channels can be selected
+        interactively by using lasso selector or clicking while holding control
+        key. The selected channels are returned along with the figure instance.
+        Defaults to ``'topomap'``.
     ch_type : None | str
-        The channel type to plot. Available options 'mag', 'grad', 'eeg',
-        'seeg', 'dbs', 'ecog', 'all'. If ``'all'``, all the available mag,
-        grad, eeg, seeg, dbs and ecog channels are plotted. If None (default),
-        then channels are chosen in the order given above.
+        The channel type to plot. Available options ``'mag'``, ``'grad'``,
+        ``'eeg'``, ``'seeg'``, ``'dbs'``, ``'ecog'``, ``'all'``. If ``'all'``,
+        all the available mag, grad, eeg, seeg, dbs and ecog channels are
+        plotted. If None (default), then channels are chosen in the order given
+        above.
     title : str | None
         Title for the figure. If None (default), equals to
         ``'Sensor positions (%%s)' %% ch_type``.
@@ -936,12 +937,10 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
     to_sphere : bool
         Whether to project the 3d locations to a sphere. When False, the
         sensor array appears similar as to looking downwards straight above the
-        subject's head. Has no effect when kind='3d'. Defaults to True.
+        subject's head. Has no effect when ``kind='3d'``. Defaults to True.
 
         .. versionadded:: 0.14.0
-    axes : instance of Axes | instance of Axes3D | None
-        Axes to draw the sensors to. If ``kind='3d'``, axes must be an instance
-        of Axes3D. If None (default), a new axes will be created.
+    %(axes_montage)s
 
         .. versionadded:: 0.13.0
     block : bool
@@ -953,10 +952,10 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
         Show figure if True. Defaults to True.
     %(sphere_topomap_auto)s
     pointsize : float | None
-        The size of the points. If None (default), will bet set to 75 if
-        ``kind='3d'``, or 25 otherwise.
+        The size of the points. If None (default), will bet set to ``75`` if
+        ``kind='3d'``, or ``25`` otherwise.
     linewidth : float
-        The width of the outline. If 0, the outline will not be drawn.
+        The width of the outline. If ``0``, the outline will not be drawn.
     %(verbose)s
 
     Returns
@@ -980,8 +979,25 @@ def plot_sensors(info, kind='topomap', ch_type=None, title=None,
     """
     from .evoked import _rgb
     _check_option('kind', kind, ['topomap', '3d', 'select'])
-    if not isinstance(info, Info):
-        raise TypeError(f'info must be an instance of Info not {type(info)}')
+    if axes is not None:
+        from matplotlib.axes import Axes
+        from mpl_toolkits.mplot3d.axes3d import Axes3D
+
+        if kind == "3d":
+            _validate_type(axes, Axes3D, "axes", extra="when 'kind' is '3d'")
+        elif kind in ("topomap", "select"):
+            _validate_type(
+                axes,
+                Axes,
+                "axes",
+                extra="when 'kind' is 'topomap' or 'select'"
+            )
+            if isinstance(axes, Axes3D):
+                raise TypeError(
+                    "axes must be an instance of Axes when 'kind' is "
+                    f"'topomap' or 'select', got {type(axes)} instead."
+                )
+    _validate_type(info, Info, "info")
     ch_indices = channel_indices_by_type(info)
     allowed_types = _DATA_CH_TYPES_SPLIT
     if ch_type is None: