BUG: Fix bug with Epochs.add_reference_channels (#10912)

Author: larsoner

doc/changes/latest.inc

@@ -143,6 +143,8 @@ Bugs
 
 - Fix bug in coregistration GUI that prevented it from starting up if only a high-resolution head model was available (:gh:`10543` by `Richard Höchenberger`_)
 
+- Fix bug with :class:`mne.Epochs.add_reference_channels` where attributes were not updated properly so subsequent `~mne.Epochs.pick_types` calls were broken (:gh:`10912` by `Eric Larson`_)
+- 
 - Fix bug in the :class:`mne.viz.Brain` tool bar that prevented the buttons to call the corresponding feature (:gh:`10560` by `Guillaume Favelier`_)
 
 - Fix issue with saving epochs once :func:`~mne.preprocessing.compute_current_source_density` has been used if a rejection threshold was used first (:gh:`10619` by `Alex Rockhill`_ and `Richard Höchenberger`_)

mne/io/reference.py

@@ -241,12 +241,16 @@ def add_reference_channels(inst, ref_channels, copy=True):
                      'loc': ref_dig_array}
         inst.info['chs'].append(chan_info)
         inst.info._update_redundant()
+    range_ = np.arange(1, len(ref_channels) + 1)
     if isinstance(inst, BaseRaw):
         inst._cals = np.hstack((inst._cals, [1] * len(ref_channels)))
-        range_ = np.arange(1, len(ref_channels) + 1)
         for pi, picks in enumerate(inst._read_picks):
             inst._read_picks[pi] = np.concatenate(
                 [picks, np.max(picks) + range_])
+    elif isinstance(inst, BaseEpochs):
+        picks = inst.picks
+        inst.picks = np.concatenate(
+            [picks, np.max(picks) + range_])
     inst.info._check_consistency()
     set_eeg_reference(inst, ref_channels=ref_channels, copy=False,
                       verbose=False)

mne/io/tests/test_reference.py

@@ -18,7 +18,7 @@
                  make_forward_solution, setup_volume_source_space,
                  pick_channels_forward, read_evokeds,
                  find_events)
-from mne.epochs import BaseEpochs
+from mne.epochs import BaseEpochs, make_fixed_length_epochs
 from mne.io import RawArray, read_raw_fif
 from mne.io.constants import FIFF
 from mne.io.proj import _has_eeg_average_ref_proj, Projection
@@ -604,6 +604,19 @@ def test_add_reference():
     with pytest.raises(TypeError, match='instance of'):
         add_reference_channels(raw, 1)
 
+    # gh-10878
+    raw = read_raw_fif(raw_fname).crop(0, 1, include_tmax=False).load_data()
+    data = raw.copy().add_reference_channels(['REF']).pick_types(eeg=True)
+    data = data.get_data()
+    epochs = make_fixed_length_epochs(raw).load_data()
+    data_2 = epochs.copy().add_reference_channels(['REF']).pick_types(eeg=True)
+    data_2 = data_2.get_data()[0]
+    assert_allclose(data, data_2)
+    evoked = epochs.average()
+    data_3 = evoked.copy().add_reference_channels(['REF']).pick_types(eeg=True)
+    data_3 = data_3.get_data()
+    assert_allclose(data, data_3)
+
 
 @pytest.mark.parametrize('n_ref', (1, 2))
 def test_add_reorder(n_ref):