Repository: mne-tools/mne-python

Issue:
set_bipolar_reference does not work for epochs data
Somewhat related to my struggles with this function in #8718 . Here is a script to reproduce on master:

```py
import os
import mne

n_picks = 10

sample_data_folder = mne.datasets.sample.data_path()
sample_data_raw_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                    'sample_audvis_raw.fif')
raw = mne.io.read_raw_fif(sample_data_raw_file, verbose=False).crop(tmax=60)

events = mne.find_events(raw, stim_channel='STI 014')

epochs = mne.Epochs(raw, events, tmin=-0.3, tmax=0.7, preload=True)

epochs.pick_types(meg=False, eeg=True)
ch_names = epochs.info['ch_names']
epochs.pick_channels(ch_names[:n_picks])
ch_names = epochs.ch_names

anode, cathode = ch_names[1:], ch_names[:-1]

evoked = epochs.average()
evoked_bipolar1 = mne.set_bipolar_reference(evoked, anode=anode, cathode=cathode)

epochs_bipolar = mne.set_bipolar_reference(epochs, anode=anode, cathode=cathode)
evoked_bipolar2 = epochs_bipolar.average()
```

I get:

```py
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/Documents/github_repos/mne-python/tutorials/epochs/plot_10_epochs_overview.py in <module>
     23 evoked_bipolar1 = mne.set_bipolar_reference(evoked, anode=anode, cathode=cathode)
     24 
---> 25 epochs_bipolar = mne.set_bipolar_reference(epochs, anode=anode, cathode=cathode)
     26 evoked_bipolar2 = epochs_bipolar.average()
     27 

<decorator-gen-232> in set_bipolar_reference(inst, anode, cathode, ch_name, ch_info, drop_refs, copy, verbose)

~/Documents/github_repos/mne-python/mne/io/reference.py in set_bipolar_reference(inst, anode, cathode, ch_name, ch_info, drop_refs, copy, verbose)
    504     if drop_refs:
    505         drop_channels = list((set(anode) | set(cathode)) & set(inst.ch_names))
--> 506         inst.drop_channels(drop_channels)
    507 
    508     return inst

~/Documents/github_repos/mne-python/mne/channels/channels.py in drop_channels(self, ch_names)
    894                    if ch in self.ch_names]
    895         idx = np.setdiff1d(np.arange(len(self.ch_names)), bad_idx)
--> 896         return self._pick_drop_channels(idx)
    897 
    898     def _pick_drop_channels(self, idx):

~/Documents/github_repos/mne-python/mne/channels/channels.py in _pick_drop_channels(self, idx)
    909 
    910         if getattr(self, 'picks', None) is not None:
--> 911             self.picks = self.picks[idx]
    912 
    913         if getattr(self, '_read_picks', None) is not None:

IndexError: index 10 is out of bounds for axis 0 with size 10
```


Discussion:

Resolving patch:
diff --git a/mne/io/reference.py b/mne/io/reference.py
--- a/mne/io/reference.py
+++ b/mne/io/reference.py
@@ -500,6 +500,9 @@ def set_bipolar_reference(inst, anode, cathode, ch_name=None, ch_info=None,
         logger.info('Bipolar channel added as "%s".' % name)
         inst.info._update_redundant()
 
+    if getattr(inst, 'picks', None) is not None:
+        del inst.picks  # picks cannot be tracked anymore
+
     # Drop remaining channels.
     if drop_refs:
         drop_channels = list((set(anode) | set(cathode)) & set(inst.ch_names))


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
@@ -15,7 +15,8 @@
                  set_eeg_reference, set_bipolar_reference,
                  add_reference_channels, create_info, make_sphere_model,
                  make_forward_solution, setup_volume_source_space,
-                 pick_channels_forward, read_evokeds)
+                 pick_channels_forward, read_evokeds,
+                 find_events)
 from mne.epochs import BaseEpochs
 from mne.fixes import nullcontext
 from mne.io import RawArray, read_raw_fif
@@ -314,23 +315,35 @@ def test_set_eeg_reference_rest():
 
 
 @testing.requires_testing_data
-def test_set_bipolar_reference():
+@pytest.mark.parametrize('inst_type', ('raw', 'epochs', 'evoked'))
+def test_set_bipolar_reference(inst_type):
     """Test bipolar referencing."""
     raw = read_raw_fif(fif_fname, preload=True)
     raw.apply_proj()
 
+    if inst_type == 'raw':
+        inst = raw
+        del raw
+    elif inst_type in ['epochs', 'evoked']:
+        events = find_events(raw, stim_channel='STI 014')
+        epochs = Epochs(raw, events, tmin=-0.3, tmax=0.7, preload=True)
+        inst = epochs
+        if inst_type == 'evoked':
+            inst = epochs.average()
+        del epochs
+
     ch_info = {'kind': FIFF.FIFFV_EOG_CH, 'extra': 'some extra value'}
     with pytest.raises(KeyError, match='key errantly present'):
-        set_bipolar_reference(raw, 'EEG 001', 'EEG 002', 'bipolar', ch_info)
+        set_bipolar_reference(inst, 'EEG 001', 'EEG 002', 'bipolar', ch_info)
     ch_info.pop('extra')
     reref = set_bipolar_reference(
-        raw, 'EEG 001', 'EEG 002', 'bipolar', ch_info)
+        inst, 'EEG 001', 'EEG 002', 'bipolar', ch_info)
     assert (reref.info['custom_ref_applied'])
 
     # Compare result to a manual calculation
-    a = raw.copy().pick_channels(['EEG 001', 'EEG 002'])
-    a = a._data[0, :] - a._data[1, :]
-    b = reref.copy().pick_channels(['bipolar'])._data[0, :]
+    a = inst.copy().pick_channels(['EEG 001', 'EEG 002'])
+    a = a._data[..., 0, :] - a._data[..., 1, :]
+    b = reref.copy().pick_channels(['bipolar'])._data[..., 0, :]
     assert_allclose(a, b)
 
     # Original channels should be replaced by a virtual one
@@ -340,7 +353,7 @@ def test_set_bipolar_reference():
 
     # Check channel information
     bp_info = reref.info['chs'][reref.ch_names.index('bipolar')]
-    an_info = reref.info['chs'][raw.ch_names.index('EEG 001')]
+    an_info = reref.info['chs'][inst.ch_names.index('EEG 001')]
     for key in bp_info:
         if key == 'loc':
             assert_array_equal(bp_info[key], 0)
@@ -352,11 +365,11 @@ def test_set_bipolar_reference():
             assert_equal(bp_info[key], an_info[key])
 
     # Minimalist call
-    reref = set_bipolar_reference(raw, 'EEG 001', 'EEG 002')
+    reref = set_bipolar_reference(inst, 'EEG 001', 'EEG 002')
     assert ('EEG 001-EEG 002' in reref.ch_names)
 
     # Minimalist call with twice the same anode
-    reref = set_bipolar_reference(raw,
+    reref = set_bipolar_reference(inst,
                                   ['EEG 001', 'EEG 001', 'EEG 002'],
                                   ['EEG 002', 'EEG 003', 'EEG 003'])
     assert ('EEG 001-EEG 002' in reref.ch_names)
@@ -364,38 +377,41 @@ def test_set_bipolar_reference():
 
     # Set multiple references at once
     reref = set_bipolar_reference(
-        raw,
+        inst,
         ['EEG 001', 'EEG 003'],
         ['EEG 002', 'EEG 004'],
         ['bipolar1', 'bipolar2'],
         [{'kind': FIFF.FIFFV_EOG_CH},
          {'kind': FIFF.FIFFV_EOG_CH}],
     )
-    a = raw.copy().pick_channels(['EEG 001', 'EEG 002', 'EEG 003', 'EEG 004'])
-    a = np.array([a._data[0, :] - a._data[1, :],
-                  a._data[2, :] - a._data[3, :]])
+    a = inst.copy().pick_channels(['EEG 001', 'EEG 002', 'EEG 003', 'EEG 004'])
+    a = np.concatenate(
+        [a._data[..., :1, :] - a._data[..., 1:2, :],
+         a._data[..., 2:3, :] - a._data[..., 3:4, :]],
+        axis=-2
+    )
     b = reref.copy().pick_channels(['bipolar1', 'bipolar2'])._data
     assert_allclose(a, b)
 
     # Test creating a bipolar reference that doesn't involve EEG channels:
     # it should not set the custom_ref_applied flag
-    reref = set_bipolar_reference(raw, 'MEG 0111', 'MEG 0112',
+    reref = set_bipolar_reference(inst, 'MEG 0111', 'MEG 0112',
                                   ch_info={'kind': FIFF.FIFFV_MEG_CH},
                                   verbose='error')
     assert (not reref.info['custom_ref_applied'])
     assert ('MEG 0111-MEG 0112'[:15] in reref.ch_names)
 
     # Test a battery of invalid inputs
-    pytest.raises(ValueError, set_bipolar_reference, raw,
+    pytest.raises(ValueError, set_bipolar_reference, inst,
                   'EEG 001', ['EEG 002', 'EEG 003'], 'bipolar')
-    pytest.raises(ValueError, set_bipolar_reference, raw,
+    pytest.raises(ValueError, set_bipolar_reference, inst,
                   ['EEG 001', 'EEG 002'], 'EEG 003', 'bipolar')
-    pytest.raises(ValueError, set_bipolar_reference, raw,
+    pytest.raises(ValueError, set_bipolar_reference, inst,
                   'EEG 001', 'EEG 002', ['bipolar1', 'bipolar2'])
-    pytest.raises(ValueError, set_bipolar_reference, raw,
+    pytest.raises(ValueError, set_bipolar_reference, inst,
                   'EEG 001', 'EEG 002', 'bipolar',
                   ch_info=[{'foo': 'bar'}, {'foo': 'bar'}])
-    pytest.raises(ValueError, set_bipolar_reference, raw,
+    pytest.raises(ValueError, set_bipolar_reference, inst,
                   'EEG 001', 'EEG 002', ch_name='EEG 003')