Repository: mne-tools/mne-python

Issue:
mne.Epochs keeping reject keys for old channel types after mne.Epochs.pick_types
#### Describe the bug
Opened at the request of @larsoner:
Having a problem with `mne.Epochs` objects retaining the keys in their `reject` and `flat` dictionaries that correspond to channel types that have been removed using `mne.Epochs.pick_types`. This can cause issues when trying to use `mne.Epochs.drop_bad` with thresholds only for the remaining channel types (e.g., those returned by `autoreject`).

#### Steps to reproduce
```
>>> import mne
>>> from mne.datasets import sample
>>> raw_fname = sample.data_path() + '/MEG/sample/sample_audvis_filt-0-40_raw.fif'
>>> event_fname = sample.data_path() + '/MEG/sample/sample_audvis_filt-0-40_raw-eve.fif'
>>> event_id = {'Auditory/Left': 1, 'Auditory/Right': 2}
>>> raw = mne.io.Raw(raw_fname)
>>> events = mne.read_events(event_fname)
>>> picks = mne.pick_types(raw.info, meg=True, eeg=True, stim=False, eog=True)                  
>>> epochs = mne.Epochs(raw, events, event_id, -0.2, 0.5, picks=picks, baseline=(None, 0), reject=dict(grad=4e-10, mag=4e-12, eeg=4e-5, eog=2.5e-4), preload=True)                                                        
>>> epochs.reject                                                                              
# Output: {'grad': 4e-10, 'mag': 4e-12, 'eeg': 4e-05, 'eog': 0.00025}
>>> epochs = epochs.pick_types(meg=True, eeg=False, eog=False)                                 
>>> epochs.reject                                                                              
# Output: {'grad': 4e-10, 'mag': 4e-12, 'eeg': 4e-05, 'eog': 0.00025}
>>> from autoreject import get_rejection_threshold                                             
>>> reject = get_rejection_threshold(epochs)
>>> reject
# Output: {'mag': 2.878764023405042e-12, 'grad': 6.981792417826413e-11}
>>> epochs.drop_bad(reject=reject)
# ValueError
```

#### Expected results
That `mne.Epochs.drop_bad` should successfully drop bad epochs based on thresholds for the remaining channel types, rather than causing issues due to channel types that have been removed.

#### Actual results
```
In [16]: epochs.drop_bad(reject=reject)                                                             
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-16-1c0adc0a26d9> in <module>
----> 1 epochs.drop_bad(reject=reject)

</Users/rnadkarni/miniconda3/envs/megssm/lib/python3.8/site-packages/mne/externals/decorator.py:decorator-gen-176> in drop_bad(self, reject, flat, verbose)

~/miniconda3/envs/megssm/lib/python3.8/site-packages/mne/utils/_logging.py in wrapper(*args, **kwargs)
     88             with use_log_level(verbose_level):
     89                 return function(*args, **kwargs)
---> 90         return function(*args, **kwargs)
     91     return FunctionMaker.create(
     92         function, 'return decfunc(%(signature)s)',

~/miniconda3/envs/megssm/lib/python3.8/site-packages/mne/epochs.py in drop_bad(self, reject, flat, verbose)
   1147                rej in (reject, flat)):
   1148             raise ValueError('reject and flat, if strings, must be "existing"')
-> 1149         self._reject_setup(reject, flat)
   1150         self._get_data(out=False)
   1151         return self

~/miniconda3/envs/megssm/lib/python3.8/site-packages/mne/epochs.py in _reject_setup(self, reject, flat)
    705             new = reject.get(key, np.inf)
    706             if new > old:
--> 707                 raise ValueError(bad_msg.format(kind='reject', key=key,
    708                                                 new=new, old=old, op='>'))
    709         for key in set(flat.keys()).union(old_flat.keys()):

ValueError: reject["eog"] == inf > 0.00015 (old value), new reject values must be at least as stringent as previous ones
```

#### Additional information
```
Platform:      macOS-10.11.6-x86_64-i386-64bit
Python:        3.8.1 (default, Jan  8 2020, 16:15:59)  [Clang 4.0.1 (tags/RELEASE_401/final)]
Executable:    /Users/rnadkarni/miniconda3/envs/megssm/bin/python
CPU:           i386: 8 cores
Memory:        16.0 GB

mne:           0.19.2
numpy:         1.18.1 {blas=openblas, lapack=openblas}
scipy:         1.4.1
matplotlib:    3.1.3 {backend=MacOSX}

sklearn:       0.22.1
numba:         Not found
nibabel:       Not found
cupy:          Not found
pandas:        1.0.1
dipy:          Not found
mayavi:        Not found
pyvista:       Not found
vtk:           Not found
```


Discussion:
epochs.pick_types should cleanup the reject to avoid this issue.

can you open a PR?

>


Resolving patch:
diff --git a/mne/channels/channels.py b/mne/channels/channels.py
--- a/mne/channels/channels.py
+++ b/mne/channels/channels.py
@@ -740,7 +740,23 @@ def pick_types(self, meg=None, eeg=False, stim=False, eog=False,
             ias=ias, syst=syst, seeg=seeg, dipole=dipole, gof=gof, bio=bio,
             ecog=ecog, fnirs=fnirs, include=include, exclude=exclude,
             selection=selection)
-        return self._pick_drop_channels(idx)
+
+        self._pick_drop_channels(idx)
+
+        # 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 during iteration
+            for ch_type in list(self.reject):
+                if ch_type not in self:
+                    del self.reject[ch_type]
+
+        if getattr(self, 'flat', None) is not None:
+            for ch_type in list(self.flat):
+                if ch_type not in self:
+                    del self.flat[ch_type]
+
+        return self
 
     def pick_channels(self, ch_names, ordered=False):
         """Pick some channels.


Test patch:
diff --git a/mne/tests/test_epochs.py b/mne/tests/test_epochs.py
--- a/mne/tests/test_epochs.py
+++ b/mne/tests/test_epochs.py
@@ -3003,4 +3003,20 @@ def test_epochs_get_data_item(preload):
     assert_array_equal(one_data, one_epo.get_data())
 
 
+def test_pick_types_reject_flat_keys():
+    """Test that epochs.pick_types removes keys from reject/flat."""
+    raw, events, _ = _get_data()
+    event_id = {'a/1': 1, 'a/2': 2, 'b/1': 3, 'b/2': 4}
+    picks = pick_types(raw.info, meg=True, eeg=True, ecg=True, eog=True)
+    epochs = Epochs(raw, events, event_id, preload=True, picks=picks,
+                    reject=dict(grad=1e-10, mag=1e-10, eeg=1e-3, eog=1e-3),
+                    flat=dict(grad=1e-16, mag=1e-16, eeg=1e-16, eog=1e-16))
+
+    assert sorted(epochs.reject.keys()) == ['eeg', 'eog', 'grad', 'mag']
+    assert sorted(epochs.flat.keys()) == ['eeg', 'eog', 'grad', 'mag']
+    epochs.pick_types(meg=True, eeg=False, ecg=False, eog=False)
+    assert sorted(epochs.reject.keys()) == ['grad', 'mag']
+    assert sorted(epochs.flat.keys()) == ['grad', 'mag']
+
+
 run_tests_if_main()