[MRG] Adding channel_wise argument to Raw.apply_function (#5875)

* Adding `channel_wise` argument to `Raw.apply_function`

* precising docstring and modifying test syntax

* DOC: rst and order fixes

* DOC: rst fixes

* updating whats_new.rst

* FIX: Line lengths

Author: hubertjb

doc/whats_new.rst

@@ -53,6 +53,8 @@ Changelog
 
 - Add ``extrapolate`` argument to :func:`mne.viz.plot_topomap` for better control of extrapolation points placement by `Mikołaj Magnuski`_
 
+- Add ``channel_wise`` argument to :func:`mne.io.Raw.apply_function` to allow applying a function on multiple channels at once by `Hubert Banville`_
+
 Bug
 ~~~
 

mne/io/base.py

@@ -1033,8 +1033,8 @@ def get_data(self, picks=None, start=0, stop=None,
         return data
 
     @verbose
-    def apply_function(self, fun, picks=None, dtype=None,
-                       n_jobs=1, *args, **kwargs):
+    def apply_function(self, fun, picks=None, dtype=None, n_jobs=1,
+                       channel_wise=True, *args, **kwargs):
         """Apply a function to a subset of channels.
 
         The function "fun" is applied to the channels defined in "picks". The
@@ -1059,15 +1059,23 @@ def apply_function(self, fun, picks=None, dtype=None,
         fun : function
             A function to be applied to the channels. The first argument of
             fun has to be a timeseries (numpy.ndarray). The function must
-            return an numpy.ndarray with the same size as the input.
+            operate on an array of shape ``(n_times,)`` if
+            ``channel_wise=True`` and ``(len(picks), n_times)`` otherwise.
+            The function must return an ndarray shaped like its input.
         picks : array-like of int (default: None)
             Indices of channels to apply the function to. If None, all data
             channels are used.
         dtype : numpy.dtype (default: None)
             Data type to use for raw data after applying the function. If None
             the data type is not modified.
         n_jobs: int (default: 1)
-            Number of jobs to run in parallel.
+            Number of jobs to run in parallel. Ignored if `channel_wise` is
+            False.
+        channel_wise: bool (default: True)
+            Whether to apply the function to each channel individually. If
+            False, the function will be applied to all channels at once.
+
+            .. versionadded:: 0.18
         *args :
             Additional positional arguments to pass to fun (first pos. argument
             of fun is the timeseries of a channel).
@@ -1094,18 +1102,23 @@ def apply_function(self, fun, picks=None, dtype=None,
         if dtype is not None and dtype != self._data.dtype:
             self._data = self._data.astype(dtype)
 
-        if n_jobs == 1:
-            # modify data inplace to save memory
-            for idx in picks:
-                self._data[idx, :] = _check_fun(fun, data_in[idx, :],
-                                                *args, **kwargs)
+        if channel_wise:
+            if n_jobs == 1:
+                # modify data inplace to save memory
+                for idx in picks:
+                    self._data[idx, :] = _check_fun(fun, data_in[idx, :],
+                                                    *args, **kwargs)
+            else:
+                # use parallel function
+                parallel, p_fun, _ = parallel_func(_check_fun, n_jobs)
+                data_picks_new = parallel(
+                    p_fun(fun, data_in[p], *args, **kwargs) for p in picks)
+                for pp, p in enumerate(picks):
+                    self._data[p, :] = data_picks_new[pp]
         else:
-            # use parallel function
-            parallel, p_fun, _ = parallel_func(_check_fun, n_jobs)
-            data_picks_new = parallel(p_fun(fun, data_in[p], *args, **kwargs)
-                                      for p in picks)
-            for pp, p in enumerate(picks):
-                self._data[p, :] = data_picks_new[pp]
+            self._data[picks, :] = _check_fun(
+                fun, data_in[picks, :], *args, **kwargs)
+
         return self
 
     @verbose

mne/io/tests/test_apply_function.py

@@ -20,6 +20,11 @@ def bad_2(x):
     return x[:-1]  # bad shape
 
 
+def bad_3(x):
+    """Fail."""
+    return x[0, :]
+
+
 def printer(x):
     """Print."""
     logger.info('exec')
@@ -35,10 +40,21 @@ def test_apply_function_verbose():
     raw = RawArray(np.zeros((n_chan, n_times)),
                    create_info(ch_names, 1., 'mag'))
     # test return types in both code paths (parallel / 1 job)
-    pytest.raises(TypeError, raw.apply_function, bad_1)
-    pytest.raises(ValueError, raw.apply_function, bad_2)
-    pytest.raises(TypeError, raw.apply_function, bad_1, n_jobs=2)
-    pytest.raises(ValueError, raw.apply_function, bad_2, n_jobs=2)
+    with pytest.raises(TypeError, match='Return value must be an ndarray'):
+        raw.apply_function(bad_1)
+    with pytest.raises(ValueError, match='Return data must have shape'):
+        raw.apply_function(bad_2)
+    with pytest.raises(TypeError, match='Return value must be an ndarray'):
+        raw.apply_function(bad_1, n_jobs=2)
+    with pytest.raises(ValueError, match='Return data must have shape'):
+        raw.apply_function(bad_2, n_jobs=2)
+
+    # test return type when `channel_wise=False`
+    raw.apply_function(printer, channel_wise=False)
+    with pytest.raises(TypeError, match='Return value must be an ndarray'):
+        raw.apply_function(bad_1, channel_wise=False)
+    with pytest.raises(ValueError, match='Return data must have shape'):
+        raw.apply_function(bad_3, channel_wise=False)
 
     # check our arguments
     with catch_logging() as sio: