diff --git a/doc/changes/dev/12219.newfeature.rst b/doc/changes/dev/12219.newfeature.rst new file mode 100644 index 00000000000..0c79f512033 --- /dev/null +++ b/doc/changes/dev/12219.newfeature.rst @@ -0,0 +1,7 @@ +Add mark_bad_epochs_by_channel method to :class:mne.Epochs for channel-specific epoch rejection + +This method allows users to mark bad epochs on a per-channel basis by setting +them to NaN. An additional nave_per_channel attribute for epochs is to reflect the number of valid epochs per channel. +Averaging epochs with NaNs will fail. + +Contributed by `Carina Forster`_. diff --git a/mne/channels/channels.py b/mne/channels/channels.py index d0e9ae73d6e..b5703cd2195 100644 --- a/mne/channels/channels.py +++ b/mne/channels/channels.py @@ -499,8 +499,13 @@ def pick(self, picks, exclude=(), *, verbose=None): The modified instance. """ picks = _picks_to_idx(self.info, picks, "all", exclude, allow_empty=False) + # get channel names + ch_names = [self.ch_names[p] for p in picks] self._pick_drop_channels(picks) + # how many epochs per channel after channel specific epoch rejection + nave_per_channel = getattr(self, "nave_per_channel", None) + # remove dropped channel types from reject and flat if getattr(self, "reject", None) is not None: # use list(self.reject) to avoid RuntimeError for changing dictionary size @@ -514,6 +519,13 @@ def pick(self, picks, exclude=(), *, verbose=None): if ch_type not in self: del self.flat[ch_type] + if nave_per_channel is not None: + # self is the epochs object, always has the same number of channels + nave_dict = dict(zip(self.info["ch_names"], nave_per_channel)) + self.nave_per_channel = np.array( + [nave_dict[ch] for ch in ch_names if ch in nave_dict] + ) + return self def reorder_channels(self, ch_names): diff --git a/mne/epochs.py b/mne/epochs.py index fa0292041b1..c73f64b8bc7 100644 --- a/mne/epochs.py +++ b/mne/epochs.py @@ -718,6 +718,59 @@ def __init__( self._check_consistency() self.set_annotations(annotations, on_missing="ignore") + def mark_bad_channels_per_epoch(self, reject_mask=None): + """Mark bad channels on an epoch by epoch basis. + + Warning: This is only useful for channel-wise analyses. + + Parameters + ---------- + reject_mask : np.ndarray, shape (n_epochs, n_channels) | None + Boolean mask where True indicates an epoch marked with + NaN for a specific channel. + If None, no epochs are marked. + + Returns + ------- + epochs : instance of Epochs + The epochs object with bad epochs marked with NaNs per channel. + Operates in-place. + """ + if reject_mask is None: + return self + + if not self.preload: + raise ValueError("Epochs must be preloaded.") + + data = self._data + assert data is not None # preloaded => _data is an ndarray + + n_epochs, n_channels, _ = data.shape + + if reject_mask.shape != (n_epochs, n_channels): + raise ValueError( + f"reject_mask must have shape ({n_epochs}, {n_channels}), " + f"got {reject_mask.shape}" + ) + + # required: bool -> boolean indexing, not int + if not np.issubdtype(reject_mask.dtype, np.bool_): + reject_mask = reject_mask.astype(bool) + + # Set bad epochs to NaN + data[reject_mask] = np.nan + + # store mask for updating nave + self.reject_mask = reject_mask + + # store nave per channel for updating nave + valid_epochs_per_channel = np.sum(~reject_mask, axis=0) + + # currently no documentation on that attribute + self.nave_per_channel = valid_epochs_per_channel + + return self + def _check_events_outside_data(self, on_outside, raw): """Warn when events fall outside the range of the recorded data (gh-12989).""" if raw is not None and hasattr(raw, "first_samp") and len(self.events) > 0: @@ -1232,6 +1285,13 @@ def _compute_aggregate(self, picks, mode="mean"): n_events = len(self.events) fun = _check_combine(mode, valid=("mean", "median", "std")) data = fun(self._data) + if np.isnan(data).any(): + raise ValueError( + "Cannot average epochs containing NaNs (possibly introduced by " + "mark_bad_channels_per_epoch): any channel with a rejected epoch " + "would average to NaN. Extract the data with get_data() and use " + "np.nanmean over the epoch axis if you need per-channel averages." + ) assert len(self.events) == len(self._data) if data.shape != self._data.shape[1:]: raise RuntimeError( @@ -1326,12 +1386,15 @@ def _evoked_from_epoch_data(self, data, info, picks, n_events, kind, comment): """Create an evoked object from epoch data.""" info = deepcopy(info) # don't apply baseline correction; we'll set evoked.baseline manually + + nave = n_events + evoked = EvokedArray( data, info, tmin=self.times[0], comment=comment, - nave=n_events, + nave=nave, kind=kind, baseline=None, ) diff --git a/mne/tests/test_epochs.py b/mne/tests/test_epochs.py index 922ee4e4d8a..c77f5cfd27f 100644 --- a/mne/tests/test_epochs.py +++ b/mne/tests/test_epochs.py @@ -5307,6 +5307,75 @@ def test_empty_error(method, epochs_empty): getattr(epochs_empty.copy(), method[0])(**method[1]) +def test_mark_bad_channels_per_epoch(): + """Test channel-specific epoch rejection.""" + # load raw and events data without loading data to disk + raw, ev, _ = _get_data(preload=False) + ep = Epochs(raw, ev, tmin=0, tmax=0.1, baseline=(0, 0), preload=False) + + # extract shape to set up reject mask (can't use shape as it loads the data) + n_epochs = len(ep.events) # number of epochs + n_channels = len(ep.ch_names) # number of channels + + # create a dummy reject mask with correct shape + reject_mask_dummy = np.zeros((n_epochs, n_channels)) + + # should throw an error + with pytest.raises(ValueError, match="must be preloaded"): + ep.mark_bad_channels_per_epoch(reject_mask_dummy) + + # load data + ep.load_data() + + # test if reject_mask == None returns epochs + assert ep == ep.mark_bad_channels_per_epoch(None) + + # set epochs to bad in reject mask + reject_mask = np.zeros((n_epochs, n_channels), dtype=bool) # all epochs are good + reject_mask[1, 0] = True # second epoch, first channel -> bad + reject_mask[1:, 1] = True # all epochs from channel two are bad + reject_mask[3, 2] = True # fourth epoch, third channel -> bad + + # this is a edge case, averaging throws an error because of empty channel + # realistically the user will drop the channel if all epochs are bad + # reject_mask[:, 1] = True # all epochs from channel two are bad + + # drop bad epochs + ep.mark_bad_channels_per_epoch(reject_mask) + + # verify bad epochs are NaN after dropping them + data = ep.get_data() + assert np.all(np.isnan(data[1, 0, :])) and np.all(np.isnan(data[3, 2, :])) + assert np.all(np.isnan(data[1:, 1, :])) + + # now we should have a nave per channel attribute + # now self.nave_per_channel should be assigned + assert hasattr(ep, "nave_per_channel") + + # sum over good epochs per channel + true_nave_per_channel = np.sum(~np.all(np.isnan(data), axis=2), axis=0) + assert np.all(ep.nave_per_channel == true_nave_per_channel) + + # channel length must match + assert len(ep.nave_per_channel) == len(ep.ch_names) + + # test mask that contains floats instead of bool + float_mask = reject_mask.astype(float) + ep.mark_bad_channels_per_epoch(float_mask) + + data = ep.get_data() + assert np.all(np.isnan(data[1, 0, :])) and np.all(np.isnan(data[3, 2, :])) + + # test wrong shape of rejection mask + bad_mask = np.zeros((n_epochs, n_channels - 1), dtype=bool) + with pytest.raises(ValueError, match="reject_mask must have shape"): + ep.mark_bad_channels_per_epoch(bad_mask) + + # make sure averaging breaks and throws intelligent error for user + with pytest.raises(ValueError, match="Cannot average epochs containing NaNs"): + ep.average() + + def test_epochs_warn_out_of_bounds_events(): """Warn when event sample numbers fall outside the recorded data (gh-12989).""" sfreq = 100.0