Repository: mne-tools/mne-bids-pipeline

Issue:
Add support for AutoReject local
With Tomas D'Amelio we implemented the AutoReject local option in ptp_reject to clean the DEAP data. Results are pretty nice! Some work is necessary to make that code fit into master after all we merged lately but we should do it.


Discussion:

Resolving patch:
diff --git a/mne_bids_pipeline/_config.py b/mne_bids_pipeline/_config.py
--- a/mne_bids_pipeline/_config.py
+++ b/mne_bids_pipeline/_config.py
@@ -1355,54 +1355,51 @@
 # Rejection based on peak-to-peak amplitude
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-reject: Optional[Union[Dict[str, float], Literal["autoreject_global"]]] = None
+reject: Optional[
+    Union[Dict[str, float], Literal["autoreject_global", "autoreject_local"]]
+] = None
 """
 Peak-to-peak amplitude limits to mark epochs as bad. This allows you to remove
 epochs with strong transient artifacts.
 
-If `None` (default), do not apply artifact rejection. If a dictionary,
-manually specify rejection thresholds (see examples).  If
-`'autoreject_global'`, use [`autoreject`](https://autoreject.github.io) to find
-suitable "global" rejection thresholds for each channel type, i.e. `autoreject`
-will generate a dictionary with (hopefully!) optimal thresholds for each
-channel type.
+!!! info
+      The rejection is performed **after** SSP or ICA, if any of those methods
+      is used. To reject epochs **before** fitting ICA, see the
+      [`ica_reject`][mne_bids_pipeline._config.ica_reject] setting.
 
+If `None` (default), do not apply artifact rejection.
+
+If a dictionary, manually specify rejection thresholds (see examples). 
 The thresholds provided here must be at least as stringent as those in
 [`ica_reject`][mne_bids_pipeline._config.ica_reject] if using ICA. In case of
 `'autoreject_global'`, thresholds for any channel that do not meet this
 requirement will be automatically replaced with those used in `ica_reject`.
 
-!!! info
-      The rejection is performed **after** SSP or ICA, if any of those methods
-      is used. To reject epochs **before** fitting ICA, see the
-      [`ica_reject`][mne_bids_pipeline._config.ica_reject] setting.
-
-If `None` (default), do not apply automated rejection. If a dictionary,
-manually specify rejection thresholds (see examples).  If `'auto'`, use
-[`autoreject`](https://autoreject.github.io) to find suitable "global"
-rejection thresholds for each channel type, i.e. `autoreject` will generate
-a dictionary with (hopefully!) optimal thresholds for each channel type. Note
-that using `autoreject` can be a time-consuming process.
+If `"autoreject_global"`, use [`autoreject`](https://autoreject.github.io) to find
+suitable "global" rejection thresholds for each channel type, i.e., `autoreject`
+will generate a dictionary with (hopefully!) optimal thresholds for each
+channel type.
 
-!!! info
-      `autoreject` basically offers two modes of operation: "global" and
-      "local". In "global" mode, it will try to estimate one rejection
-      threshold **per channel type.** In "local" mode, it will generate
-      thresholds **for each individual channel.** Currently, the BIDS Pipeline
-      only supports the "global" mode.
+If `"autoreject_local"`, use "local" `autoreject` to detect (and potentially repair) bad
+channels in each epoch. Use [`autoreject_n_interpolate`][mne_bids_pipeline._config.autoreject_n_interpolate]
+to control how many channels are allowed to be bad before an epoch gets dropped.
 
 ???+ example "Example"
     ```python
-    reject = {'grad': 4000e-13, 'mag': 4e-12, 'eog': 150e-6}
-    reject = {'eeg': 100e-6, 'eog': 250e-6}
+    reject = {"grad": 4000e-13, 'mag': 4e-12, 'eog': 150e-6}
+    reject = {"eeg": 100e-6, "eog": 250e-6}
     reject = None  # no rejection based on PTP amplitude
+    reject = "autoreject_global"  # find global (per channel type) PTP thresholds
+    reject = "autoreject_local"  # find local (per channel) thresholds and repair epochs
     ```
 """
 
 reject_tmin: Optional[float] = None
 """
 Start of the time window used to reject epochs. If `None`, the window will
-start with the first time point.
+start with the first time point. Has no effect if
+[`reject`][mne_bids_pipeline._config.reject] has been set to `"autoreject_local"`.
+
 ???+ example "Example"
     ```python
     reject_tmin = -0.1  # 100 ms before event onset.
@@ -1412,13 +1409,32 @@
 reject_tmax: Optional[float] = None
 """
 End of the time window used to reject epochs. If `None`, the window will end
-with the last time point.
+with the last time point. Has no effect if
+[`reject`][mne_bids_pipeline._config.reject] has been set to `"autoreject_local"`.
+
 ???+ example "Example"
     ```python
     reject_tmax = 0.3  # 300 ms after event onset.
     ```
 """
 
+autoreject_n_interpolate: FloatArrayLike = [4, 8, 16]
+"""
+The maximum number of bad channels in an epoch that `autoreject` local will try to
+interpolate. The optimal number among this list will be estimated using a
+cross-validation procedure; this means that the more elements are provided here, the
+longer the `autoreject` run will take. If the number of bad channels in an epoch
+exceeds this value, the channels won't be interpolated and the epoch will be dropped.
+
+!!! info
+    This setting only takes effect if [`reject`][mne_bids_pipeline._config.reject] has
+    been set to `"autoreject_local"`.
+
+!!! info
+    Channels marked as globally bad in the BIDS dataset (in `*_channels.tsv)`) will not
+    be considered (i.e., will remain marked as bad and not analyzed by autoreject).
+"""
+
 ###############################################################################
 # DECODING
 # --------
diff --git a/mne_bids_pipeline/_config_import.py b/mne_bids_pipeline/_config_import.py
--- a/mne_bids_pipeline/_config_import.py
+++ b/mne_bids_pipeline/_config_import.py
@@ -259,7 +259,7 @@ def _check_config(config: SimpleNamespace, config_path: Optional[PathLike]) -> N
         if (
             ica_reject is not None
             and reject is not None
-            and reject != "autoreject_global"
+            and reject not in ["autoreject_global", "autoreject_local"]
         ):
             for ch_type in reject:
                 if ch_type in ica_reject and reject[ch_type] > ica_reject[ch_type]:
@@ -354,6 +354,7 @@ def _default_factory(key, val):
         {"custom": (8, 24.0, 40)},  # decoding_csp_freqs
         {"suffix": "ave"},  # source_info_path_update
         ["evoked"],  # inverse_targets
+        [4, 8, 16],  # autoreject_n_interpolate
     ]
     for typ in (dict, list):
         if isinstance(val, typ):
diff --git a/mne_bids_pipeline/steps/preprocessing/_08_ptp_reject.py b/mne_bids_pipeline/steps/preprocessing/_08_ptp_reject.py
--- a/mne_bids_pipeline/steps/preprocessing/_08_ptp_reject.py
+++ b/mne_bids_pipeline/steps/preprocessing/_08_ptp_reject.py
@@ -11,6 +11,9 @@
 from types import SimpleNamespace
 from typing import Optional
 
+import numpy as np
+import autoreject
+
 import mne
 from mne_bids import BIDSPath
 
@@ -79,46 +82,72 @@ def drop_ptp(
 
     # Get rejection parameters and drop bad epochs
     epochs = mne.read_epochs(in_files.pop("epochs"), preload=True)
-    reject = _get_reject(
-        subject=subject,
-        session=session,
-        reject=cfg.reject,
-        ch_types=cfg.ch_types,
-        param="reject",
-        epochs=epochs,
-    )
-    if cfg.spatial_filter == "ica":
-        ica_reject = _get_reject(
+
+    if cfg.reject == "autoreject_local":
+        msg = "Using autoreject to find and repair bad epochs"
+        logger.info(**gen_log_kwargs(message=msg))
+
+        ar = autoreject.AutoReject(
+            n_interpolate=np.array(cfg.autoreject_n_interpolate),
+            random_state=cfg.random_state,
+            n_jobs=exec_params.n_jobs,
+            verbose=False,
+        )
+        n_epochs_before_reject = len(epochs)
+        epochs, reject_log = ar.fit_transform(epochs, return_log=True)
+        n_epochs_after_reject = len(epochs)
+        assert (
+            n_epochs_before_reject - n_epochs_after_reject
+            == reject_log.bad_epochs.sum()
+        )
+
+        msg = (
+            f"autoreject marked {reject_log.bad_epochs.sum()} epochs as bad "
+            f"(cross-validated n_interpolate limit: {ar.n_interpolate_})"
+        )
+        logger.info(**gen_log_kwargs(message=msg))
+    else:
+        reject = _get_reject(
             subject=subject,
             session=session,
-            reject=cfg.ica_reject,
+            reject=cfg.reject,
             ch_types=cfg.ch_types,
-            param="ica_reject",
+            param="reject",
+            epochs=epochs,
         )
-    else:
-        ica_reject = None
-
-    if ica_reject is not None:
-        for ch_type, threshold in ica_reject.items():
-            if ch_type in reject and threshold < reject[ch_type]:
-                # This can only ever happen in case of
-                # reject = 'autoreject_global'
-                msg = (
-                    f"Adjusting PTP rejection threshold proposed by "
-                    f"autoreject, as it is greater than ica_reject: "
-                    f"{ch_type}: {reject[ch_type]} -> {threshold}"
-                )
-                logger.info(**gen_log_kwargs(message=msg))
-                reject[ch_type] = threshold
-
-    msg = f"Using PTP rejection thresholds: {reject}"
-    logger.info(**gen_log_kwargs(message=msg))
 
-    n_epochs_before_reject = len(epochs)
-    epochs.reject_tmin = cfg.reject_tmin
-    epochs.reject_tmax = cfg.reject_tmax
-    epochs.drop_bad(reject=reject)
-    n_epochs_after_reject = len(epochs)
+        if cfg.spatial_filter == "ica":
+            ica_reject = _get_reject(
+                subject=subject,
+                session=session,
+                reject=cfg.ica_reject,
+                ch_types=cfg.ch_types,
+                param="ica_reject",
+            )
+        else:
+            ica_reject = None
+
+        if ica_reject is not None:
+            for ch_type, threshold in ica_reject.items():
+                if ch_type in reject and threshold < reject[ch_type]:
+                    # This can only ever happen in case of
+                    # reject = 'autoreject_global'
+                    msg = (
+                        f"Adjusting PTP rejection threshold proposed by "
+                        f"autoreject, as it is greater than ica_reject: "
+                        f"{ch_type}: {reject[ch_type]} -> {threshold}"
+                    )
+                    logger.info(**gen_log_kwargs(message=msg))
+                    reject[ch_type] = threshold
+
+        msg = f"Using PTP rejection thresholds: {reject}"
+        logger.info(**gen_log_kwargs(message=msg))
+
+        n_epochs_before_reject = len(epochs)
+        epochs.reject_tmin = cfg.reject_tmin
+        epochs.reject_tmax = cfg.reject_tmax
+        epochs.drop_bad(reject=reject)
+        n_epochs_after_reject = len(epochs)
 
     if 0 < n_epochs_after_reject < 0.5 * n_epochs_before_reject:
         msg = (
@@ -155,6 +184,24 @@ def drop_ptp(
     with _open_report(
         cfg=cfg, exec_params=exec_params, subject=subject, session=session
     ) as report:
+        if cfg.reject == "autoreject_local":
+            caption = (
+                f"Autoreject was run to produce cleaner epochs. "
+                f"{reject_log.bad_epochs.sum()} epochs were rejected because more than "
+                f"{ar.n_interpolate_} channels were bad (cross-validated n_interpolate "
+                f"limit; excluding globally bad and non-data channels, shown in white)."
+            )
+            report.add_figure(
+                fig=reject_log.plot(
+                    orientation="horizontal", aspect="auto", show=False
+                ),
+                title="Epochs: Autoreject cleaning",
+                caption=caption,
+                tags=("epochs", "autoreject"),
+                replace=True,
+            )
+            del caption
+
         report.add_epochs(
             epochs=epochs,
             title="Epochs: after cleaning",
@@ -176,6 +223,8 @@ def get_config(
         spatial_filter=config.spatial_filter,
         ica_reject=config.ica_reject,
         reject=config.reject,
+        autoreject_n_interpolate=config.autoreject_n_interpolate,
+        random_state=config.random_state,
         ch_types=config.ch_types,
         _epochs_split_size=config._epochs_split_size,
         **_bids_kwargs(config=config),


Test patch:
diff --git a/mne_bids_pipeline/tests/configs/config_ERP_CORE.py b/mne_bids_pipeline/tests/configs/config_ERP_CORE.py
--- a/mne_bids_pipeline/tests/configs/config_ERP_CORE.py
+++ b/mne_bids_pipeline/tests/configs/config_ERP_CORE.py
@@ -71,13 +71,18 @@
 t_break_annot_start_after_previous_event = 3.0
 t_break_annot_stop_before_next_event = 1.5
 
-ica_reject = dict(eeg=350e-6, eog=500e-6)
-reject = "autoreject_global"
+if task == "N400":  # test autoreject local without ICA
+    spatial_filter = None
+    reject = "autoreject_local"
+    autoreject_n_interpolate = [2, 4]
+else:
+    ica_reject = dict(eeg=350e-6, eog=500e-6)
+    reject = "autoreject_global"
 
-spatial_filter = "ica"
-ica_max_iterations = 1000
-ica_eog_threshold = 2
-ica_decim = 2  # speed up ICA fitting
+    spatial_filter = "ica"
+    ica_max_iterations = 1000
+    ica_eog_threshold = 2
+    ica_decim = 2  # speed up ICA fitting
 
 run_source_estimation = False
 
@@ -104,7 +109,6 @@
     }
 
     eeg_reference = ["P9", "P10"]
-    ica_n_components = 30 - len(eeg_reference)
     epochs_tmin = -0.2
     epochs_tmax = 0.8
     epochs_metadata_tmin = 0