Repository: mne-tools/mne-python

Issue:
Picking after adding a reference channel is broken for Epochs
First reported on the forum at https://mne.discourse.group/t/can-not-reorder-channels-after-adding-a-reference-channel/5180


Works:
```python
# %%
import mne

sample_dir = mne.datasets.sample.data_path()
sample_fname = sample_dir / 'MEG' / 'sample' / 'sample_audvis_raw.fif'

raw = mne.io.read_raw_fif(sample_fname).crop(tmax=10).load_data()
raw.add_reference_channels(['REF'])
raw.pick_types(eeg=True)
```

Doesn't work:
```python
# %%
import mne

sample_dir = mne.datasets.sample.data_path()
sample_fname = sample_dir / 'MEG' / 'sample' / 'sample_audvis_raw.fif'

raw = mne.io.read_raw_fif(sample_fname).crop(tmax=60).load_data()
events = mne.find_events(raw, stim_channel='STI 014')
epochs = mne.Epochs(raw, events=events, preload=True)

epochs.add_reference_channels(['REF'])
epochs.pick_types(eeg=True)
```
and yields:
```python
IndexError                                Traceback (most recent call last)
Input In [30], in <cell line: 12>()
     10 epochs = mne.Epochs(raw, events=events, preload=True)
     11 epochs.add_reference_channels(['REF'])
---> 12 epochs.pick_types(eeg=True)

File <decorator-gen-43>:12, in pick_types(self, meg, eeg, stim, eog, ecg, emg, ref_meg, misc, resp, chpi, exci, ias, syst, seeg, dipole, gof, bio, ecog, fnirs, csd, dbs, include, exclude, selection, verbose)

File ~/Development/mne-python/mne/channels/channels.py:681, in UpdateChannelsMixin.pick_types(self, meg, eeg, stim, eog, ecg, emg, ref_meg, misc, resp, chpi, exci, ias, syst, seeg, dipole, gof, bio, ecog, fnirs, csd, dbs, include, exclude, selection, verbose)
    597 """Pick some channels by type and names.
    598 
    599 Parameters
   (...)
    672 .. versionadded:: 0.9.0
    673 """
    674 idx = pick_types(
    675     self.info, meg=meg, eeg=eeg, stim=stim, eog=eog, ecg=ecg, emg=emg,
    676     ref_meg=ref_meg, misc=misc, resp=resp, chpi=chpi, exci=exci,
    677     ias=ias, syst=syst, seeg=seeg, dipole=dipole, gof=gof, bio=bio,
    678     ecog=ecog, fnirs=fnirs, csd=csd, dbs=dbs, include=include,
    679     exclude=exclude, selection=selection)
--> 681 self._pick_drop_channels(idx)
    683 # remove dropped channel types from reject and flat
    684 if getattr(self, 'reject', None) is not None:
    685     # use list(self.reject) to avoid RuntimeError for changing
    686     # dictionary size during iteration

File <decorator-gen-46>:12, in _pick_drop_channels(self, idx, verbose)

File ~/Development/mne-python/mne/channels/channels.py:855, in UpdateChannelsMixin._pick_drop_channels(self, idx, verbose)
    852     _check_preload(self, msg)
    854 if getattr(self, 'picks', None) is not None:
--> 855     self.picks = self.picks[idx]
    857 if getattr(self, '_read_picks', None) is not None:
    858     self._read_picks = [r[idx] for r in self._read_picks]

IndexError: index 376 is out of bounds for axis 0 with size 376
```



Discussion:
I'm not sure this should be a blocker for 1.1. If we can fix it, great, but if not, we can fix early in 1.2 and backport
I agree, let's get 1.1 out the door ASAP! 

Resolving patch:
diff --git a/mne/io/reference.py b/mne/io/reference.py
--- a/mne/io/reference.py
+++ b/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)


Test patch:
diff --git a/mne/io/tests/test_reference.py b/mne/io/tests/test_reference.py
--- a/mne/io/tests/test_reference.py
+++ b/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):