[BUG] Fix computation of welch's PSD when the number of jobs exceeds the number of channels

doc/changes/latest.inc

@@ -41,6 +41,7 @@ Bugs
 - Fix bug in the ``.compute_psd()`` methods where the number of unaggregated Welch segments was wrongly computed for some inputs, leading to an assertion error when computing the PSD (:gh:`11248` by `Daniel McCloy`_)
 - Fix bug in the :func:`~mne.viz.plot_evoked_topo` and :meth:`~mne.Evoked.plot_topo`, where legend colors where shown incorrectly on newer matplotlib versions (:gh:`11258` by `Erkka Heinila`_)
 - Fix bug where EEGLAB channel positions were read as meters, while they are commonly in millimeters, leading to head outlies of the size of one channel when plotting topomaps. Now ``montage_units`` argument has been added to :func:`~mne.io.read_raw_eeglab` and :func:`~mne.read_epochs_eeglab` to control in what units EEGLAB channel positions are read. The default is millimeters, ``'mm'`` (:gh:`11283` by `Mikołaj Magnuski`_)
+- Fix bug where computing PSD with welch's method with more jobs than channels would fail (:gh:`11298` by `Mathieu Scheltienne`_)
 - Fix channel selection edge-cases in `~mne.preprocessing.ICA.find_bads_muscle` (:gh:`11300` by `Mathieu Scheltienne`_)
 
 API changes

mne/time_frequency/psd.py

@@ -180,7 +180,7 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     parallel, my_spect_func, n_jobs = parallel_func(_spect_func, n_jobs=n_jobs)
     func = partial(spectrogram, noverlap=n_overlap, nperseg=n_per_seg,
                    nfft=n_fft, fs=sfreq, window=window)
-    x_splits = np.array_split(x, n_jobs)
+    x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
     f_spect = parallel(my_spect_func(d, func=func, freq_sl=freq_sl,
                                      average=average)
                        for d in x_splits)

mne/time_frequency/tests/test_psd.py

@@ -192,3 +192,10 @@ def test_compares_psd():
     assert (np.sum(freqs_scipy < 0) == 0)
     assert (np.sum(psds_mne < 0) == 0)
     assert (np.sum(psds_scipy < 0) == 0)
+
+
+def test_psd_array_welch_n_jobs():
+    """Test that n_jobs works even with more jobs than channels."""
+    data = np.empty((1, 2048))
+    psd_array_welch(data, 1024, n_jobs=1)
+    psd_array_welch(data, 1024, n_jobs=2)