Repository: mne-tools/mne-python

Issue:
ICA plot_components on provided axes fails
### Description of the problem

```
from matplotlib import pyplot as plt
from mne.datasets import sample
from mne.io import read_raw_fif
from mne.preprocessing import ICA


# create a toy dataset
fname = sample.data_path() / "MEG" / "sample" / "sample_audvis_raw.fif"
raw = read_raw_fif(fname, preload=False)
raw.pick_types(eeg=True)
raw.crop(0, 10)
raw.load_data()
raw.filter(1.0, 40.0)
ica = ICA(n_components=5, method="picard")
ica.fit(raw)

# create figure
f, ax = plt.subplots(2, 2)
ica.plot_components(inst=raw, picks=[1, 2, 3, 4], axes=ax)
```

The code snippet above raises:

```
AttributeError: 'numpy.ndarray' object has no attribute 'set_title'
```

If you flatten the provided `ax`, it fails with:

```
UnboundLocalError: local variable 'fig' referenced before assignment
```

Working on it atm.


Discussion:

Resolving patch:
diff --git a/mne/viz/topomap.py b/mne/viz/topomap.py
--- a/mne/viz/topomap.py
+++ b/mne/viz/topomap.py
@@ -1136,7 +1136,7 @@ def plot_ica_components(
         Defaults to True, which plots one standard deviation above/below.
         If set to float allows to control how many standard deviations are
         plotted. For example 2.5 will plot 2.5 standard deviation above/below.
-    reject : 'auto' | dict | None
+    reject : ``'auto'`` | dict | None
         Allows to specify rejection parameters used to drop epochs
         (or segments if continuous signal is passed as inst).
         If None, no rejection is applied. The default is 'auto',
@@ -1167,7 +1167,12 @@ def plot_ica_components(
         .. versionadded:: 1.3
     %(colorbar_topomap)s
     %(cbar_fmt_topomap)s
-    %(axes_evoked_plot_topomap)s
+    axes : Axes | array of Axes | None
+        The subplot(s) to plot to. Either a single Axes or an iterable of Axes
+        if more than one subplot is needed. The number of subplots must match
+        the number of selected components. If None, new figures will be created
+        with the number of subplots per figure controlled by ``nrows`` and
+        ``ncols``.
     title : str | None
         The title of the generated figure. If ``None`` (default) and
         ``axes=None``, a default title of "ICA Components" will be used.
@@ -1199,6 +1204,8 @@ def plot_ica_components(
     topomap (this option is only available when the ``inst`` argument is
     supplied).
     """  # noqa E501
+    from matplotlib.pyplot import Axes
+
     from ..io import BaseRaw
     from ..epochs import BaseEpochs
 
@@ -1226,7 +1233,12 @@ def plot_ica_components(
         figs = []
         cut_points = range(max_subplots, n_components, max_subplots)
         pick_groups = np.split(range(n_components), cut_points)
-        for _picks in pick_groups:
+        for k, _picks in enumerate(pick_groups):
+            _axes = axes.flatten() if isinstance(axes, np.ndarray) else axes
+            try:  # either an iterable, 1D numpy array or others
+                _axes = _axes[k * max_subplots: (k + 1) * max_subplots]
+            except TypeError:  # None or Axes
+                _axes = axes
             fig = plot_ica_components(
                 ica, picks=_picks, ch_type=ch_type, inst=inst,
                 plot_std=plot_std, reject=reject, sensors=sensors,
@@ -1234,7 +1246,7 @@ def plot_ica_components(
                 sphere=sphere, image_interp=image_interp,
                 extrapolate=extrapolate, border=border, res=res, size=size,
                 cmap=cmap, vlim=vlim, cnorm=cnorm, colorbar=colorbar,
-                cbar_fmt=cbar_fmt, axes=axes, title=title, nrows=nrows,
+                cbar_fmt=cbar_fmt, axes=_axes, title=title, nrows=nrows,
                 ncols=ncols, show=show, image_args=image_args,
                 psd_args=psd_args, verbose=verbose)
             figs.append(fig)
@@ -1260,6 +1272,10 @@ def plot_ica_components(
     if not user_passed_axes:
         fig, axes, _, _ = _prepare_trellis(len(data), ncols=ncols, nrows=nrows)
         fig.suptitle(title)
+    else:
+        axes = axes.flatten() if isinstance(axes, np.ndarray) else axes
+        axes = [axes] if isinstance(axes, Axes) else axes
+        fig = axes[0].get_figure()
 
     subplot_titles = list()
     for ii, data_, ax in zip(picks, data, axes):


Test patch:
diff --git a/mne/viz/tests/test_ica.py b/mne/viz/tests/test_ica.py
--- a/mne/viz/tests/test_ica.py
+++ b/mne/viz/tests/test_ica.py
@@ -6,10 +6,10 @@
 import sys
 from pathlib import Path
 
+import matplotlib.pyplot as plt
 import numpy as np
-from numpy.testing import assert_equal, assert_array_equal
 import pytest
-import matplotlib.pyplot as plt
+from numpy.testing import assert_equal, assert_array_equal
 
 from mne import (read_events, Epochs, read_cov, pick_types, Annotations,
                  make_fixed_length_events)
@@ -59,9 +59,11 @@ def test_plot_ica_components():
     res = 8
     fast_test = {"res": res, "contours": 0, "sensors": False}
     raw = _get_raw()
-    ica = ICA(noise_cov=read_cov(cov_fname), n_components=2)
+    ica = ICA(noise_cov=read_cov(cov_fname), n_components=8)
     ica_picks = _get_picks(raw)
-    with pytest.warns(RuntimeWarning, match='projection'):
+    with pytest.warns(
+        RuntimeWarning, match="(projection)|(unstable mixing matrix)"
+    ):
         ica.fit(raw, picks=ica_picks)
 
     for components in [0, [0], [0, 1], [0, 1] * 2, None]:
@@ -106,6 +108,18 @@ def test_plot_ica_components():
     title = topomap_ax.get_title()
     assert (lbl.split(' ')[0] == title.split(' ')[0])
 
+    # test provided axes
+    _, ax = plt.subplots(1, 1)
+    ica.plot_components(axes=ax, picks=0, **fast_test)
+    _, ax = plt.subplots(2, 1)
+    ica.plot_components(axes=ax, picks=[0, 1], **fast_test)
+    _, ax = plt.subplots(2, 2)
+    ica.plot_components(axes=ax, picks=[0, 1, 2, 3], **fast_test)
+    _, ax = plt.subplots(3, 2)
+    ica.plot_components(
+        axes=ax, picks=[0, 1, 2, 3, 4, 5], nrows=2, ncols=2, **fast_test
+    )
+
     ica.info = None
     with pytest.raises(RuntimeError, match='fit the ICA'):
         ica.plot_components(1, ch_type='mag')