Repository: mne-tools/mne-python

Issue:
compute_psd: ValueError when n_jobs=-1 and only 1 channel
### Description of the problem

when [`compute_psd`](https://mne.tools/stable/generated/mne.io.Raw.html) is passed `n_jobs=-1` and the instance contains only _one_ channel, the method fails.

### Steps to reproduce

```Python
import numpy as np
import mne
data = np.random.randn(1, 4096) 
info = mne.create_info(ch_names=["EEG 01"], ch_types="eeg")
raw = mne.io.RawArray(data, info)
spectrum = raw.compute_psd(method="welch", n_jobs=-1)
```
```


### Link to data

_No response_

### Expected results

no error

### Actual results

```
ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)
```

### Additional information

Platform:         Linux-5.19.16-76051916-generic-x86_64-with-glibc2.35
Python:           3.10.6 (main, Aug 10 2022, 11:40:04) [GCC 11.3.0]
Executable:       /usr/bin/python3
CPU:              x86_64: 8 cores
Memory:           15.3 GB

mne:              1.2.1
numpy:            1.23.4 {OpenBLAS 0.3.20 with 8 threads}
scipy:            1.9.3
matplotlib:       3.6.0 {backend=module://matplotlib_inline.backend_inline}

sklearn:          1.1.3
numba:            Not found
nibabel:          Not found
nilearn:          Not found
dipy:             Not found
openmeeg:         Not found
cupy:             Not found
pandas:           1.5.1
pyvista:          Not found
pyvistaqt:        Not found
ipyvtklink:       Not found
vtk:              Not found
qtpy:             2.2.1 {PyQt5=5.15.3}
ipympl:           Not found
pyqtgraph:        Not found
pooch:            v1.6.0

mne_bids:         Not found
mne_nirs:         Not found
mne_features:     Not found
mne_qt_browser:   Not found
mne_connectivity: Not found
mne_icalabel:     Not found


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

Hi @sam-s, thanks for opening the issue!
Could you paste the whole error traceback that you get?
@mmagnuski Here you go:

```
Effective window size : 0.250 (s)
[Parallel(n_jobs=2)]: Using backend LokyBackend with 2 concurrent workers.
[Parallel(n_jobs=2)]: Done   2 out of   2 | elapsed:    0.9s remaining:    0.0s
[Parallel(n_jobs=2)]: Done   2 out of   2 | elapsed:    0.9s finished
Traceback (most recent call last):

  File "C:\Users\Mathieu\AppData\Local\Temp\ipykernel_22068\4243930617.py", line 8, in <cell line: 8>
    spectrum = raw.compute_psd(method="welch", n_jobs=2)

  File "<decorator-gen-237>", line 12, in compute_psd

  File "c:\users\mathieu\documents\git\mne-tools\mne-python\mne\io\base.py", line 1841, in compute_psd
    return Spectrum(

  File "c:\users\mathieu\documents\git\mne-tools\mne-python\mne\time_frequency\spectrum.py", line 942, in __init__
    self._compute_spectra(data, fmin, fmax, n_jobs, method_kw, verbose)

  File "c:\users\mathieu\documents\git\mne-tools\mne-python\mne\time_frequency\spectrum.py", line 380, in _compute_spectra
    result = self._psd_func(

  File "<decorator-gen-198>", line 12, in psd_array_welch

  File "c:\users\mathieu\documents\git\mne-tools\mne-python\mne\time_frequency\psd.py", line 187, in psd_array_welch
    psds = np.concatenate(f_spect, axis=0)

  File "<__array_function__ internals>", line 180, in concatenate

ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)
```

With `n_jobs=2`.
Nevermind, I can reproduce the issue:
```python-traceback
  File "C:\Users\mmagn\AppData\Local\Temp\ipykernel_2108\1312025537.py", line 12, in <module>
    spectrum = raw.compute_psd(method="welch", n_jobs=-1)

  File "<decorator-gen-240>", line 12, in compute_psd

  File "c:\src\mne-python\mne\io\base.py", line 1845, in compute_psd
    return Spectrum(

  File "c:\src\mne-python\mne\time_frequency\spectrum.py", line 942, in __init__
    self._compute_spectra(data, fmin, fmax, n_jobs, method_kw, verbose)

  File "c:\src\mne-python\mne\time_frequency\spectrum.py", line 380, in _compute_spectra
    result = self._psd_func(

  File "<decorator-gen-201>", line 12, in psd_array_welch

  File "c:\src\mne-python\mne\time_frequency\psd.py", line 187, in psd_array_welch
    psds = np.concatenate(f_spect, axis=0)

  File "<__array_function__ internals>", line 180, in concatenate

ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)
```
Oh, @mscheltienne we were simultaneous! :)
Looks like an easy fix here on `x_splits`:

```
    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)
    f_spect = parallel(my_spect_func(d, func=func, freq_sl=freq_sl,
                                     average=average)
                       for d in x_splits)
    psds = np.concatenate(f_spect, axis=0)
```

Resolving patch:
diff --git a/mne/time_frequency/psd.py b/mne/time_frequency/psd.py
--- a/mne/time_frequency/psd.py
+++ b/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)


Test patch:
diff --git a/mne/time_frequency/tests/test_psd.py b/mne/time_frequency/tests/test_psd.py
--- a/mne/time_frequency/tests/test_psd.py
+++ b/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)