implement channel specific epoch rejection

mne/channels/channels.py

@@ -815,7 +815,7 @@ def interpolate_bads(
         exclude=(),
         verbose=None,
     ):
-        """Interpolate bad MEG and EEG channels.
+        """Interpolate bad MEG, EEG and fNIRS channels.
 
         Operates in place.
 
@@ -858,6 +858,7 @@ def interpolate_bads(
         exclude : list | tuple
             The channels to exclude from interpolation. If excluded a bad
             channel will stay in bads.
+
         %(verbose)s
 
         Returns

mne/epochs.py

@@ -693,6 +693,37 @@ def __init__(
         self._check_consistency()
         self.set_annotations(annotations, on_missing="ignore")
 
+    def set_bad_epochs_to_NaN(self, bad_epochs_indices: list = None):
+        """
+        Define bad epochs based on indices list and set to NaN.
+
+        Parameters
+        ----------
+        self : instance of Epochs
+        bad_epochs_indices : list of arrays
+            List of arrays with indices of bad epochs per channel.
+
+        Notes
+        -----
+        This function operates in-place.
+        """
+        if not self.preload:
+            raise ValueError("Data must be preloaded.")
+
+        if len(bad_epochs_indices) != self.get_data().shape[1]:
+            raise RuntimeError(
+                "The length of the list of bad epochs indices "
+                "must match the number of channels."
+            )
+        # extract the data from the epochs object: shape (n_epochs, n_channels, n_times)
+        data = self.get_data()
+        # loop over channels of specific type
+        for ch in range(data.shape[1]):
+            # use those indices to set the epochs per channel to NaN for all timepoints
+            data[bad_epochs_indices[ch], ch, :] = np.nan
+        # put back into epochs structure
+        self.data = data
+
     def _check_consistency(self):
         """Check invariants of epochs object."""
         if hasattr(self, "events"):

mne/tests/test_epochs.py

@@ -5240,3 +5240,19 @@ def test_empty_error(method, epochs_empty):
         pytest.importorskip("pandas")
     with pytest.raises(RuntimeError, match="is empty."):
         getattr(epochs_empty.copy(), method[0])(**method[1])
+
+
+def test_set_bad_epochs_to_nan():
+    """Test channel specific epoch rejection."""
+    # preload=False
+    raw, ev, _ = _get_data(preload=False)
+    ep = Epochs(raw, ev, tmin=0, tmax=0.1, baseline=(0, 0))
+    bads = [[]] * ep.info["nchan"]
+    bads[0] = [1]
+    with pytest.raises(ValueError, match="must be preloaded"):
+        ep.set_bad_epochs_to_NaN(bads)
+
+    # preload=True
+    ep.load_data()
+    ep.set_bad_epochs_to_NaN(bads)
+    _ = ep.average()

mne/utils/check.py

@@ -919,12 +919,12 @@ def _check_combine(mode, valid=("mean", "median", "std"), axis=0):
     if mode == "mean":
 
         def fun(data):
-            return np.mean(data, axis=axis)
+            return np.nanmean(data, axis=axis)
 
     elif mode == "std":
 
         def fun(data):
-            return np.std(data, axis=axis)
+            return np.nanstd(data, axis=axis)
 
     elif mode == "median" or mode == np.median: