compute_psd: ValueError when n_jobs=-1 and only 1 channel

[sam-s]

Description of the problem

when compute_psd is passed n_jobs=-1 and the instance contains only one channel, the method fails.

Steps to reproduce

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

[welcome]

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

[mmagnuski]

Hi @sam-s, thanks for opening the issue!
Could you paste the whole error traceback that you get?

[mscheltienne]

@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.

[mmagnuski]

Nevermind, I can reproduce the issue:

  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)

[mmagnuski]

Oh, @mscheltienne we were simultaneous! :)

[mscheltienne]

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)