Repository: mne-tools/mne-python

Issue:
`mne.time_frequency.psd_welch` does not support negative numbers for `n_jobs` parameter.
Since it should be a rather unambiguous phenomenon, so please allow me to describe the bug without your recommended format.

Negative numbers like -1 cannot be used for `n_jobs` in the function `psd_welch` under the module `mne.time_frequency`, which fails to enable users to adaptively activate all cores when calling this function, unlike those functions (such as `plot_psd`) that allows -1 for `n_jobs` in MNE.


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

This probably means we need a `n_jobs = check_n_jobs(n_jobs)` in the code. Would you be up for making a PR to fix this?
@larsoner I have carefully checked your code and I think I have found the real reason. The code in question is `x_splits = np.array_split(x, n_jobs)`, which lies in line **156**, the function `psd_array_welch`, in the file `mne\time_frequency\psd.py`.
If I understand your code right, the negative n_jobs will be converted by the function `check_n_jobs`, which will be invoked by the function `parallel_func`. The problem is you calling `parallel_func` in line **162**, after the array split action, which causes a `ValueError` in calling `np.array_split`.
I propose a straightforward fix that simply swaps the line 156 under the line 162, because the result `x_splits` won't be used until line 165. Please confirm if my solution is right.
yes go for it

thx

>

When I am fixing this bug, I can't help but check if there is any similar problem in `psd_array_multitaper` and turns out there is. However this time a simple swap cannot work.

The codes in question are extracted as below, from line 435 to line 448 of `mne/time_frequency/multitaper.py`:
```
    n_chunk = max(50000000 // (len(freq_mask) * len(eigvals) * 16), n_jobs)
    offsets = np.concatenate((np.arange(0, x.shape[0], n_chunk), [x.shape[0]]))
    for start, stop in zip(offsets[:-1], offsets[1:]):
        x_mt = _mt_spectra(x[start:stop], dpss, sfreq)[0]
        if not adaptive:
            weights = np.sqrt(eigvals)[np.newaxis, :, np.newaxis]
            psd[start:stop] = _psd_from_mt(x_mt[:, :, freq_mask], weights)
        else:
            n_splits = min(stop - start, n_jobs)
            parallel, my_psd_from_mt_adaptive, n_jobs = \
                parallel_func(_psd_from_mt_adaptive, n_splits)
            out = parallel(my_psd_from_mt_adaptive(x, eigvals, freq_mask)
                           for x in np.array_split(x_mt, n_splits))
            psd[start:stop] = np.concatenate(out)
```
Notice that `n_jobs` is converted at `parallel, my_psd_from_mt_adaptive, n_jobs = parallel_func(_psd_from_mt_adaptive, n_splits)` which is inside an `else` and a `for`, but the very first line `n_chunk = max(50000000 // (len(freq_mask) * len(eigvals) * 16), n_jobs)`, which is outside the loop, uses `n_jobs` immediately, which invalidates a swap as the solution.
From my coarse level of understanding, maybe a direct invocation of `check_n_jobs` in place is the only feasible solution to this. I do not know if it suits your project and I hope there is a more elegant way to solve it.
> maybe a direct invocation of check_n_jobs in place is the only feasible solution to this.

I think calling `check_n_jobs` is the standard way we deal with `n_jobs` arguments, so yes I would fix it by adding a `n_jobs = check_n_jobs(n_jobs)` in whatever methods currently do not have it.

Resolving patch:
diff --git a/mne/time_frequency/multitaper.py b/mne/time_frequency/multitaper.py
--- a/mne/time_frequency/multitaper.py
+++ b/mne/time_frequency/multitaper.py
@@ -7,7 +7,7 @@
 import numpy as np
 
 from ..fixes import _import_fft
-from ..parallel import parallel_func
+from ..parallel import parallel_func, check_n_jobs
 from ..utils import sum_squared, warn, verbose, logger, _check_option
 
 
@@ -432,6 +432,7 @@ def psd_array_multitaper(x, sfreq, fmin=0, fmax=np.inf, bandwidth=None,
 
     psd = np.zeros((x.shape[0], freq_mask.sum()))
     # Let's go in up to 50 MB chunks of signals to save memory
+    n_jobs = check_n_jobs(n_jobs)
     n_chunk = max(50000000 // (len(freq_mask) * len(eigvals) * 16), n_jobs)
     offsets = np.concatenate((np.arange(0, x.shape[0], n_chunk), [x.shape[0]]))
     for start, stop in zip(offsets[:-1], offsets[1:]):
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
@@ -5,7 +5,7 @@
 from functools import partial
 import numpy as np
 
-from ..parallel import parallel_func
+from ..parallel import parallel_func, check_n_jobs
 from ..io.pick import _picks_to_idx
 from ..utils import logger, verbose, _time_mask, _check_option
 from .multitaper import psd_array_multitaper
@@ -153,6 +153,7 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     freqs = freqs[freq_sl]
 
     # Parallelize across first N-1 dimensions
+    n_jobs = check_n_jobs(n_jobs)
     x_splits = np.array_split(x, n_jobs)
     logger.debug(
         f'Spectogram using {n_fft}-point FFT on {n_per_seg} samples with '


Test patch: